package de.jungblut.classification.bayes;

import gnu.trove.set.hash.TIntHashSet;

import java.util.Iterator;

import com.google.common.base.Preconditions;

import de.jungblut.classification.AbstractClassifier;
import de.jungblut.math.DoubleVector;
import de.jungblut.math.DoubleVector.DoubleVectorElement;
import de.jungblut.math.dense.DenseDoubleMatrix;
import de.jungblut.math.dense.DenseDoubleVector;

/**
 * Simple multinomial naive bayes classifier.
 * 
 * @author thomas.jungblut
 * 
 */
public final class MultinomialNaiveBayesClassifier extends AbstractClassifier {

  private DenseDoubleMatrix probabilityMatrix;
  private DenseDoubleVector classProbability;

  /**
   * Default constructor to construct this classifier.
   */
  public MultinomialNaiveBayesClassifier() {
  }

  /**
   * Deserialization constructor to instantiate an already trained classifier
   * from the internal representations.
   * 
   * @param probabilityMatrix the probability matrix.
   * @param classProbability the prior class probabilities.
   */
  public MultinomialNaiveBayesClassifier(DenseDoubleMatrix probabilityMatrix,
      DenseDoubleVector classProbability) {
    super();
    this.probabilityMatrix = probabilityMatrix;
    this.classProbability = classProbability;
  }

  @Override
  public void train(DoubleVector[] features, DoubleVector[] outcome) {
    Preconditions.checkArgument(features.length > 0,
        "Features must contain at least a single item!");
    Preconditions.checkArgument(features.length == outcome.length,
        "There must be an equal amount of features and prediction outcomes!");
    TIntHashSet distinctClasses = new TIntHashSet();
    int[] classes = new int[outcome.length];
    for (int i = 0; i < outcome.length; i++) {
      if (outcome[i].getDimension() == 1) {
        classes[i] = (int) outcome[i].get(0);
      } else {
        classes[i] = outcome[i].maxIndex();
      }
      distinctClasses.add(classes[i]);
    }
    final int numDistinctClasses = distinctClasses.size();
    // add-one smooth
    probabilityMatrix = new DenseDoubleMatrix(numDistinctClasses,
        features[0].getDimension(), 1.0d);

    int[] tokenPerClass = new int[numDistinctClasses];
    int[] numDocumentsPerClass = new int[numDistinctClasses];

    // init the probability with the document length = word count for each token
    for (int columnIndex = 0; columnIndex < features.length; columnIndex++) {
      final DoubleVector document = features[columnIndex];
      final int predictedClass = classes[columnIndex];
      tokenPerClass[predictedClass] += document.getLength();
      numDocumentsPerClass[predictedClass]++;

      Iterator<DoubleVectorElement> iterateNonZero = document.iterateNonZero();
      while (iterateNonZero.hasNext()) {
        DoubleVectorElement next = iterateNonZero.next();
        double currentCount = probabilityMatrix.get(predictedClass,
            next.getIndex());
        probabilityMatrix.set(predictedClass, next.getIndex(), currentCount
            + next.getValue());
      }
    }

    // know we know the token distribution per class, we can calculate the
    // probability. It is intended for them to be negative in some cases
    for (int row = 0; row < numDistinctClasses; row++) {
      for (int tokenColumn = 0; tokenColumn < probabilityMatrix
          .getColumnCount(); tokenColumn++) {
        double currentWordCount = probabilityMatrix.get(row, tokenColumn);
        double logLikelyhood = Math.log(currentWordCount
            / (tokenPerClass[row] + probabilityMatrix.getColumnCount() - 1));
        probabilityMatrix.set(row, tokenColumn, logLikelyhood);
      }
    }

    classProbability = new DenseDoubleVector(numDistinctClasses);
    for (int i = 0; i < numDistinctClasses; i++) {
      classProbability.set(i, (numDocumentsPerClass[i])
          / (double) features.length);
    }
  }

  /**
   * Returns the maximum likely class.
   */
  public int classify(DoubleVector document) {
    return getProbabilityDistribution(document).maxIndex();
  }

  @Override
  public DoubleVector predict(DoubleVector features) {
    return getProbabilityDistribution(features);
  }

  private double getProbabilityForClass(DoubleVector document, int classIndex) {
    double probabilitySum = 0.0d;
    Iterator<DoubleVectorElement> iterateNonZero = document.iterateNonZero();
    while (iterateNonZero.hasNext()) {
      DoubleVectorElement next = iterateNonZero.next();
      double wordCount = next.getValue();
      double probabilityOfToken = probabilityMatrix.get(classIndex,
          next.getIndex());
      probabilitySum += (wordCount * probabilityOfToken);
    }
    return probabilitySum;
  }

  private DenseDoubleVector getProbabilityDistribution(DoubleVector document) {

    int numClasses = classProbability.getLength();
    DenseDoubleVector distribution = new DenseDoubleVector(numClasses);
    // loop through all classes and get the max probable one
    for (int i = 0; i < numClasses; i++) {
      double probability = getProbabilityForClass(document, i);
      distribution.set(i, probability);
    }

    double maxProbability = distribution.max();
    double probabilitySum = 0.0d;
    // we normalize it back
    for (int i = 0; i < numClasses; i++) {
      double probability = distribution.get(i);
      double normalizedProbability = Math.exp(probability - maxProbability)
          * classProbability.get(i);
      distribution.set(i, normalizedProbability);
      probabilitySum += normalizedProbability;
    }

    // since the sum is sometimes not 1, we need to divide by the sum
    distribution = (DenseDoubleVector) distribution.divide(probabilitySum);

    return distribution;
  }

  /**
   * @return the internal prior class probability.
   */
  public DenseDoubleVector getClassProbability() {
    return this.classProbability;
  }

  /**
   * @return the internal probability matrix.
   */
  public DenseDoubleMatrix getProbabilityMatrix() {
    return this.probabilityMatrix;
  }

}

@@ -40,48 +40,8 @@
 s;%0A%0A
-import gnu.trove.set.hash.TIntHashSet;%0A%0A
 impo
@@ -1572,21 +1572,16 @@
     
-TIntHashSet d
+int numD
 isti
@@ -1597,137 +1597,16 @@
 s = 
-new TIntHashSet();%0A    int%5B%5D classes = new int%5Boutcome.length%5D;%0A    for (int i = 0; i %3C outcome.length; i++) %7B%0A      if (
 outcome%5B
 i%5D.g
@@ -1597,25 +1597,25 @@
 s = outcome%5B
-i
+0
 %5D.getDimensi
@@ -1622,188 +1622,65 @@
 on()
- == 1) %7B%0A        classes%5Bi%5D = (int) outcome%5Bi%5D.get(0);%0A      %7D else %7B%0A        classes%5Bi%5D = outcome%5Bi%5D.maxIndex();%0A      %7D%0A      distinctClasses.add(classes%5Bi%5D);%0A    %7D%0A    final int
+;%0A    // respect the binary case%0A    numDistinctClasses =
  num
@@ -1696,18 +1696,30 @@
 lasses =
- d
+= 1 ? 2 : numD
 istinctC
@@ -1728,15 +1728,8 @@
 sses
-.size()
 ;%0A  
@@ -2200,22 +2200,16 @@
 ;%0A      
-final 
 int pred
@@ -2225,29 +2225,145 @@
 s = 
-classes%5BcolumnIndex%5D;
+outcome%5BcolumnIndex%5D.maxIndex();%0A      if (numDistinctClasses == 2) %7B%0A        predictedClass = (int) outcome%5BcolumnIndex%5D.get(0);%0A      %7D
 %0A   

remove ancient caching logic