code: 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(); 
  	}
  
  }
  