@Override
    public SelectedFeatureSet select(Dataset<Label> dataset) {
        FSMatrix data = FSMatrix.buildMatrix(dataset,numBins);
        ImmutableFeatureMap fmap = data.getFeatureMap();
        int max = k == -1 ? fmap.size() : Math.min(k,fmap.size());
        int numFeatures = fmap.size();

        boolean[] unselectedFeatures = new boolean[numFeatures];
        Arrays.fill(unselectedFeatures, true);
        int[] selectedFeatures = new int[max];
        double[] selectedScores = new double[max];

        ForkJoinPool fjp = null;

        Pair<Integer,Double> maxPair;
        if (numThreads > 1) {
            fjp = new ForkJoinPool(numThreads);
            try {
                maxPair = fjp.submit(() -> IntStream.range(0, numFeatures).parallel().mapToObj(i -> new Pair<>(i, data.mi(i))).max(Comparator.comparingDouble(Pair::getB)).get()).get();
            } catch (InterruptedException | ExecutionException e) {
                throw new RuntimeException(e);
            }
        } else {
            maxPair = IntStream.range(0, numFeatures).mapToObj(i -> new Pair<>(i,data.mi(i))).max(Comparator.comparingDouble(Pair::getB)).get();
        }

        int selectedIdx = maxPair.getA();
        selectedFeatures[0] = selectedIdx;
        unselectedFeatures[selectedIdx] = false;
        selectedScores[0] = maxPair.getB();
        logger.log(Level.INFO,"Itr 0: selected feature " + fmap.get(selectedIdx).getName() + ", score = " + selectedScores[0]);

        double[] jmiScore = new double[numFeatures];

        //
        // Select features in max JMI order
        for (int i = 1; i < max; i++) {
            if (numThreads > 1) {
                final int prevIdx = selectedFeatures[i-1];
                try {
                    double[] updates = fjp.submit(() -> IntStream.range(0, numFeatures).parallel().mapToDouble(j -> unselectedFeatures[j] ? data.jmi(j, prevIdx) : 0.0).toArray()).get();
                    for (int j = 0; j < jmiScore.length; j++) {
                        jmiScore[j] += updates[j];
                    }
                    maxPair = fjp.submit(() -> IntStream.range(0, numFeatures).parallel().filter(j -> unselectedFeatures[j]).mapToObj(j -> new Pair<>(j, jmiScore[j])).max(Comparator.comparingDouble(Pair::getB)).get()).get();
                } catch (InterruptedException | ExecutionException e) {
                    throw new RuntimeException(e);
                }
            } else {
                int maxIndex = -1;
                double maxScore = Double.NEGATIVE_INFINITY;
                for (int j = 0; j < numFeatures; j++) {
                    if (unselectedFeatures[j]) {
                        int prevIdx = selectedFeatures[i-1];
                        jmiScore[j] += data.jmi(j,prevIdx);
                        if (jmiScore[j] > maxScore) {
                            maxScore = jmiScore[j];
                            maxIndex = j;
                        }
                    }
                }
                maxPair = new Pair<>(maxIndex,maxScore);
            }
            int maxIdx = maxPair.getA();
            selectedFeatures[i] = maxIdx;
            unselectedFeatures[maxIdx] = false;
            selectedScores[i] = maxPair.getB() / i;

            logger.log(Level.INFO,"Itr " + i + ": selected feature " + fmap.get(maxIdx).getName() + ", score = " + maxPair.getB() + ", average score = " + selectedScores[i]);
        }

        if (fjp != null) {
            fjp.shutdown();
        }

        ArrayList<String> names = new ArrayList<>();
        ArrayList<Double> scores = new ArrayList<>();
        for (int i = 0; i < max; i++) {
            names.add(fmap.get(selectedFeatures[i]).getName());
            scores.add(selectedScores[i]);
        }

        FeatureSetProvenance provenance = new FeatureSetProvenance(SelectedFeatureSet.class.getName(),dataset.getProvenance(),getProvenance());
        return new SelectedFeatureSet(names,scores,isOrdered(),provenance);
    }
@Test
    public void parallelismTest() {
        Dataset<Label> dataset = MIMTest.createDataset();

        JMI sequential = new JMI(-1,5,1);
        SelectedFeatureSet sequentialSet = sequential.select(dataset);

        JMI parallel = new JMI(-1,5,4);
        SelectedFeatureSet parallelSet = parallel.select(dataset);

        assertEquals(sequentialSet.featureNames(),parallelSet.featureNames());
        double[] sequentialScores = Util.toPrimitiveDouble(sequentialSet.featureScores());
        double[] parallelScores = Util.toPrimitiveDouble(parallelSet.featureScores());
        assertArrayEquals(sequentialScores,parallelScores,1e-15);
    }