package cluster;

import generator.GraphGenerator;
import graph.APSP;
import graph.DistanceMeasure;
import graph.Graph;
import graph.IGraph;
import graph.ITimeGraph;
import graph.NodeWeightedTimeGraph;
import graph.WeightedTimeGraph;

import java.io.Externalizable;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.BitSet;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashSet;
import java.util.PriorityQueue;
import java.util.Random;
import java.util.TreeSet;

import misc.ExtToolsDriver;
import misc.IO;

import cluster.Common.GDIST;

public class TreeSplitClustering {
	
	public static String ERF = "L1"; // "L1"
	
	/** BUDGET/MAXSSE CLUSTERING
	  * Spits until stop criteria
    **/
	public static SplitTreeNode scluster(IGraph g, DistanceMeasure dm, StopCriteria sc, boolean probeall) {
		BitSet filt = new BitSet(g.getn());
		filt.set(0, g.getn());
		float avg = avg(g, filt);
		float sse = err(g, filt); 
		SplitTreeNode root = new SplitTreeNode(filt,null, avg, sse);
		PriorityQueue<SplitTreeNode> Q = new PriorityQueue<SplitTreeNode>(g.getn(),Collections.reverseOrder());
		Q.add(root);
		if (probeall) splitBudgetProbeAll(Q, root, g, dm, sc);
		else splitBudget(Q,root, g, dm, sc);
		return root;
	}
	
	public static void updateSplits(SplitTreeNode root, IGraph g, DistanceMeasure dm, StopCriteria sc, boolean probeall) { 
		PriorityQueue<SplitTreeNode> Q = new PriorityQueue<SplitTreeNode>(g.getn(),Collections.reverseOrder());
		ArrayList<SplitTreeNode> toadd = new ArrayList<SplitTreeNode>();
		root.getAllNotSplitDecendants(toadd);
		Q.addAll(toadd);
		if (probeall) splitBudgetProbeAll(Q, root, g, dm, sc);
		else splitBudget(Q,root, g, dm, sc);
	}
	
	public static void splitBudgetProbeAll(PriorityQueue<SplitTreeNode> Q, SplitTreeNode root,  
			IGraph g, DistanceMeasure dm, StopCriteria sc) {
		if (sc.stop(root)) return;
		BitSet fl = null, fr = null;
		SplitTreeNode tosplit = null;
		int vpb=0;
		float deltab = 0,distb=0;
		
		for (SplitTreeNode c: Q) {
			int bestvp = -1;
			float bestd = -1;
			float minsse = Float.POSITIVE_INFINITY, tmpsse = -1;
			
			// First, find the best possible split
			// TODO(): Very expensive and to be optimized
			for (int vp=0; vp < g.getn(); vp++) {
				ArrayList<Double> vec = new ArrayList<Double>();
				for (float d = 0; d < dm.max()+1; d+=1) {
					fl = filter(c.els,vp,d,true,dm,g);
					fr = filter(c.els,vp,d,false,dm,g);
					if (fl.cardinality()*fr.cardinality() == 0) {
						vec.add(0.0);
						continue; // no split
					}
					tmpsse = err(g,fl) + err(g,fr);
					vec.add((double)tmpsse);
					if (minsse > tmpsse) {
						minsse = tmpsse;
						bestvp = vp;
						bestd = d;
					}
				}
//				if (vp==0) {
					System.err.print(vp + "\t" + root.numSplitDecendants());
					ExtToolsDriver.plotVector(vec,"" + root.numSplitDecendants());
					System.err.print("\n");
//				}
			}
			if (c.sse - minsse > deltab) {
				deltab = c.sse - minsse;
				tosplit = c;
				vpb = bestvp;
				distb = bestd;
			}
		}
		Q.remove(tosplit);
		tosplit.vantagepoint = vpb;
		tosplit.dist = distb;
		
		if (vpb >= 0) {
			// Left child (within)
			fl = filter(tosplit.els,vpb,distb,true,dm,g);
			tosplit.lc = new SplitTreeNode(fl, tosplit, avg(g,fl), err(g,fl));
			Q.add(tosplit.lc);
			
			// Right child (out)
			fr = filter(tosplit.els,vpb,distb,false,dm,g);
			tosplit.rc = new SplitTreeNode(fr, tosplit, avg(g,fr), err(g,fr));
			Q.add(tosplit.rc);
			
			if (!sc.stop(root)) splitBudgetProbeAll(Q, root, g, dm, sc);
		}
	}
	
	public static void splitBudget(PriorityQueue<SplitTreeNode> Q, SplitTreeNode root,  IGraph g, DistanceMeasure dm, StopCriteria sc) {
		if (sc.stop(root)) return;
		
		SplitTreeNode c = Q.poll();
		
		int bestvp = -1;
		float bestd = -1;
		float minsse = Float.POSITIVE_INFINITY, tmpsse = -1;
		BitSet fl = null, fr = null;
		// First, find the best possible split
		// TODO(): Very expensive and to be optimized
		for (int vp=0; vp < g.getn(); vp++) {
			ArrayList<Double> vec = new ArrayList<Double>();
			for (float d = 0; d < dm.max(); d+=1) {
				fl = filter(c.els,vp,d,true,dm,g);
				fr = filter(c.els,vp,d,false,dm,g);
				if (fl.cardinality()*fr.cardinality() == 0) continue; // no split
				tmpsse = err(g,fl) + err(g,fr);
				vec.add((double)tmpsse);
				if (minsse > tmpsse) {
					minsse = tmpsse;
					bestvp = vp;
					bestd = d;
				}
			}
			if (vp==0) {
				System.err.print(root.numSplitDecendants());
				ExtToolsDriver.plotVector(vec,"" + root.numSplitDecendants());
				System.err.print("\n");
			}
		}
		
		c.vantagepoint = bestvp;
		c.dist = bestd;
		
		if (bestvp >= 0) {
			// Left child (within)
			fl = filter(c.els,bestvp,bestd,true,dm,g);
			c.lc = new SplitTreeNode(fl, c, avg(g,fl), err(g,fl));
			Q.add(c.lc);
			
			// Right child (out)
			fr = filter(c.els,bestvp,bestd,false,dm,g);
			c.rc = new SplitTreeNode(fr, c, avg(g,fr), err(g,fr));
			Q.add(c.rc);
			
			if (!sc.stop(root)) splitBudget(Q, root, g, dm, sc);
		}
	}
	
	/** SPLIT then PRUNE
	 * Split-based clustering
	 * First splits till sse improvement is less than deltasse
	 * Then prune based on BIC
	 * TODO() one can stop splitting directly based on BIC
	 **/
	public static void timeGraphClustering(ITimeGraph g, DistanceMeasure dm, float deltasse) {
		
		// Tree splits data structure
		ArrayList<SplitTreeNode> roots = new ArrayList<SplitTreeNode>(); 
		SplitTreeNode tree = null;
		
		long start;
		
		// compute the split clustering for every slice  
		for (int t = 0; t < g.gett(); t++) {
			
			start = System.currentTimeMillis();
			
			// TODO(): extract experiment with budgeted
			// vary budget and compute permutation of weights
			
			// move node weights
			if (NodeWeightedTimeGraph.class.equals(g.getClass())){
				for (int i = 0; i < g.getn(); i++) 
					g.getGraph().wn[i] = g.getTimeseries()[t][i];
			} else if (WeightedTimeGraph.class.equals(g.getClass())) { // or edge weights
				for (int i = 0; i < g.getm(); i++) 
					g.getGraph().we[i] = g.getTimeseries()[t][i];
			}
			tree = scluster(g, dm, deltasse);
			//tree = sclusterBudget(tg, dm, 10);
			// add to the the roots array
			roots.add(tree);
			
			if (Common.DEBUG) {
				if (NodeWeightedTimeGraph.class.equals(g.getClass())){
					for(int i=0 ; i < g.getn(); i++) 
						System.err.print(String.format("%.2f", g.getGraph().wn[i]) + " ");
				} else if (WeightedTimeGraph.class.equals(g.getClass())) {
					for(int i=0 ; i < g.getm(); i++) 
						System.err.print(String.format("%.2f", g.getGraph().we[i]) + " ");
				}
				System.err.print("\n");
				System.err.print(tree.getPartitions() + "\n");
				tree.visualize(); 
			}
			
			// prune
			pruneTree(tree,g);
			
			if (Common.DEBUG) {
				tree.visualize();
			}
			
			System.err.print(t + ",sz="+ tree.numSplitDecendants() + "[" + (System.currentTimeMillis()- start) + "ms]\n");
		} 
		
	}
	
	static SplitTreeNode scluster(IGraph g, DistanceMeasure dm,
			float deltasse) {
		BitSet filt = new BitSet(g.getn());
		filt.set(0, g.getn());
		float avg = avg(g, filt);
		float sse = err(g, filt); 
		SplitTreeNode root = new SplitTreeNode(filt,null, avg, sse);
		split(root, g, dm, deltasse);
		return root;
	}
	
	public static void split(SplitTreeNode c, IGraph g, DistanceMeasure dm, float deltasse) {
		
		int bestvp = -1;
		float bestd = -1;
		float minsse = Float.POSITIVE_INFINITY, tmpsse = -1;
		BitSet fl = null, fr = null;
		// First, find the best possible split
		// TODO(): Very expensive and to be optimized
		for (int vp=0; vp < g.getn(); vp++) {
			for (float d = 0; d < dm.max(); d+=1) {
				fl = filter(c.els,vp,d,true,dm,g);
				fr = filter(c.els,vp,d,false,dm,g);
				if (fl.cardinality()*fr.cardinality() == 0) continue; // no split
				tmpsse = err(g,fl) + err(g,fr);
				if (minsse > tmpsse) {
					minsse = tmpsse;
					bestvp = vp;
					bestd = d;
				}
			}
		}
		
		if (c.sse - minsse > deltasse) {
			c.vantagepoint = bestvp;
			c.dist = bestd;
			
			// Left child (within)
			fl = filter(c.els,bestvp,bestd,true,dm, g);
			c.lc = new SplitTreeNode(fl, c, avg(g,fl), err(g,fl));
			split(c.lc, g, dm, deltasse);
			
			// Right child (out)
			fr = filter(c.els,bestvp,bestd,false,dm, g);
			c.rc = new SplitTreeNode(fr, c, avg(g,fr), err(g,fr));
			split(c.rc, g, dm, deltasse);
		}
	}

	// PRUNE
	public static void pruneTree(SplitTreeNode root, IGraph g) {
		float rsse = root.aggregateSSE(); 
		int rnumn = root.numSplitDecendants();
		int n = root.els.cardinality();
		if (isEdgeW(g)) n /= 2;
		
		float minBIC = 0;
		int index = -1;
		
		ArrayList<SplitTreeNode> snds = new ArrayList<SplitTreeNode>();
		ArrayList<SplitTreeNode> decend = new ArrayList<SplitTreeNode>();
		root.getAllSplitDecendants(snds);
		
		do {
			// number of values
			
			minBIC = getModelComplexity(rsse, rnumn, n);
			index = -1;
			for (int i=0; i < snds.size(); i++) {
				float tmp = getModelComplexity(rsse-snds.get(i).aggregateSSE()+snds.get(i).sse, 
								   rnumn - snds.get(i).numSplitDecendants(), n);
				if (tmp < minBIC) {
					index = i;
					minBIC = tmp;
				}
			}
			if (index >= 0) {
				SplitTreeNode nd = snds.get(index);
				if (Common.DEBUG) {
					System.err.print("Pruning under " + nd.toString() + "\n" + nd.getPartitions() + "\n");
				}
				// update global model complexity and error for the new tree
				rsse = rsse-nd.aggregateSSE()+nd.sse;
				rnumn = rnumn - nd.numSplitDecendants();
				// update list
				decend.clear();
				nd.getAllDescendants(decend);
				snds.removeAll(decend);
				
				// prune tree
				nd.lc.pruneSubtree();
				nd.rc.pruneSubtree();
				nd.lc = null;
				nd.rc = null;
				if (Common.DEBUG) {
					root.visualize();
				}
			}
			
		} while (index >= 0 && snds.size() > 1);
	} 
	
	// UTILS
	// This is an important method. We should apply an index for Radius search here.
	// VP tree on the graph might not be good enough. If the radii are discrete we can 
	private static BitSet filter(BitSet els, int vp, float d, 
			  boolean within, DistanceMeasure dm, 
			  IGraph g) {
		BitSet result = new BitSet(els.size());
		result.or(els);
		if (isNodeW(g)) { // node weighted
			for (int i = els.nextSetBit(0); i >= 0 && i < dm.numelems(); i = els.nextSetBit(i+1)) {
				if (within != (dm.getD(vp, i))<=d) result.set(i, false);
			}
		} else if (isEdgeW(g)) { // edge weighted
			for (int n = 0; n < g.getn(); n++) {
				for (int i = g.ind(n); i < g.ind(n+1); i++) {
					if (!els.get(i)) continue;
					// if none of the adjacent vertices is within d and within=false exclude the edge 
					if (within != (dm.getD(vp, n)<=d || dm.getD(vp, g.endv(i))<=d)) 
						result.set(i, false);
				}
			}
		}
		return result;
	}
	
	private static float avg(IGraph g, BitSet els) {
		float res = 0;
		if (isNodeW(g)) {
			for (int i = els.nextSetBit(0); i >= 0 && i < g.getn(); i = els.nextSetBit(i+1)) {
				res += g.wn(i)/els.cardinality();
			}
		} else if (isEdgeW(g)) {
			for (int n = 0; n < g.getn(); n++) {
				for (int i = g.ind(n); i < g.ind(n+1); i++) {
					if (!els.get(i)) continue;
					res += g.we(i)/els.cardinality();
				}
			}
		}
		return res;
	}
	
	public static float err(IGraph g,BitSet els) {
		if (ERF.equals("SSE")) 
			return sse_err(g, els);
		else if  (ERF.equals("L1")) 
			return l1_err(g, els);
		else {
			System.err.print("Unknown error function\n");
			System.exit(1);
			return -1;
		}
	}

	private static float sse_err(IGraph g, BitSet els) {
		float res = 0;
		float avg = avg(g,els);
		if (isNodeW(g)) {
			for (int i = els.nextSetBit(0); i >= 0 && i < g.getn(); i = els.nextSetBit(i+1)) {
				res += Math.pow(avg - g.wn(i),2);
			}
		} else if (isEdgeW(g)) {
			for (int n = 0; n < g.getn(); n++) {
				for (int i = g.ind(n); i < g.ind(n+1); i++) {
					if (!els.get(i)) continue;
					res += Math.pow(avg - g.we(i), 2)/2.0;
				}
			}
		}
		
		if (els.cardinality() == 1 && res > 0) {
			System.err.print(els.toString() + "\t" + avg + "\t" + res + "\n");
		}
		return res;
	}
	
	private static float l1_err(IGraph g, BitSet els) {
		float res = 0;
		float avg = avg(g,els);
		if (isNodeW(g)) {
			for (int i = els.nextSetBit(0); i >= 0 && i < g.getn(); i = els.nextSetBit(i+1)) {
				res += Math.abs(avg - g.wn(i));
			}
		} else if (isEdgeW(g)) {
			for (int n = 0; n < g.getn(); n++) {
				for (int i = g.ind(n); i < g.ind(n+1); i++) {
					if (!els.get(i)) continue;
					res += Math.abs(avg - g.we(i))/2.0;
				}
			}
		}
		
		if (els.cardinality() == 1 && res > 0) {
			System.err.print(els.toString() + "\t" + avg + "\t" + res + "\n");
		}
		return res;
	}
	
	public static float getModelComplexity(float sse, int params, int n) {
		// Schwarz's Bayesian Criterion (SBC) = ln(MeanSE) + kln(n)/n	
		float alpha = 0.8F;
		return (float) (alpha*Math.log(sse*1.0/n) + (1-alpha)*params*Math.log(n)*1.0/n);
	}
	
	public static boolean isNodeW(IGraph g) {
		return NodeWeightedTimeGraph.class.equals(g.getClass());
	}
	public static boolean isEdgeW(IGraph g) {
		return WeightedTimeGraph.class.equals(g.getClass());
	}
	
	// MAIN AND DATA
	public static NodeWeightedTimeGraph getSmallTestGraph(int ts) {
		Graph gr = GraphGenerator.getTestGraph();
		double[][] cl = {{1.0,1.0,0.0,0.0,0.0,0.0,1.0},
				 {0.0,0.0,1.0,1.0,1.0,1.0,0.0}};

		NodeWeightedTimeGraph ntg = 
			//GraphGenerator.generateRandomNWSoftClusters(ts, gr, cl, 10, 0.9, 10, 0.5 );
			GraphGenerator.generateSineNWClusters(ts, gr, cl, 10, 0.5 );
			
		Arrays.fill(ntg.we,1.0);
		return ntg;
	}
	
	public static NodeWeightedTimeGraph getRandomTestGraph(int ts, int n, int m, int nclusters) {
		Random r = new Random();
		Graph gr = GraphGenerator.generateRandom(n, m);
		Arrays.fill(gr.we,1.0);
		
		double[][] cl = new double[nclusters][n];
		TreeSet<Integer> cluster = new TreeSet<Integer>();
		HashSet<Integer> exclude = new HashSet<Integer>();
		
		// sample some connected clusters
		for (int i = 0; i < nclusters-1; i++) {
			cluster.clear();
			int center = r.nextInt(gr.getn());
			while (exclude.contains(center)) center = r.nextInt(gr.getn()); 
			cluster.addAll(gr.getBFSNodesExcludeNodes(center,
					       (int)(0.9*n/nclusters),
					       exclude));
			exclude.addAll(cluster);
			System.err.print(cluster.toString() + "\n");
			for (Integer nd: cluster) cl[i][nd] = 1.0; 
		}
		
		// put the rest of the nodes in the last cluster (Background)
		// normalize the clusters
		for (int nd = 0; nd < gr.getn(); nd++) {
			double sum = 0;
			for (int i=0; i < nclusters - 1; i++) sum += cl[i][nd];
			if (sum == 0) cl[nclusters-1][nd] = 1.0;
			else {
				for (int i=0; i < nclusters - 1; i++) cl[i][nd] /= sum;
			}
		}

		NodeWeightedTimeGraph ntg = 
			//GraphGenerator.generateRandomNWSoftClusters(ts, gr, cl, 100, 0.9, 100, 0.5 );
			GraphGenerator.generateSineNWClusters(ts, gr, cl, 100, 0.5 );
		Arrays.fill(ntg.we,1.0);
		return ntg;
	}
	
	public static void main(String[] args) throws IOException{
		
		NodeWeightedTimeGraph ntg = null;
		WeightedTimeGraph etg = null;
		
//		ntg = getSmallTestGraph(100);
//		int t = 100, n = 100, m = 600, nclusters = 10;
//		ntg = getRandomTestGraph(100,n,m,nclusters);
		
		// fb
//		etg = IO.readEWGraph("/home/petko/data/seine/01_facebook/preprocessed/graph1000", 
//							 "/home/petko/data/seine/01_facebook/preprocessed/ew1000");
		
		// 100 nodes traffic
//		etg = IO.readEWGraph("/home/petko/data/seine/15_traffic/preprocessed/graphall", 
//		 "/home/petko/data/seine/15_traffic/preprocessed/ewall");
		
		ntg = IO.readNWGraph("/home/petko/data/seine/15_traffic/preprocessed/graphall", 
		 "/home/petko/data/seine/15_traffic/preprocessed/nwall");
		
		ntg = IO.readNWGraph("/home/petko/data/seine/15_traffic/preprocessed/100.speeds.graph", 
		 "/home/petko/data/seine/15_traffic/preprocessed/100.speeds.data");
		for (int i = 0; i < ntg.getm(); i++) ntg.we[i] = 1;
		
		
		Common.DEBUG = true;
		DistanceMeasure dm = new APSP(ntg,true);
		for (int t = 0; t < ntg.gett()-1; t++) {
			ntg.aggregateByTime(t, t);
			SplitTreeNode tree = scluster(ntg, dm, new StopCriteria(5, -1),false);
			tree.visualize("g1");
			ntg.aggregateByTime(t+1, t+1);
			SplitTreeNode treen = scluster(ntg, dm, new StopCriteria(5, -1),false);
			treen.visualize("g2");
			ExtToolsDriver.kill("eog");
		}
	}
}
