import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Created by sol on 2015-10-21.
 */
public class NaiveBayes {

    List<String> categories;

    int totalWords = 0;
    int totalDocuments = 0;
    Map<String,Map<String, Integer>> occurrencesPerWordAndCategory;
    Map<String, Integer> documentsPerCategory;
    Map<String, Integer> numberOfWordsPerCategory;

    public NaiveBayes(List<String> categories) {
        this.categories = categories;
        occurrencesPerWordAndCategory = new HashMap();
        documentsPerCategory = new HashMap();
        numberOfWordsPerCategory = new HashMap();
        setupMaps();
    }

    /**
     * Initiate maps in order to avoid nullpointer
     */
    private void setupMaps() {
        for (String category : categories) {
            occurrencesPerWordAndCategory.put(category, new HashMap<String, Integer>());
            numberOfWordsPerCategory.put(category, 0);
            documentsPerCategory.put(category, 0);
        }
    }

    public void train(String category, String document){
//        document = document.toLowerCase();
        Map<String, Integer> occurrencesPerWordLocal = occurrencesPerWordAndCategory.get(category);
        String[] words = document.split(" ");
        for (String word : words) {
            //Count occurrences of word in category
            Integer occurrences = occurrencesPerWordLocal.get(word);
            if (occurrences == null) {
                occurrencesPerWordLocal.put(word, 1);
            } else {
                occurrencesPerWordLocal.put(word, occurrences + 1);
            }

        }
        //total number of words in category
        numberOfWordsPerCategory.put(category, numberOfWordsPerCategory.get(category) + words.length);
        totalWords = totalWords + words.length;
        documentsPerCategory.put(category, documentsPerCategory.get(category) + 1);
        totalDocuments++;
    }

    public String classify(String document){
        //calculate probability for each category given the document
        String bestCategory = "";
        double bestProbability = 0;
        double probability = 0;
        for (String category : categories) {
            System.out.print(category + " ");
            probability = probabilty(category, document);
            System.out.println(probability);
            if (probability > bestProbability) {
                bestProbability = probability;
                bestCategory = category;
            }
        }
        return bestCategory;
    }

    private double probabilty(String category, String document) {
        return documentBelongsToCategoryProbability(category, document) * isCategoryProbability(category);
    }

    private double isCategoryProbability(String category) {
        return (double) documentsPerCategory.get(category) / totalDocuments;
    }

    private double documentBelongsToCategoryProbability(String category, String document) {
        double probability = 1;
        String[] words = document.split(" ");
        for (String word : words) {
            probability = probability * wordBelongsCategoryProbability(category, word);
        }
        return probability;
    }

    private double wordBelongsCategoryProbability(String category, String word) {
        Integer wordCount = occurrencesPerWordAndCategory.get(category).get(word);
        if (wordCount == null) {
            wordCount = 1;
        } else {
            wordCount++;
        }
        return (double) wordCount / numberOfWordsPerCategory.get(category);
    }

}

@@ -469,32 +469,37 @@
 gories;%0A        
+this.
 occurrencesPerWo
@@ -517,32 +517,61 @@
 ry = new HashMap
+%3CString,Map%3CString, Integer%3E%3E
 ();%0A        docu
@@ -558,32 +558,37 @@
 er%3E%3E();%0A        
+this.
 documentsPerCate
@@ -597,32 +597,49 @@
 ry = new HashMap
+%3CString, Integer%3E
 ();%0A        numb
@@ -626,32 +626,37 @@
 ger%3E();%0A        
+this.
 numberOfWordsPer
@@ -677,16 +677,33 @@
  HashMap
+%3CString, Integer%3E
 ();%0A    

Remove warnings of unchecked types