import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.InputStreamReader;
import java.io.OutputStreamWriter;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Random;
import java.util.TreeMap;


//implements representation testing for EECS 395/495: Probabilistic Graphical Models

public class RepresentationTester {

	private HashMap<String, ArrayList<String>> _instances;
	
	private HashMap<String, ArrayList<double []>> _doubleReps; // holds continuous-vector word representations 
		
	private Random _rand;
	
	
	
	public RepresentationTester() {
		_instances = null;
	}
	
	public RepresentationTester(String instancesDir) throws Exception {
		_instances = readInstances(instancesDir);
	}

	public RepresentationTester(String instancesDir, boolean limitOneWord) throws Exception {
		_instances = readInstances(instancesDir, limitOneWord);
	}
	
	private static double precOf(int [] cts) {
		return cts[1] / (double)(cts[0] + cts[1]);
	}
	
	
	public void initDoubleReps(String repsFile) throws Exception {
		HashMap<String, double[]> wordReps = new HashMap<String, double[]>();
		String sLine;
		BufferedReader brIn = new BufferedReader(new InputStreamReader(new FileInputStream(repsFile), "UTF8"));
		
		int numDims = 0;
		int ct=0;
		while((sLine = brIn.readLine())!=null) {
			String [] wordVec = sLine.split("\t| ");
			String word = wordVec[0].replace("CN_", "");
			double [] dims = new double[wordVec.length-1];
			for(int i=1; i<wordVec.length;i++) {
				dims[i-1] = Double.parseDouble(wordVec[i]);
			}
			numDims = dims.length;
			wordReps.put(word, dims);
			if(++ct % 100 == 0)
				System.out.println(ct);
		}
		brIn.close();
		
		double [] backOff = new double[numDims];
		for(int i=0; i<backOff.length; i++) {
			backOff[i] = 1.0 / (double) numDims;
		}
		
		_doubleReps = new HashMap<String, ArrayList<double[]>>();
		boolean success = true;
		for(String s : _instances.keySet()) {
			ArrayList<String> is = _instances.get(s);
			_doubleReps.put(s, new ArrayList<double[]>());
			ArrayList<double[]> al = _doubleReps.get(s);
			for(int i=0; i<is.size(); i++) {
				String toFind = is.get(i);
				if(!wordReps.containsKey(toFind)) {
					System.err.println("Error: could not find representation for " + s + " " + toFind);
					al.add(backOff);
					success = false;
				}
				else
					al.add(wordReps.get(toFind));
			}
		}
		if(success)
			System.out.println("Found all representations.");
		else
			System.out.println("Failed to find all representations.");
	}
	
	//takes in instance strings (for debugging output), scores, and +/- designators (clss)
	//returns AUC
	public static double computeAuc(ArrayList<String> is, ArrayList<Double> scs, ArrayList<Boolean> clss) {
		TreeMap<Double, int []> tm = new TreeMap<Double, int[]>();
		  //in arrays a, a[0] is - instances with score, a[1] is + instances with score
		
		for(int i=0; i<scs.size(); i++) {
			String inst = is.get(i);
			double sc = scs.get(i);
			boolean cls = clss.get(i);
			if(!tm.containsKey(sc)) {
				tm.put(sc, new int[2]);
			}
			int [] counts = tm.get(sc);
			counts[cls?1:0]++;
		}
		Iterator<Double> tmIt = tm.keySet().iterator();
		double lastPrec = -1.0;
		double area = 0.0;
		int [] negPosSoFar = new int[2];
		while(tmIt.hasNext()) {
			double sc = tmIt.next();
			int [] negpos = tm.get(sc);
			if(lastPrec<0.0) {
				lastPrec = precOf(negpos);
			}
			negPosSoFar[0] += negpos[0];
			negPosSoFar[1] += negpos[1];
			double curPrec = precOf(negPosSoFar);
			area += negpos[1]*(lastPrec + curPrec) / 2.0;
			lastPrec = curPrec;
		}
		return area / (double)negPosSoFar[1];
	}
	
	public static boolean testAuc() {
		ArrayList<String> is = new ArrayList<String>();
		ArrayList<Double> scs = new ArrayList<Double>();
		ArrayList<Boolean> clss = new ArrayList<Boolean>();
		is.add("blah");
		scs.add(0.5);
		clss.add(true);
		is.add("blah");
		scs.add(0.5);
		clss.add(false);
		is.add("blah");
		scs.add(0.5);
		clss.add(true);
		is.add("blah");
		scs.add(0.25);
		clss.add(true);
		is.add("blah");
		scs.add(0.25);
		clss.add(true);
		is.add("blah");
		scs.add(0.2);
		clss.add(false);
		is.add("blah");
		scs.add(0.2);
		clss.add(false);
		double auc = computeAuc(is, scs, clss);
		System.out.println(auc + " should be " + (0.6667 * 2 + 0.6667+0.8)/4.0);
		is.clear();
		scs.clear();
		clss.clear();
		
		is.add("blah");
		scs.add(0.5);
		clss.add(true);
		is.add("blah");
		scs.add(0.5);
		clss.add(true);
		is.add("blah");
		scs.add(0.5);
		clss.add(true);
		is.add("blah");
		scs.add(0.25);
		clss.add(false);
		is.add("blah");
		scs.add(0.25);
		clss.add(false);
		is.add("blah");
		scs.add(0.2);
		clss.add(false);
		is.add("blah");
		scs.add(0.2);
		clss.add(false);
		auc = computeAuc(is, scs, clss);
		System.out.println(auc + " should be 1.0");
		return true;
	}
	
	public double Euc(double [] a, double [] b) {
		double ret = 0.0;
		for(int i=0; i<a.length; i++) {
			ret += (a[i] - b[i])*(a[i] - b[i]);
		}
		return ret;
	}
	
	public double KL(double [] a, double [] b) {
		double ret = 0.0;
		for(int i=0; i<a.length; i++) {
			ret += a[i] * Math.log(a[i]/b[i]);
		}
		return ret;
	}
	
	public double Cos(double [] a, double [] b) {
		double ret = 0.0;
		double aSum = 0.0;
		double bSum = 0.0;
		for(int i=0; i<a.length; i++) {
			ret += a[i] * b[i];
			aSum += a[i];
			bSum += b[i];
		}
		return 1.0 - ret / (aSum * bSum);
	}
	
	//returns hashmap with scores for each instance, for given seeds within the target class
	//if polysemy method is specified, uses class averaging method
	//if randomize, assigns random scores instead
	public HashMap<String, ArrayList<Double>> getScoresForSeedsDoubleReps(HashSet<String> seeds, String target, boolean polysemyMethod,
			boolean randomize) {
		HashMap<String, ArrayList<Double>> scores = new HashMap<String, ArrayList<Double>>();
		
		ArrayList<double []> seedVectors = new ArrayList<double[]>(); 
		
		//get seed distributions:
		for(String s : _instances.keySet()) {
			ArrayList<String> is = _instances.get(s);
			boolean cls = s.equals(target);
			for(int i=0; i<is.size(); i++) {
				if(cls && seeds.contains(is.get(i))) {
					seedVectors.add(_doubleReps.get(s).get(i));
				}
			}
		}
		
		if(!polysemyMethod || randomize) {
			double [] seedDist = new double[seedVectors.get(0).length];
			Iterator<double []> seedIt = seedVectors.iterator();
			while(seedIt.hasNext()) {
				double [] sd = seedIt.next();
				for(int i=0; i<seedDist.length; i++) {
					seedDist[i] += sd[i] / (double)seedVectors.size();
				}
			}
			for(String s : _instances.keySet()) {
				ArrayList<String> is = _instances.get(s);
				scores.put(s, new ArrayList<Double>());
				ArrayList<Double> toFill = scores.get(s);
				for(int i=0; i<is.size(); i++) {
					if(randomize)
						toFill.add(_rand.nextDouble());
					else
						toFill.add(Euc(seedDist, _doubleReps.get(s).get(i)));
				}
			}
		}
		else {
			double [] seedDist = new double[seedVectors.get(0).length];
			Iterator<double []> seedIt = seedVectors.iterator();
			double maxDist = Double.NEGATIVE_INFINITY;
			while(seedIt.hasNext()) {
				double [] sd = seedIt.next();
				for(int i=0; i<seedDist.length; i++) {
					seedDist[i] += Math.log(sd[i]);
					if(seedDist[i] > maxDist)
						maxDist = seedDist[i];
				}
			}
			//adjust by maxDist to avoid underflow, and exponentiate:
			for(int i=0; i<seedDist.length; i++) {
				seedDist[i] -= maxDist;
				seedDist[i] = Math.exp(seedDist[i]);
			}
			
			for(String s : _instances.keySet()) {
				ArrayList<String> is = _instances.get(s);
				scores.put(s, new ArrayList<Double>());
				ArrayList<Double> toFill = scores.get(s);
				for(int i=0; i<is.size(); i++) {
					//compute dot prod:
					double dp = 0.0;
					double [] reps = _doubleReps.get(s).get(i);
					for(int j=0; j<seedDist.length;j++)
						dp += seedDist[j]*reps[j];
					toFill.add(-1.0*dp);
				}
			}			
		}
		
		return scores;
	}
	
	//takes in scores of similarity computed by representation, returns AUC-PR 
	public double scoreRankings(HashMap<String, ArrayList<Double>> scores, HashSet<String> seeds, String target) {
		
		//create +/- set:
		
		ArrayList<String> alInstances = new ArrayList<String>();
		ArrayList<Double> alScores = new ArrayList<Double>();
		ArrayList<Boolean> alClass = new ArrayList<Boolean>();
		for(String s : scores.keySet()) {
			ArrayList<String> is = _instances.get(s);
			boolean cls = s.equals(target);
			for(int i=0; i<is.size(); i++) {
				if(!seeds.contains(is.get(i))) {
					alInstances.add(is.get(i));
					alScores.add(scores.get(s).get(i));
					alClass.add(cls);
				}
			}
		}
		
		double auc = computeAuc(alInstances, alScores, alClass);
		return auc;
	}
	
	public HashMap<String, ArrayList<String>> readInstances(String inDir) throws Exception {
		return readInstances(inDir, false);
	}
	
    public static ArrayList<String> readLinesToStringArrayList(File inFile) throws Exception {
    	BufferedReader brIn = new BufferedReader(new InputStreamReader(new FileInputStream(inFile), "UTF8"));
    	String sLine;
    	ArrayList<String> alLines = new ArrayList<String>();
    	while((sLine = brIn.readLine())!=null) {
    		alLines.add(sLine);
    	}
    	return alLines;
    }
	
	public HashMap<String, ArrayList<String>> readInstances(String inDir, boolean limitOneWordAndLCase) throws Exception {
		HashMap<String, ArrayList<String>> out = new HashMap<String, ArrayList<String>>();
		File [] fs = (new File(inDir)).listFiles();
		for(int i=0; i<fs.length; i++) {
			String fName = fs[i].getName();
			ArrayList<String> res=readLinesToStringArrayList(fs[i]);
			res.remove(0); //first entry is name of class
			if(limitOneWordAndLCase) {
				ArrayList<String> resFixed = new ArrayList<String>();
				Iterator<String> it = res.iterator();
				while(it.hasNext()) {
					String s = it.next();
					if(!s.contains("+")) {
						//resFixed.add(s.toLowerCase());
						resFixed.add(s);
					}
				}
				out.put(fName, resFixed);
			}
			else
				out.put(fName, res);
		}
		return out;
	}
	
	//selects num integers from set 0,...,size-1
	//goes slow if numToSelect is close to size
	public HashSet<Integer> selectIdxs(int num, int size) {
		if(num > size)
			return null;
		HashSet<Integer> out = new HashSet<Integer>(num);
		while(out.size() < num)
			out.add(_rand.nextInt(size));
		return out;
	}
	
	//executes tests with random seed selections, for each class.  Note, requires that
	//skips cls if numSeeds is as large as cls size (and goes slow if numSeeds is close to cls size)
	//also requires that reps be initialized 
	//if randomize, will execute a test with random orderings for instances
	public void executeRandomTests(int numSeeds, int numTrials, boolean randomize) {
		Iterator<String> clsIt = _instances.keySet().iterator();
		double avgAuc = 0.0;
		double runs = 0.0;
		while(clsIt.hasNext()) {
			String cls = clsIt.next();
			ArrayList<String> sInstances = _instances.get(cls);
			int numInstances = sInstances.size();
			double a=0.0;
			double r=0.0;
			if(numSeeds < numInstances) {
				System.out.print(cls);
				for(int t=0;t<numTrials; t++) {
					//pick seeds:
					HashSet<Integer> chosenIdxs = selectIdxs(numSeeds, numInstances);
					HashSet<String> seeds = new HashSet<String>(numSeeds);
					for(int i : chosenIdxs)
						seeds.add(sInstances.get(i));
					HashMap<String, ArrayList<Double>> scores = getScoresForSeedsDoubleReps(seeds, cls, false, randomize);
					double auc = scoreRankings(scores, seeds, cls);
					a += auc;
					r++;
				}
				avgAuc += a;
				runs += r;
				System.out.println("\t" + a/r);
			}
		}
		avgAuc /= runs;
		System.out.println("----------------------");
		if(randomize)
			System.out.println("Random average: " + avgAuc);
		else
			System.out.println("Representation average: " + avgAuc);
		System.out.println();
	}
	
	public static void makeSimpleRep(String inFile, String outFile) throws Exception {
		BufferedReader brIn = new BufferedReader(new InputStreamReader(new FileInputStream(inFile), "UTF8"));
		BufferedWriter bwOut = new BufferedWriter(new OutputStreamWriter(new FileOutputStream(outFile), "UTF8"));
		String sLine ="";
		while((sLine=brIn.readLine())!= null) {
			bwOut.write(sLine + "\t" + sLine.length() + "\r\n");
		}
		brIn.close();
		bwOut.close();
	}
	
	/**
	 * @param args
	 */
	public static void main(String[] args) throws Exception {
		if(args.length != 2) {
			System.out.println("Usage: java RepresentationTester <test directory> <representation file>");
			return;
		}

		RepresentationTester et = new RepresentationTester(args[0], true);
		System.out.println(et._instances.size() + " instances.");
		et.initDoubleReps(args[1]);
		et._rand = new Random(4);
		et.executeRandomTests(5, 30, false);
		et._rand = new Random(4);
		et.executeRandomTests(5, 30, true);
		
		//scraps:
//		makeSimpleRep("x:\\d\\teach\\395_Fall13\\project_files\\devset.txt",
//		"x:\\d\\teach\\395_Fall13\\project_files\\simplerep.txt");
		//ExpansionTester.testAuc();
	}


}