class InformationGainFeatureRanker extends AbstractFeatureRanker {

    @Override
    public FeatureRanking rankFeatures(Dataset dataset, ProgressReporter progress) {
        Validate.notNull(dataset, "dataset must not be null");
        Map<String, Double> informationGainValues = new HashMap<>();

        progress.startTask("Information Gain", -1);
        LOGGER.debug("Calculating discretization");
        Discretization discretization = new Discretization(dataset, progress.createSubProgress(0.5));
        Dataset preparedData = dataset.transform(discretization);

        NominalValueStatistics categoryStatistics = new DatasetStatistics(dataset).getCategoryStatistics();
        CategoryEntries categoryCounts = new CategoryEntriesBuilder(categoryStatistics.getMap()).create();
        double entropy = categoryCounts.getEntropy();
        Set<String> featureNames = dataset.getFeatureInformation().getFeatureNames();

        ProgressReporter informationGainProgress = progress.createSubProgress(0.5);
        LOGGER.debug("Calculating gain");
        informationGainProgress.startTask("Calculating gain", featureNames.size());
        for (String featureName : featureNames) {
            double gain = entropy - conditionalEntropy(featureName, preparedData);
            informationGainValues.put(featureName, gain);
            informationGainProgress.increment();
        }
        informationGainProgress.finishTask();
        return new FeatureRanking(informationGainValues);
    }

}