package alg;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.BitSet;
import java.util.HashMap;
import java.util.HashSet;
import java.util.TreeMap;
import java.util.TreeSet;

import misc.Common;
import misc.ComparableEdge;

import graph.*;
import generator.GraphGenerator;

public class ST {
	
	// takes positive and negative edges and maximizes 
	// the sum of edge weights
	public static Tree computeMaxST(Graph g) {
		//assert(g.isUndirected());
		
		// order edges
		long start = System.currentTimeMillis();
		TreeMap<Double, ArrayList<String>> edges = new TreeMap<Double, ArrayList<String>>();
		int[] comp = new int[g.getn()];
		for (int i = 0; i < g.getn(); i++) comp[i]=i;
		int[] tree_edges = new int[2*(g.getn()-1)];
		
		// Order edges by increasing cost
		for (int i = 0; i < g.getm(); i++) {
			if (!edges.containsKey(g.we[i])) edges.put(g.we[i], new ArrayList<String>());
			edges.get(g.we[i]).add(g.getOrigin(i)  + Common.SEP + g.endv[i]);
		}
		long ordert = System.currentTimeMillis() - start;
		
		start = System.currentTimeMillis();
		// Get the MST
		int added = 0;
		int v1,v2,c;
		double sum = 0;
		for (double d:edges.descendingKeySet()) {
			for (String s:edges.get(d)) {
				v1 = Integer.parseInt(s.split(Common.SEP)[0]);
				v2 = Integer.parseInt(s.split(Common.SEP)[1]);
				if (comp[v1] != comp[v2]) {
					sum += d;
					// add the edge indices
					tree_edges[added]=g.getEdgeIndex(v1, v2);added++;
					tree_edges[added]=g.getEdgeIndex(v2, v1);added++;
					c=comp[v2];
					for (int i = 0; i < comp.length; i++) {
						if (comp[i] == c) comp[i] = comp[v1];
					}
					if (added >= 2*(g.getn()-1)) break;
				}
			}
			if (added >= 2*(g.getn()-1)) break;
		}
		long kruskalt = System.currentTimeMillis() - start;
		
		start = System.currentTimeMillis();
		Tree t = new Tree(g.subgraph(tree_edges));
		long constt = System.currentTimeMillis() - start;
		System.err.print(sum + "( ord: " + ordert + " + krus: " + kruskalt + " + constr: " + constt + ")" );
		return t;
	}
	
	// takes positive and negative edges and maximizes 
	// the sum of edge weights
	public static Tree computeMaxSTFast(Graph g, TreeSet<ComparableEdge> edges) {
		long start = System.currentTimeMillis();
		TreeSet<ComparableEdge> tmpedges = new TreeSet<ComparableEdge>();
		int[] comp = new int[g.getn()];
		for (int i = 0; i < g.getn(); i++) comp[i]=i;
		
		start = System.currentTimeMillis();
		// Get the MST
		int c2 = -1;
		double sum = 0;
		for (ComparableEdge e: edges.descendingSet()) {
			if (comp[e.v1] != comp[e.v2]) {
				// add the edge indices
				tmpedges.add(e);
				c2 = comp[e.v2];
				for (int i = 0; i < comp.length; i++) {
					if (comp[i] == c2) { comp[i] = comp[e.v1];}
				}
				if (tmpedges.size() >= (g.getn()-1)) break;
			}
		}
		long kruskalt = System.currentTimeMillis() - start;
		
		start = System.currentTimeMillis();
		Tree t = new Tree(g.subgraph(tmpedges));
		long constt = System.currentTimeMillis() - start;
		
		//System.err.print(sum + "( ord: " + ordert + " + krus: " + kruskalt + " + constr: " + constt + ")" );
		return t;
	}
	
	// takes positive edges as costs
	public static Tree computeMST(Graph g){
		assert(g.isUndirected());
		// TODO(): assert that the graph is connected
		TreeMap<Double, ArrayList<String>> edges = new TreeMap<Double, ArrayList<String>>();
		int[] comp = new int[g.getn()];
		for (int i = 0; i < g.getn(); i++) comp[i]=i;
		int[] tree_edges = new int[2*(g.getn()-1)];
		
		
		// Order edges by increasing cost
		for (int i = 0; i < g.getm(); i++) {
			if (!edges.containsKey(g.we[i])) edges.put(g.we[i], new ArrayList<String>());
			edges.get(g.we[i]).add(g.getOrigin(i)  + Common.SEP + g.endv[i]);
		}
		
		// Get the MST
		int added = 0;
		int v1,v2,c;
		for (double d:edges.keySet()) {
			for (String s:edges.get(d)) {
				v1 = Integer.parseInt(s.split(Common.SEP)[0]);
				v2 = Integer.parseInt(s.split(Common.SEP)[1]);
				if (comp[v1] != comp[v2]) {
					// add the edge indices
					tree_edges[added]=g.getEdgeIndex(v1, v2);added++;
					tree_edges[added]=g.getEdgeIndex(v2, v1);added++;
					c=comp[v2];
					for (int i = 0; i < comp.length; i++) {
						if (comp[i] == c) comp[i] = comp[v1];
					}
					if (added >= 2*(g.getn()-1)) break;
				}
			}
			if (added >= 2*(g.getn()-1)) break;
		}
		return new Tree(g.subgraph(tree_edges));
	}
	
	// This is the 2-approximation of the Steiner tree
	public static Tree computeSteiner(Graph g, ArrayList<String> terminals) {
		assert(g.isUndirected()):"Called for directed graph";
		if(terminals.size()==1) 
			return new Tree(new Graph(terminals.get(0),g.wn[g.getNodeIndex(terminals.get(0))]));
		APSP apsp = g.getMetic();
		
		
		TreeMap<Double, ArrayList<String>> edges = new TreeMap<Double, ArrayList<String>>();
		int[] comp = new int[terminals.size()];
		for (int i = 0; i < terminals.size(); i++) comp[i]=i;
		HashSet<String> added_edges = new HashSet<String>();
		
		int v1,v2,c=-1;
		double val;
		// Order edges by increasing cost
		for (int i = 0; i < terminals.size(); i++) {
			for (int j = 0; j < terminals.size(); j++) {
				if (i == j) continue;
				val = apsp.getD(g.getNodeIndex(terminals.get(i)), g.getNodeIndex(terminals.get(j)));
				if (!edges.containsKey(val)) edges.put(val, new ArrayList<String>());
				edges.get(val).add(terminals.get(i)  + Common.SEP + terminals.get(j));
			}
		}
		
		// Get the MST
		int disj_comps = comp.length, ind;
		HashSet<Integer> comp_ind = new HashSet<Integer>();
		String es;
		for (double d:edges.keySet()) {
			for (String s:edges.get(d)) {
				v1 = g.getNodeIndex(s.split(Common.SEP)[0]);
				v2 = g.getNodeIndex(s.split(Common.SEP)[1]);
				if (comp[terminals.indexOf(g.names[v1])] != comp[terminals.indexOf(g.names[v2])]) {
					ArrayList<Integer> path = apsp.getSPath(v1, v2);
					// add the whole path
					for (int i = 0; i < path.size()-1; i++ ) {
						es=((path.get(i)<path.get(i+1))
								?path.get(i) + Common.SEP + path.get(i+1)
							    :path.get(i+1) + Common.SEP + path.get(i)); 
						added_edges.add(es);
					}
					
					// update components
					for (int i = 1; i < path.size(); i++ ) {
						// is path.get(i) a terminal
						ind = terminals.indexOf(g.names[path.get(i)]);
						if (ind > -1) {
							c=comp[ind];
							for (int j = 0; j < comp.length; j++) {
								if (comp[j] == c) comp[j] = comp[terminals.indexOf(g.names[v1])];
							}
						}
					}
					// count disjoint components
					comp_ind.clear();
					for (int i:comp) comp_ind.add(i);
					disj_comps = comp_ind.size();
					if (disj_comps ==1 ) break;
				}
			}
			if (disj_comps ==1 ) break;
		}
		
		int pos = 0,n1,n2;
		int[] tree_edges = new int[2*(added_edges.size())];
		for (String s: added_edges) {
			n1 = Integer.parseInt(s.split(Common.SEP)[0]);
			n2 = Integer.parseInt(s.split(Common.SEP)[1]);
			tree_edges[pos] = g.getEdgeIndex(n1, n2);pos++;
			tree_edges[pos] = g.getEdgeIndex(n2, n1);pos++;
		}
		//System.out.println(g.edgesToString(tree_edges));
		return new Tree(g.subgraph(tree_edges));
	}
	
	public static void main(String[] argv) {
		Graph g = GraphGenerator.getTestGraph();
		System.out.print(g.toString());
		//ST.computeMST();
		ArrayList<String> terminals = new ArrayList<String>();
		terminals.add("6");terminals.add("3"); terminals.add("1");terminals.add("0");
		//terminals.add("6");terminals.add("5"); //terminals.add(2); terminals.add(6);
		System.out.print(ST.computeSteiner(g,terminals).toString());
	} 
	
}
