package util;

import java.io.FileNotFoundException;
import java.io.PrintStream;
import java.util.Arrays;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.stream.Collectors;

import driver.Main;
import pag.node.GNode;
import pag.node.alloc.Alloc_Node;
import pta.ContextSensPTA;
import pta.PTA;
import pta.context.ContextElements;
import pta.pts.PTSetInternal;
import pta.pts.PTSetVisitor;
import soot.Kind;
import soot.MethodOrMethodContext;
import soot.SootMethod;
import soot.Type;
import soot.Unit;
import soot.jimple.toolkits.callgraph.Edge;

public class PTAComparator {
	static final boolean verbose=true;
	public static int count,count2;
	public static void main(String[] args) {
		Comparator comparator = new PTSComparator();
		Main.main(args);
		comparator.record(Main.pta);
		
		String[] newArgs=new String[args.length+1];
		System.arraycopy(args, 0, newArgs, 0, args.length);
		newArgs[args.length]="-pta=rc";
		
		Main.main(newArgs);
		comparator.compare(Main.pta);
//		comparator.out.close();
		count = comparator.count;
		count2 = comparator.count2;
	}
	abstract static class Comparator{
		protected int count,count2;
		PrintStream out;
		public Comparator() {
			out = System.out;
		}
		public Comparator(String filename) {
			try {
				out = new PrintStream(filename);
			} catch (FileNotFoundException e) {
				e.printStackTrace();
			}
		}
		abstract void record(PTA pta);
		abstract void compare(PTA pta);
	}
	static class CGComparator extends Comparator{
		private static Set<H> myCG;
		@Override
		void record(PTA pta) {
			myCG = new HashSet<>();
			pta.getCallGraph().forEach(e->{
				if(e.toString().contains("FakeMain"))return;
				H h = new H(e);
				if(!myCG.add(h)){
					System.out.println(e);
					myCG.forEach(oh->{
						if(oh.equals(h)){
							out.println(oh.e.getSrc()==e.getSrc());
							out.println(oh.e.srcUnit().equals(e.srcUnit()));
							out.println(oh.e.getTgt()==e.getTgt());
							out.println(oh.e.kind()==e.kind());
						}
							
					});
					out.println();
					throw new RuntimeException();
				}
			});
		}

		@Override
		void compare(PTA pta) {
			count=0;
			pta.getCallGraph().forEach(e->{
				if(e.toString().contains("FakeMain"))return;
				if(!myCG.remove(new H(e))){
					out.println(e);
//					PTAUtils.dumpSlicedCallGraph(pta.getCallGraph(), e.getSrc());
//					System.exit(0);
					count++;
				}
			});
			out.println("\nTotal cg over:" + count);
			count2=0;
			out.println("\n old set remained:::::");
			myCG.forEach(h->{
				out.println(h);
				count2++;
			});
			out.println("\nTotal cg lose:" + count2);
		}
		static class H{
			private Edge e;
			protected String srcMtd;
		    protected Unit srcUnit;
		    protected String tgtMtd;
		    protected Kind kind;
			H(Edge e){
				this.e=e;
				MethodOrMethodContext src = e.getSrc();
				srcMtd = src.method().toString();
				srcUnit = e.srcUnit();
				MethodOrMethodContext tgt=e.getTgt();
				tgtMtd = tgt.method().toString();
				kind = e.kind();
			}
			@Override
			public int hashCode() {
				return 
						srcMtd.hashCode()
						+srcUnit.hashCode()
						+tgtMtd.hashCode()
						+kind.hashCode();
			}
			@Override
			public boolean equals(Object other) {
				if (other == null)
					return false;
				H o = (H) other;
				return 
					srcMtd.equals(o.srcMtd)
					&&srcUnit.equals(o.srcUnit)
					&&tgtMtd.equals(o.tgtMtd)
					&&kind.equals(o.kind);
			}
			@Override
			public String toString() {
				return kind.toString()+" edge: "+srcUnit+" in "+"Method " + srcMtd +" ==> "+"Method " + tgtMtd ;
			}
		}
		static class SH extends H{
			private List<Object> srcContext;
		    private List<Object> tgtContext;
			SH(Edge e){
				super(e);
				MethodOrMethodContext src = e.getSrc();
				srcContext = Arrays.asList(
						((ContextElements) src.context())
						.getElements()).stream()
						.map(a->a==null?null:((Alloc_Node)a).getNewExpr())
						.collect(Collectors.toList());
				MethodOrMethodContext tgt=e.getTgt();
				tgtContext =  Arrays.asList(((ContextElements) tgt.context()).getElements()).stream().map(a->a==null?null:((Alloc_Node)a).getNewExpr()).collect(Collectors.toList());
			}
			@Override
			public int hashCode() {
				return super.hashCode()+srcContext.hashCode()
						+tgtContext.hashCode();
			}
			@Override
			public boolean equals(Object other) {
				if (other == null)
					return false;
				SH o = (SH) other;
				return super.equals(other)&&srcContext.equals(o.srcContext)
					&&tgtContext.equals(o.tgtContext);
			}
			@Override
			public String toString() {
				return kind.toString()+" edge: "+srcUnit+" in "+"Method " + srcMtd + " in context " + srcContext+" ==> "+"Method " + tgtMtd + " in context " + tgtContext;
			}
		}
		
	}
	static class PTSComparator extends Comparator{
		private Map<Object, Set<Object>> valToOldAllocs= new HashMap<>();
		private Map<Object, Set<Object>> valToNewAllocs= new HashMap<>();
		private Set<Object> getAllocs(Map<Object, Set<Object>> map, Object val){
			Set<Object> ret = map.get(val);
			if(ret==null)
				map.put(val, ret=new HashSet<>());
			return ret;
		}
		private Set<Object> getOldAllocs(Object val){
			return getAllocs(valToOldAllocs, val);
		}
		private Set<Object> getNewAllocs(Object val){
			return getAllocs(valToNewAllocs, val);
		}
		@Override
		void record(PTA pta) {
			pta.getPag().getLocalPointers().forEach(l->{
				final Set<Object> allocSites = getOldAllocs(l);
				((PTSetInternal) pta.reachingObjects(l)).forall(new PTSetVisitor() {
					@Override
					public void visit(GNode n) {
						allocSites.add(((Alloc_Node) n).getNewExpr());
					}
				});
			});
		}

		@Override
		void compare(PTA pta) {
			count=0;
			pta.getPag().getLocalPointers().forEach(l->{
				final Set<Object> allocSites = getNewAllocs(l);
				((PTSetInternal) pta.reachingObjects(l)).forall(new PTSetVisitor() {
					@Override
					public void visit(GNode n) {
						allocSites.add(((Alloc_Node) n).getNewExpr());
					}
				});
			});
			
			Set<Object> locals = valToOldAllocs.keySet().stream().collect(Collectors.toSet());
			locals.addAll(valToNewAllocs.keySet());
			
			locals.forEach(val->{
				Set<Object> oldAllocSites = getOldAllocs(val);
				Set<Object> allocSites = getNewAllocs(val);
				if(!oldAllocSites.equals(allocSites)){
					if(verbose){
						out.println("\nPoint-to sets not equal!");
						out.println(val);
						out.println("Old Points-to set:"+oldAllocSites.size());
//						oldAllocSites.forEach(a->System.out.println(pta.getPag().valToAllocNode.get(a)));
						out.println("New Points-to set:"+allocSites.size());
//						allocSites.forEach(a->System.out.println(pta.getPag().valToAllocNode.get(a)));
					}
					Set<Object> oldAllocSitesCopy = oldAllocSites.stream().collect(Collectors.toSet());
					oldAllocSites.removeAll(allocSites);
					allocSites.removeAll(oldAllocSitesCopy);
					count+=allocSites.size();
					count2+=oldAllocSites.size();
					out.println("difference:"+allocSites.size()+"/"+oldAllocSites.size());
					if(verbose){
						out.println("over:");
						allocSites.forEach(a->out.println(pta.getPag().valToAllocNode.get(a)));
						out.println("loss:");
						oldAllocSites.forEach(a->out.println(pta.getPag().valToAllocNode.get(a)));
					}
				}
			});
			out.println("\nTotal pts over:" + count);
			out.println("\nTotal pts lose:" + count2);
		}
		
	}
	static class CSAComparator extends Comparator{
		private Map<Type, Integer> typeToSize= new HashMap<>();
		private Map<SootMethod, Integer> methodToSize= new HashMap<>();
		private Integer getOldSize(Type type){
			Integer ret = typeToSize.get(type);
			if(ret==null)return 0;
			return ret;
		}
		private Integer getOldSize(SootMethod method){
			Integer ret = methodToSize.get(method);
			if(ret==null)return 0;
			return ret;
		}
		
		public CSAComparator() {
			super();
		}
		
		public CSAComparator(String string) {
			super(string);
		}

		@Override
		void record(PTA pta) {
			((ContextSensPTA)pta).getContextAllocNodeMap().forEach((a,set)->{
//				Type type = a.getType();
//				typeToSize.put(type, getOldSize(type)+set.size());
				SootMethod method = a.getMethod();
				methodToSize.put(method, getOldSize(method)+set.size());
			});
		}

		@Override
		void compare(PTA pta) {
			((ContextSensPTA)pta).getContextAllocNodeMap().forEach((a,set)->{
//				Type type = a.getType();
//				typeToSize.put(type, getOldSize(type)-set.size());
				SootMethod method = a.getMethod();
				methodToSize.put(method, getOldSize(method)-set.size());
			});
			java.util.Comparator<Entry<?, Integer>> paramComparator = new java.util.Comparator<Entry<?, Integer>>() {
				@Override
				public int compare(Entry<?, Integer> o1, Entry<?, Integer> o2) {
					return o2.getValue()-o1.getValue();
				}
			};
//			typeToSize.entrySet().stream().sorted(paramComparator).forEach(e->{
//				out.println(e.getKey()+" : "+e.getValue());
//			});
			methodToSize.entrySet().stream().sorted(paramComparator).forEach(e->{
				out.println(e.getKey()+" : "+e.getValue());
			});
		}
		
	}
}
