package interactivelearner.featureselection;

import interactivelearner.util.Category;
import interactivelearner.util.Corpus;

import java.util.*;
import java.util.stream.Collectors;

public class ChiSquare {

    private static final int HIGHEST = 300;

    /**
     * Creates a list of words to be used for processing text documents.
     * @param corpus the corpus to be used for selecting the meaningful words
     * @return list of meaningful words
     */
    public static List<String> createVocabulary(Corpus corpus) {
        Set<String> words = corpus.getWords();
        HashMap<String, Double> wordScores = new HashMap<>();
        for (String word : words) {
            wordScores.put(word, chiSquareValue(word, corpus));
        }
        return getHighestValues(wordScores);
    }

    /**
     * Calculates the chi-square value of a word.
     * @param word the word for which the chi-square value needs to be calculated
     * @param corpus the word set in calculating the chi-square value
     * @return the chi-square value of the given word
     */
    private static double chiSquareValue(String word, Corpus corpus) {
        double result = 0;
        int categoryAmount = corpus.getCategories().size();
        int[][] table = new int[categoryAmount + 1][3];

        for (int i = 0; i < categoryAmount; i++) {
            Category category = corpus.getCategories().get(i);
            table[i][0] = category.documentCount(word);
            table[categoryAmount][0] += table[i][0];
            table[i][2] = category.totalDocuments();
            table[i][1] = table[i][2] - table[i][0];
            table[categoryAmount][1] += table[i][1];
        }

        table[categoryAmount][2] = table[categoryAmount][0] + table[categoryAmount][1];

        for (int i = 0; i < categoryAmount; i++) {
            for (int j = 0; j < 2; j++) {
                double expectedValue = ((double) table[categoryAmount][j] * (double) table[i][2]) / (double) table[categoryAmount][2];
                double realValue = table[i][j];
                result += Math.pow(realValue - expectedValue, 2) / expectedValue;
            }
        }
        return result;
    }

    /**
     * Returns a list of words corresponding to the highest scored chi-square values.
     * @param wordScores map of words and chi-square values
     * @return list of best words to use as vocabulary
     */
    private static List<String> getHighestValues(HashMap<String, Double> wordScores) {
        int loopCount = (wordScores.values().size() > HIGHEST) ? HIGHEST : wordScores.values().size();
        List<String> result = new ArrayList<>();

        Set<Double> chiValues = new HashSet<>(wordScores.values());
        List<Double> sortedValues = new ArrayList<>(chiValues);
        sortedValues.sort(Collections.reverseOrder());

        for (int i = 0; i < loopCount; i++) {
            result.addAll(getKeysByValue(wordScores, sortedValues.get(i)));
        }
        return result;
    }

    /**
     * Custom method to convert a value of a map into its key(s).
     * @param map
     * @param value
     * @return a set of keys
     */
    private static Set<String> getKeysByValue(HashMap<String, Double> map, Double value) {
        return map.entrySet()
                .stream()
                .filter(entry -> Objects.equals(entry.getValue(), value))
                .map(Map.Entry::getKey)
                .collect(Collectors.toSet());
    }


}

@@ -211,40 +211,147 @@
     
-private static final int HIGHEST
+/**%0A     * The maximum amount of selected words in the resulting Vocabulary.%0A     */%0A    private static final int MAXIMUM_VOCABULARY_LENGTH
  = 3
@@ -2595,871 +2595,184 @@
     
-int loopCount = (wordScores.values().size() %3E HIGHEST) ? HIGHEST : wordScores.values().size();%0A        List%3CString%3E result = new ArrayList%3C%3E();%0A%0A        Set%3CDouble%3E chiValues = new HashSet%3C%3E(wordScores.values());%0A        List%3CDouble%3E sortedValues = new ArrayList%3C%3E(chiValues);%0A        sortedValues.sort(Collections.reverseOrder());%0A%0A        for (int i = 0; i %3C loopCount; i++) %7B%0A            result.addAll(getKeysByValue(wordScores, sortedValues.get(i)));%0A        %7D%0A        return result;%0A    %7D%0A%0A    /**%0A     * Custom method to convert a value of a map into its key(s).%0A     * @param map%0A     * @param value%0A     * @return a set of keys%0A     */%0A    private static Set%3CString%3E getKeysByValue(HashMap%3CString, Double%3E map, Double value) %7B%0A        return map.entrySet()%0A                .stream()%0A                .filter(entry -%3E Objects.equals(e
+return wordScores.entrySet()%0A                .stream()%0A                .sorted((o1, o2) -%3E -Double.compare(o1.getValue(), o2.getValue()))%0A                .map(Map.E
 ntry
-.getValue(), value)
+::getKey
 )%0A  
@@ -2786,37 +2786,47 @@
        .
-map(Map.Entry::getKey
+limit(MAXIMUM_VOCABULARY_LENGTH
 )%0A      
@@ -2857,18 +2857,19 @@
 ctors.to
-Se
+Lis
 t());%0A  
@@ -2872,12 +2872,10 @@
 ;%0A    %7D%0A
-%0A%0A
 %7D%0A

Change way of sorting Map values to use shorter lambda