package xyz.lejon.utils;

import cc.mallet.classify.Classification;
import cc.mallet.classify.Trial;
import cc.mallet.classify.evaluate.ConfusionMatrix;
import cc.mallet.types.Instance;
import cc.mallet.types.LabelAlphabet;
import cc.mallet.types.LabelVector;
import cc.mallet.types.Labeling;


public class EnhancedConfusionMatrix extends ConfusionMatrix {
	protected String MATRIX_TYPE = "Confusion Matrix";
	protected int numClasses;
	// Since it is package private in parent...
	protected Trial myTrial;
	protected int [][] myValues;
	protected int numCorrect = 0;
	protected int total = 0;
	protected double averageAccuracy = 0.0;
	
	public EnhancedConfusionMatrix(Trial t) {
		super(t);
		
		this.myTrial = t;
		Labeling tempLabeling =
			((Classification)t.get(0)).getLabeling();
		this.numClasses = tempLabeling.getLabelAlphabet().size();
		
		myValues = new int[numClasses][numClasses];
		for(int i=0; i < t.size(); i++)
		{
			LabelVector lv =
				((Classification)t.get(i)).getLabelVector();
			Instance inst = ((Classification)t.get(i)).getInstance();
			int bestIndex = lv.getBestIndex();
			int correctIndex = inst.getLabeling().getBestIndex();
			assert(correctIndex != -1);
			//System.out.println("Best index="+bestIndex+". Correct="+correctIndex);
			myValues[correctIndex][bestIndex]++;
			total++;
		}		
		
		for (int i = 0; i < t.size(); i++)
			if (t.get(i).bestLabelIsCorrect())
				numCorrect++;
		
		averageAccuracy = (double) numCorrect / total;
	}
	
	public EnhancedConfusionMatrix(Trial [] ts) {
		super(ts[0]);
		
		MATRIX_TYPE = "Combined Confusion Matrix";
		
		Trial ttmp = ts[0];
		this.myTrial = ttmp;
		Labeling tempLabeling =
				((Classification)ttmp.get(0)).getLabeling();
		this.numClasses = tempLabeling.getLabelAlphabet().size();
		myValues = new int[numClasses][numClasses];
		
		for (Trial t : ts) {
			for(int i=0; i < t.size(); i++)
			{
				LabelVector lv =
					((Classification)t.get(i)).getLabelVector();
				Instance inst = ((Classification)t.get(i)).getInstance();
				int bestIndex = lv.getBestIndex();
				int correctIndex = inst.getLabeling().getBestIndex();
				assert(correctIndex != -1);
				//System.out.println("Best index="+bestIndex+". Correct="+correctIndex);
				myValues[correctIndex][bestIndex]++;
				total++;
				if (t.get(i).bestLabelIsCorrect())
					numCorrect++;
			}					
		}
		averageAccuracy = (double) numCorrect / total;
	}
	
	public String toCsv(String sep) {
		StringBuffer sb = new StringBuffer ();
		int maxLabelNameLength = 0;
		LabelAlphabet labelAlphabet = myTrial.getClassifier().getLabelAlphabet();
		for (int i = 0; i < numClasses; i++) {
			int len = labelAlphabet.lookupLabel(i).toString().length();
			if (maxLabelNameLength < len)
				maxLabelNameLength = len;
		}

		sb.append ("Label (R=true C=Predicted)" + sep);
		for (int c2 = 0; c2 < numClasses; c2++) sb.append (labelAlphabet.lookupLabel(c2).toString() + sep);
		sb.append ("total\n");
		for (int c = 0; c < numClasses; c++) {
			String labelName = labelAlphabet.lookupLabel(c).toString();
			sb.append (labelName + sep);
			for (int c2 = 0; c2 < numClasses; c2++) {
				sb.append (myValues[c][c2] + sep);
			}
			sb.append (MatrixOps.sum(myValues[c]));
			sb.append ('\n');
		}
		sb.append (sep);
		int [] colsums = MatrixOps.colSum(myValues);
		for (int c = 0; c < numClasses; c++) {
			sb.append(colsums[c] + sep);
		}
		sb.append (MatrixOps.sum(myValues));
		sb.append ('\n');
		return sb.toString(); 
	}
	
	@Override
	public String toString() {
		StringBuffer sb = new StringBuffer ();
		int maxLabelNameLength = 0;
		LabelAlphabet labelAlphabet = myTrial.getClassifier().getLabelAlphabet();
		for (int i = 0; i < numClasses; i++) {
			int len = labelAlphabet.lookupLabel(i).toString().length();
			if (maxLabelNameLength < len)
				maxLabelNameLength = len;
		}

		sb.append (MATRIX_TYPE + ", row=true, column=predicted  accuracy=" + averageAccuracy + " Correct=" + numCorrect + "\n");
		for (int i = 0; i < maxLabelNameLength-5+4; i++) sb.append (' ');
		sb.append ("label");
		for (int c2 = 0; c2 < Math.min(10,numClasses); c2++)	sb.append ("   "+c2);
		for (int c2 = 10; c2 < numClasses; c2++)	sb.append ("  "+c2);
		sb.append ("  |total\n");
		for (int c = 0; c < numClasses; c++) {
			appendJustifiedInt (sb, c, false);
			String labelName = labelAlphabet.lookupLabel(c).toString();
			for (int i = 0; i < maxLabelNameLength-labelName.length(); i++) sb.append (' ');
			sb.append (" "+labelName+" ");
			for (int c2 = 0; c2 < numClasses; c2++) {
				appendJustifiedInt (sb, myValues[c][c2], true);
				sb.append (' ');
			}
			sb.append (" |"+ MatrixOps.sum(myValues[c]));
			sb.append ('\n');
		}
		sb.append ("----");
		for (int i = 0; i < maxLabelNameLength; i++) sb.append ('-');
		sb.append ("+");
		for (int c = 0; c < numClasses; c++) {
			sb.append ("----");
		}
		sb.append (" |");
		sb.append ('\n');
		sb.append ("     ");
		int [] colsums = MatrixOps.colSum(myValues);
		for (int i = 0; i < maxLabelNameLength; i++) sb.append (' ');
		for (int c = 0; c < numClasses; c++) {
			appendJustifiedInt (sb, colsums[c], true);
			sb.append (' ');
		}
		sb.append (" |"+ MatrixOps.sum(myValues));

		return sb.toString(); 
	}
	
	// So, I'll copy paste this since it is private...
	static private void appendJustifiedInt (StringBuffer sb, int i, boolean zeroDot) {
		if (i < 100) sb.append (' ');
		if (i < 10) sb.append (' ');
		if (i == 0 && zeroDot)
			sb.append (".");
		else
			sb.append (""+i);
	}
	
	/**
	 * Assumes square <code>myValues</code> matrix with corrects on top-left to 
	 * lower-right diagonal
	 * 
	 * @param sep
	 * @param labelAlphabet
	 * @param myValues
	 * @return
	 */
	public static String enhancedConfusionMatrixToCsv(String sep, LabelAlphabet labelAlphabet, int [][] myValues) {
		int numClasses = myValues.length;
		StringBuffer sb = new StringBuffer ();
		int maxLabelNameLength = 0;
		for (int i = 0; i < numClasses; i++) {
			int len = labelAlphabet.lookupLabel(i).toString().length();
			if (maxLabelNameLength < len)
				maxLabelNameLength = len;
		}

		sb.append ("Label (R=true C=Predicted)" + sep);
		for (int c2 = 0; c2 < numClasses; c2++) sb.append (labelAlphabet.lookupLabel(c2).toString() + sep);
		sb.append ("total\n");
		for (int c = 0; c < numClasses; c++) {
			String labelName = labelAlphabet.lookupLabel(c).toString();
			sb.append (labelName + sep);
			for (int c2 = 0; c2 < numClasses; c2++) {
				sb.append (myValues[c][c2] + sep);
			}
			sb.append (MatrixOps.sum(myValues[c]));
			sb.append ('\n');
		}
		sb.append (sep);
		int [] colsums = MatrixOps.colSum(myValues);
		for (int c = 0; c < numClasses; c++) {
			sb.append(colsums[c] + sep);
		}
		sb.append (MatrixOps.sum(myValues));
		sb.append ('\n');
		return sb.toString(); 
	}
	
	/**
	 * Assumes square <code>myValues</code> matrix with corrects on top-left to 
	 * lower-right diagonal
	 * 
	 * @param sep
	 * @param labelAlphabet
	 * @param myValues
	 * @return
	 */
	public static String enhancedConfMatrixToString(int numClasses, LabelAlphabet labelAlphabet, String MATRIX_TYPE, 
			int [][] myValues) {
		int numCorrect = 0;
		int total = 0;
		for (int i = 0; i < myValues.length; i++) {
			for (int j = 0; j < myValues.length; j++) {
				if(i==j) {
					numCorrect += myValues[i][j];
				}
				total += myValues[i][j];
			}
		}
		double averageAccuracy = ((double) numCorrect) / total;
		StringBuffer sb = new StringBuffer ();
		int maxLabelNameLength = 0;
		for (int i = 0; i < numClasses; i++) {
			int len = labelAlphabet.lookupLabel(i).toString().length();
			if (maxLabelNameLength < len)
				maxLabelNameLength = len;
		}

		sb.append (MATRIX_TYPE + ", row=true, column=predicted  accuracy=" + averageAccuracy + " Correct=" + numCorrect + "\n");
		for (int i = 0; i < maxLabelNameLength-5+4; i++) sb.append (' ');
		sb.append ("label");
		for (int c2 = 0; c2 < Math.min(10,numClasses); c2++)	sb.append ("   "+c2);
		for (int c2 = 10; c2 < numClasses; c2++)	sb.append ("  "+c2);
		sb.append ("  |total\n");
		for (int c = 0; c < numClasses; c++) {
			appendJustifiedInt (sb, c, false);
			String labelName = labelAlphabet.lookupLabel(c).toString();
			for (int i = 0; i < maxLabelNameLength-labelName.length(); i++) sb.append (' ');
			sb.append (" "+labelName+" ");
			for (int c2 = 0; c2 < numClasses; c2++) {
				appendJustifiedInt (sb, myValues[c][c2], true);
				sb.append (' ');
			}
			sb.append (" |"+ MatrixOps.sum(myValues[c]));
			sb.append ('\n');
		}
		sb.append ("----");
		for (int i = 0; i < maxLabelNameLength; i++) sb.append ('-');
		sb.append ("+");
		for (int c = 0; c < numClasses; c++) {
			sb.append ("----");
		}
		sb.append (" |");
		sb.append ('\n');
		sb.append ("     ");
		int [] colsums = MatrixOps.colSum(myValues);
		for (int i = 0; i < maxLabelNameLength; i++) sb.append (' ');
		for (int c = 0; c < numClasses; c++) {
			appendJustifiedInt (sb, colsums[c], true);
			sb.append (' ');
		}
		sb.append (" |"+ MatrixOps.sum(myValues));

		return sb.toString(); 
	}

}
