package alg;

import generator.GraphGenerator;
import graph.Graph;
import graph.TimeGraph;

import java.io.BufferedWriter;
import java.io.FileWriter;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Random;

import misc.Common;

// Clusters the nodes of 
public class SAClusterer
{
	private static final double INIT_TEMP = 100.0;
	private static final double TEMP_DELTA = 0.00025;
	private static final double MIN_TEMP = 0.00001;
	private static final double LOW_TEMP = 0.001;
	private static final int MAX_ITERS = 100000;
	
	public static int DEBUG_LEVEL = 0;
	
	// ========================================================================
	
	private final int K;
	/*
	 * Annealing parameters
	 */
	private final double initialTemp;
	private final double tempDelta;
	private final int maxIters;
	private final double minTemp;
	private final boolean isGreedy;
	
	private double globalEnergy;
	private Graph graph;
	private int[] labels;
	private int f_coclust = 0;
	private int total_coclust = 0;

	// ========================================================================
	
	public SAClusterer(Graph graph_in, int k)
	{
		this(graph_in, k, INIT_TEMP, MIN_TEMP, TEMP_DELTA, MAX_ITERS, false);
	}
	
	// ========================================================================
	
	public SAClusterer(Graph graph_in, int k, boolean run_greedy)
	{
		this(graph_in, k, INIT_TEMP, MIN_TEMP, TEMP_DELTA, MAX_ITERS, run_greedy);
	}
	
	// ========================================================================
	
	public SAClusterer(Graph graph_in, int k, 
	                   double init_temp, double min_temp, double temp_delta,
	                   int iters, boolean run_greedy)
	{
		// we label 0..K,
		K = k;
		graph = graph_in;
		
		// SA parameters
		initialTemp = init_temp;
		minTemp = min_temp;
		tempDelta = temp_delta;
		maxIters = iters;
		
		globalEnergy = 0.0;
		isGreedy = run_greedy;
		
		f_coclust = 0;
		total_coclust = 0;
	}

	// ========================================================================
	
	public void printMLClusterLabels(HashMap<String, Double[]> nodeProbs)
	{
		Iterator<String> iter = nodeProbs.keySet().iterator();
		int i;
		String name;
		double[] max = new double[] {0.0, 0.0};
		Double[] p;
		
		while (iter.hasNext())
		{
			name = iter.next();
			p = nodeProbs.get(name);
			
			System.out.format("%20s\t", name);
			max[0] = 0.0;
			max[1] = 0.0;
			for (i = 0; i < K; ++i)
			{
				if (p[i] > max[0])
				{
					max[0] = p[i];
					max[1] = (double) i;
				}
			}
			System.out.format("%d", (int)max[1]);
			System.out.println();
		}
	}
	
	// ========================================================================
	
	public void printProbClusterLabels(HashMap<String, Double[]> nodeProbs)
	{
		Iterator<String> iter = nodeProbs.keySet().iterator();
		int i;
		String name;
		Double[] p;
		
		while (iter.hasNext())
		{
			name = iter.next();
			p = nodeProbs.get(name);
			
			System.out.format("%20s\t", name);
			for (i = 0; i < K; ++i)
			{
				System.out.format("%d:%.3f, ", i, p[i].doubleValue());	
			}
			System.out.println();
		}
	}
	
	// ========================================================================
	
	public void clusterSampler()
	{
		HashMap<String, Integer> clustSamples = new HashMap<String, Integer>(); 
		ArrayList<Double> energy = new ArrayList<Double>(); 
		int idx, iter = 0;
		double temperature;
		double[] p = null;
		
		/*
		 * Initialize
		 */
		globalEnergy = initializeGraphLabels();
		temperature = 10.0;
		/*
		 * Main loop for clustering
		 */
		do
		{
			temperature = 0.5 / java.lang.Math.log(2 + iter);
//			temperature - (0.5 * tempDelta * temperature);
			if (temperature < minTemp)
				temperature = minTemp;
			
			for (int n= 0 ; n < graph.getn(); n++)
			{
				/*
				 * Compute distribution over possible node relabelings
				 * Sample from that set and reset this node's label
				 */
				p = getLabelDistribution(n, temperature);
				idx = sample(p);
				labels[n] = idx;
				if ((iter > 10000)) {
					for (int i = graph.ind[n]; i < graph.ind[n+1]; i++)
					if (labels[graph.endv[i]] == labels[n])
						++f_coclust;
					++total_coclust;
				}
			}
			++iter;
			
			/*
			 * --------------------------- DEBUG
			 */
			if ((iter > 10000) && ((iter % 10) == 0))
			{
				int label;
				String clustState = "";
				
				relabelClusters();
				for (int i=0; i < graph.getn(); i++) clustState += labels[i] + ",";
				
				if (! clustSamples.containsKey(clustState))
					clustSamples.put(clustState, 1);
				else
				{
					System.out.println("Incrementing Prev state!");
					int count = clustSamples.get(clustState);
					clustSamples.put(clustState, count + 1);
				}
			}
			
			
			if ((iter % 100) == 0)
			{
				computeGlobalEnergy();
				energy.add(globalEnergy);
			}
			
			if ((iter % 2000) == 0)
			{
				System.out.println("ITERATION    : " + iter);
				System.out.println("Global Energy: " + globalEnergy);
				System.out.println("Temperature  : " + temperature);
			}
			globalEnergy = 0.0;
		}
		while (iter < maxIters);
		
		computeGlobalEnergy();

		System.out.println("\n\n-------------------\n");
       	System.out.println("Done...");
		System.out.println("Total iterations: " + iter);
		System.out.println("Global energy   : " + globalEnergy);
		System.out.println("\n");
		
		//Petko: not needed here to do consistency
		//graph.printLableConsistency(this.K);
	}
	
	// ========================================================================
	
	public int[] clusteringTrials(int trials)
	{
		int[][] clusterings = new int[graph.getn()][trials];
		
		int[] temp = null;
		int[] bestClust = null;
		double bestGlobalEnergy = Double.MAX_VALUE;
		int i;
		
		/*
		 * Run the clustering for a set number of trials,
		 * mark the frequency that each pair of users is put 
		 * into the same cluster
		 */
		for (i = 0; i < trials; ++i)
		{
			temp = cluster();
			for (int j = 0; j < graph.getn(); j++){
				clusterings[j][i] = temp[j];
			}
			if (globalEnergy < bestGlobalEnergy)
			{
				bestClust = temp;
				bestGlobalEnergy = globalEnergy;
			}
		}
		globalEnergy = bestGlobalEnergy;
		System.out.println("\n\nBest of " + trials + " trials: " + globalEnergy);
		
		return bestClust;
	}
	
	// ========================================================================
	
	public int[] bestClustering(int trials)
	{	
		int[] temp = null;
		int[] bestClust = null;
		double bestGlobalEnergy = Double.MAX_VALUE;
		int i;
		
		for (i = 0; i < trials; ++i)
		{
			System.out.println("Trial: " + i);
			temp = cluster();
			if (globalEnergy < bestGlobalEnergy)
			{
				bestClust = temp;
				bestGlobalEnergy = globalEnergy;
			}
		}
		globalEnergy = bestGlobalEnergy;
		System.out.println("\n\nBest of " + trials + " trials: " + globalEnergy);
		
		return bestClust;
	}

	// ========================================================================
	
	public void relabelClusters()
	{
		HashMap<Integer, Integer> labelReplacement = new HashMap<Integer, Integer>();
		int idx = 0;
		for (int i=0; i < graph.getn(); i++) {
			if(!labelReplacement.containsKey(labels[i])){
				labelReplacement.put(labels[i], idx++);
			}
		}
		for (int i=0; i < graph.getn(); i++) 
			labels[i] = labelReplacement.get(labels[i]);
	}
	
	// ========================================================================
	
//	private ArrayList<String> randomizeNodeList(ArrayList<String> nodeIds)
//	{
//		ArrayList<String> ids = new ArrayList<String>();
//		Random rnd = new Random();
//		int idx;
//		
//		while (nodeIds.size() > 0)
//		{
//			idx = rnd.nextInt(nodeIds.size());
//			ids.add(nodeIds.remove(idx));
//		}
//		return ids;
//	}

	// ========================================================================
	
	public int[] cluster()
	{
		final int MAX_CONV_COUNT = 10;
		int[] best_labels = new int[graph.getn()];
		ArrayList<Double> energy = new ArrayList<Double>();
		int i, iter = 0, convCount = 0, acceptedCount;
		double bestEnergy = Double.MAX_VALUE, temperature;
		boolean isConverged = false;

		/*
		 * Initialize
		 */
		globalEnergy = initializeGraphLabels(false);
		temperature = initialTemp;
		
		/*
		 * Main loop for clustering
		 */
		while ((iter < maxIters) && (! isConverged)) 
		{
			temperature = temperature - tempDelta * temperature;
			if (temperature < minTemp)
			{
				temperature = minTemp;
				isConverged = true;
			}

			if (convCount >= MAX_CONV_COUNT)
				isConverged = true;
			
			/*
			 * Iterate through the set of nodes and propose a change for each one
			 */
			acceptedCount = 0;
			for (i = 0; i < graph.getn(); ++i)
			{
				/*
				 * Proposal to change the current clustering state
				 * and decide whether or not to keep it
				 */
				if (temperature <= LOW_TEMP) 
				{
					if (proposeBestChange(i, temperature))
						++acceptedCount;
				}
				else
				{
					if (proposeChange(i, temperature))
						++acceptedCount;
				}
			}
			/*
			 * Recompute the global energey
			 */
			computeGlobalEnergy();
			energy.add(globalEnergy);
			
			/*
			 * Keep track of best solution seen so far
			 */
			if (globalEnergy < bestEnergy)
			{
				bestEnergy = globalEnergy;
				best_labels = Arrays.copyOf(labels, graph.getn());
			}
			
			++iter;
			if ((DEBUG_LEVEL > 0) && ((iter % 1000) == 0))
			{
				System.out.println("ITERATION    : " + iter);
				System.out.println("Temperature  : " + temperature); 
				System.out.println("Global Energy: " + globalEnergy);
				System.out.println("Best Energy  : " + bestEnergy);
			}

			if ((DEBUG_LEVEL > 0) && (((iter % 2000) == 0) || (isConverged)))
				System.out.format("(%7d): accepted moves %6d / %6d\ttemperature %.6f\n", 
				                  iter, acceptedCount, graph.getn(), temperature);
			
			if (((double)acceptedCount / graph.getn()) <= 0.01)
				++convCount;
			else
				convCount = 0;
		}
		
		// compute energy
		computeGlobalEnergy();
		energy.add(globalEnergy);
		
		System.out.println("\n------------------------------------------");
		System.out.println("Total iterations: " + iter);
		System.out.println("Global energy   : " + globalEnergy);
		
		relabelClusters();
		return best_labels;
	}
	
	// ========================================================================
	
	/*
	 * This function is 'more deterministic' as it searches over all
	 * possible states and picks the one that gives the best improvement.
	 * This function is only good for optimizing (not exploring) and 
	 * is only called when the temperature is low 
	 */
	private boolean proposeBestChange(int n, double temp)
	{
		boolean isAccept = false;
		Random rnd = new Random();
		int propLabel, bestLabel;
		double[] propNegCount = new double[K];
		double curNegCount, bestNegCount;

		curNegCount = getNegCount(n, labels[n]);
		bestNegCount = curNegCount;
		bestLabel = labels[n];
		
		/*
		 * Iterate through the set of proposed changes
		 */
		for (propLabel = 0; propLabel < K; ++propLabel)
		{
			propNegCount[propLabel] = getNegCount(n, propLabel);
			if (propNegCount[propLabel] < bestNegCount)
			{
				bestLabel = propLabel;
				bestNegCount = propNegCount[propLabel];
			}
		}

		if (bestNegCount == curNegCount)
		{
			// don't change...
		}		
		else if (bestNegCount < curNegCount)
		{
			// accept change
			labels[n] = bestLabel;
			isAccept = true;
		}
		else if (!isGreedy)
		{
			double prob, urand;
			
			/*
			 * Do simulated annealing 
			 */
			// accept change with probability
			prob = java.lang.Math.exp((curNegCount - bestNegCount) / temp);
			urand = rnd.nextDouble();
			
			if (prob >= urand)
			{
				labels[n] = bestLabel;
				isAccept = true;
			}
		}
		return isAccept;
	}
	
	// ========================================================================
	
	private boolean proposeChange(int n, double temp)
	{
		boolean isAccept = false;
		Random rnd = new Random();
		int propLabel;
		double curNegCount, propNegCount, prob, urand;
		
		propLabel = getSampleExcludeVal(labels[n]);
		
		curNegCount = getNegCount(n, labels[n]);
		propNegCount = getNegCount(n, propLabel);

		if (propNegCount < curNegCount)
		{
			// accept change
			labels[n] = propLabel;
			isAccept = true;
		}
		else if (! isGreedy)
		{
			/*
			 * Do simulated annealing 
			 */
			// accept change with probability
			prob = java.lang.Math.exp((curNegCount - propNegCount) / temp);
			urand = rnd.nextDouble();
			
			if (prob >= urand)
			{
				labels[n] = propLabel;
				isAccept = true;
			}
		}
		return isAccept;
	}
	
	// ========================================================================
	
	private void computeGlobalEnergy()
	{
		globalEnergy = 0.0;
		for (int i = 0; i < graph.getn(); i++){
			globalEnergy += getLocalEnergy(i);
		}
	}
	
	// ========================================================================
	
	public int getSampleExcludeVal(int exclude)
	{
		Random rnd = new Random();
		int sample;
		
		if (exclude == 0)
		{
			sample = rnd.nextInt(K - 1) + 1;
		}
		else if (exclude == (K - 1))
		{
			sample = rnd.nextInt(K - 1);
		}
		else
		{
			double p = (double) exclude / (double)K;
			if (p >= rnd.nextDouble())
			{
				// sample from 0..exclude
				sample = rnd.nextInt(exclude);
			}
			else
			{
				// sample from (exclude + 1)..K
				sample = rnd.nextInt(K - exclude - 1) + exclude + 1;
			}
		}
		return sample;
	}
	
	// ========================================================================
	
	public void printNodeLabelingToFile(HashMap<String, Integer> labels, String out)
	{
		Iterator<String> iter = null;
		String node;
		BufferedWriter br;
		
		try
		{
			br = new BufferedWriter(new FileWriter(out));
		
			iter = labels.keySet().iterator();
			while (iter.hasNext())
			{
				node = iter.next();
				br.write(node + Common.SEP + labels.get(node) + "\n");
			}
			br.close();
		}
		catch(Exception e)
		{
			e.printStackTrace();
			System.exit(1);
		}
	}
	
	// ========================================================================
	
	public void printNodeLabeling(HashMap<String, Integer> labels)
	{
		Iterator<String> iter = null;
		String node;
		int i;
		
		for (i = 0; i < K; ++i)
		{
			iter = labels.keySet().iterator();
			while (iter.hasNext())
			{
				node = iter.next();
				if (labels.get(node) == i)
					System.out.format("%20s\t%3d\n", node, i);
			}
		}
	}
	
	// ========================================================================
	
	private int sample(double[] dist)
	{
		Random x = new Random();
		double v, sum = 0.0;
		int idx = 1;

		sum = dist[0];
		v = x.nextDouble();
		while ((sum < v) && (idx < dist.length))
		{
			sum += dist[idx];
			++idx;
		}
		--idx;
		return idx;
	}

	// ========================================================================
	
	private double[] getLabelDistribution(int n, double temp)
	{
		int i;
		double[] p = new double[K];
		double sum = 0.0, curCount = getNegCount(n, labels[n]);
		
		// compute scores
		for (i = 0; i < K; ++i)
		{
			p[i] = scoreNode(n, i, curCount, temp);
			sum += p[i];
		}
		
		if (sum == 0)
		{
			// normalize
			for (i = 0; i < K; ++i)
				p[i] = 1.0 / K;
		}
		else
		{
			// normalize
			for (i = 0; i < K; ++i)
				p[i] /= sum;
		}
		return p;
	}

	// ========================================================================
	
	private double scoreNode(int n, int label, double curCount, double temp)
	{
		double negCount = getNegCount(n, label);
		return java.lang.Math.exp(-(negCount - curCount) / temp);
	}
	
	// ========================================================================
	
	private double getLocalEnergy(int n)
	{
		return getLocalEnergy(n, labels[n]);
	}
	
	// ========================================================================
	
	private double getLocalEnergy(int n, int label)
	{
		double negCount = getNegCount(n, label);
		return negCount;
	}
	
	// ========================================================================
	
	/*
	 * Compute the count of bad edges vs good edges for a node.
	 * A Higher value for negCount means this is a poor configuration.
	 * Low (or even negative) values mean this is a good configuration.
	 */
	private double getNegCount(int node, int label)
	{
		double incVal;
		int lab;
		double negCount = 0.0;

		for (int i = graph.ind[node]; i < graph.ind[node+1]; i++)
		{
			lab = getLabel(graph.endv[i]);
			incVal = java.lang.Math.abs(graph.we[i]);
			
			if (label == lab)
			{
				if ( graph.we[i] < 0.0)
					negCount += incVal;
				else
					negCount -= incVal;
			}
			else
			{
				if (graph.we[i] >= 0.0)
					negCount += incVal;
				else
					negCount -= incVal;
			}
		}
		return negCount;
	}
	
	private int getLabel(int node) {
		return labels[node];
	}

	// ========================================================================
	
	private double initializeGraphLabels()
	{
		return initializeGraphLabels(true);
	}
	
	// ========================================================================
	
	private double initializeGraphLabels(boolean resetCoClustCounts)
	{
		double energy = 0.0;
		Random r = new Random();
		labels  = new int[graph.getn()];
		for (int i=0; i < graph.getn(); i++) labels[i] = r.nextInt(K);
		for (int i=0; i < graph.getn(); i++) energy += getLocalEnergy(i);
		return energy;
	}
	
	// ========================================================================
	
	public double optimalEnergy()
	{	
		double optimalEnergy = Double.MAX_VALUE;
		int bestCounter = 0;
		int idx, label, counter = 0;
		
		if (K > 2)
			return 0;
		
		while (counter < java.lang.Math.pow(2, graph.getn() - 1))
		{
			for (int i=0; i < graph.getn(); i++)
			{
				label = (counter & (int)java.lang.Math.pow(2.0, i));
				if (label != 0)
					labels[i] = 1;
				else
					labels[i] = 0;
			}
			// compute global energy
			computeGlobalEnergy();
		
			// save best so far
			if (globalEnergy < optimalEnergy)
			{
				optimalEnergy = globalEnergy;
				bestCounter = counter;
			}
			++counter;
		}
	
		// reset to optimal configuration
		for (int i=0; i < graph.getn(); i++)
		{
			label = (bestCounter & (int)java.lang.Math.pow(2.0, i));
			
			if (label != 0)
				labels[i] = 1;
			else
				labels[i] = 0;
		}
		
		for (int i = 0; i < graph.getn(); i++)
		{
			System.out.println(i + "\t" + labels[i]);
		}
		
		return optimalEnergy;
	}
	
	// ========================================================================
	
	public double getGlobalEnergy()
	{
		return globalEnergy;
	}
	
	// ========================================================================
	
	public Graph getGraph()
	{
		return graph;
	}
	
	// ========================================================================
	
	public static void main(String[] Args) {
		// first probability is persist, second seed
////	Graph g = GraphGenerator.generateScaleFree(300);
//	Graph g = GraphGenerator.generateRandom(100, 190);
////	double Pn = 0.005, Pp = 0.05, Ps = 0.001; // settings
////	TimeGraph tg = generateTimeGraphLocalized(40, Pn, Pp, Ps, g);
		Graph g = GraphGenerator.generateRandom(200, 500);
		double N = 0.3, S = 0.1;
		TimeGraph tg = GraphGenerator.generateTimeHMM(40, N, S, g);
		tg.aggregateEdgeWeights(3, 8, g.we);
		SAClusterer sac = new SAClusterer(g, 2);
		int[] a = sac.cluster();
	}
	
}




