package ch.ethz.sae;

import java.util.HashMap;
import java.util.Iterator;
import java.util.List;

import apron.Abstract1;
import apron.ApronException;
import apron.Environment;
import apron.Manager;
import apron.MpqScalar;
import apron.Polka;
import apron.Tcons1;
import apron.Texpr1BinNode;
import apron.Texpr1CstNode;
import apron.Texpr1Intern;
import apron.Texpr1Node;
import apron.Texpr1VarNode;
import soot.IntegerType;
import soot.Local;
import soot.SootClass;
import soot.SootField;
import soot.Unit;
import soot.Value;
import soot.jimple.BinopExpr;
import soot.jimple.DefinitionStmt;
import soot.jimple.IntConstant;
import soot.jimple.Stmt;
import soot.jimple.internal.AbstractBinopExpr;
import soot.jimple.internal.JAddExpr;
import soot.jimple.internal.JEqExpr;
import soot.jimple.internal.JGeExpr;
import soot.jimple.internal.JGtExpr;
import soot.jimple.internal.JIfStmt;
import soot.jimple.internal.JLeExpr;
import soot.jimple.internal.JLtExpr;
import soot.jimple.internal.JMulExpr;
import soot.jimple.internal.JNeExpr;
import soot.jimple.internal.JSubExpr;
import soot.jimple.internal.JimpleLocal;
import soot.jimple.toolkits.annotation.logic.Loop;
import soot.toolkits.graph.LoopNestTree;
import soot.toolkits.graph.UnitGraph;
import soot.toolkits.scalar.ForwardBranchedFlowAnalysis;
import soot.util.Chain;

/* numerical analysis */
public class Analysis extends ForwardBranchedFlowAnalysis<AWrapper> {

	/* Apron abstract domain instance */
	public static Manager man;

	/* Apron environment */
	public static Environment env;

	public void run() {
		doAnalysis();
	}

	@Override
	protected void flowThrough(AWrapper inWrapper, Unit op,
			List<AWrapper> fallOutWrappers, List<AWrapper> branchOutWrappers) {

		Stmt stmt = (Stmt) op;
		Abstract1 o_fallout = null, o_branchout = null;

		try {
			o_fallout = new Abstract1(man, inWrapper.get());
			o_branchout = new Abstract1(man, inWrapper.get());

			if (stmt instanceof DefinitionStmt) {
				/* assignment statement */
				o_fallout = new Abstract1(man, inWrapper.get());
				o_branchout = new Abstract1(man, inWrapper.get());

				DefinitionStmt sd = (DefinitionStmt) stmt;
				Value lhs = sd.getLeftOp();
				Value rhs = sd.getRightOp();

				if (lhs instanceof JimpleLocal) {
					/* assignment to a local var */
					String lhs_varName = ((JimpleLocal) lhs).getName();
					System.out.println("debug: lhs = JimpleLocal: lhs_varName: " + lhs_varName);

					if (rhs instanceof IntConstant) {
						/* assign constant to a local var (e.g., n = 10) */
						IntConstant c = ((IntConstant) rhs);
						System.out.println("debug: IntConstant - consant: " + c);
						Texpr1Node rAr = new Texpr1CstNode(new MpqScalar(c.value));
						Texpr1Intern xp = new Texpr1Intern(env, rAr);
						o_fallout.assign(man, lhs_varName, xp, null);
						/* o_branchout does not matter for Definition statements */
					}

					if (rhs instanceof JimpleLocal) {
						/* assign local var to a local var */
						JimpleLocal var = ((JimpleLocal) rhs);
						String rhs_varName = var.getName();
						System.out.println("debug: jimplelocal - rhs_varName: " + rhs_varName);
						if (env.hasVar(rhs_varName)){ //if is a pointer assignment, we must not change o_fallout (i.e, it stays inWrapper.get())
							Texpr1Node rAr = new Texpr1VarNode(rhs_varName);
							Texpr1Intern xp = new Texpr1Intern(env, rAr);
							o_fallout.assign(man, lhs_varName, xp, null);
						}

						/* o_branchout does not matter for Definition statements */
					}

					if (rhs instanceof BinopExpr) {
						/* assign BinOp expression to a local var */
						BinopExpr bexpr = ((BinopExpr) rhs);
						int o = -1;
						if (bexpr instanceof JAddExpr) {
							o = Texpr1BinNode.OP_ADD;
						} else if (bexpr instanceof JSubExpr) {
							o = Texpr1BinNode.OP_SUB;
						} else if (bexpr instanceof JMulExpr) {
							o = Texpr1BinNode.OP_MUL;
						}

						Value op1 = bexpr.getOp1();
						Value op2 = bexpr.getOp2();
						System.out.println("debug: BinopExpr - opq1: " + op1 + ", op2: " + op2);
						Texpr1Node lAr = null;
						Texpr1Node rAr = null;

						if (op1 instanceof JimpleLocal) {
							JimpleLocal var = (JimpleLocal) op1;
							String rhs_varName = var.getName();
							lAr = new Texpr1VarNode(rhs_varName);
						} else if (op1 instanceof IntConstant) {
							IntConstant c = ((IntConstant) op1);
							lAr = new Texpr1CstNode(new MpqScalar(c.value));
						}

						if (op2 instanceof JimpleLocal) {
							JimpleLocal var = (JimpleLocal) op2;
							String rhs_varName = var.getName();
							rAr = new Texpr1VarNode(rhs_varName);
						} else if (op2 instanceof IntConstant) {
							IntConstant c = ((IntConstant) op2);
							rAr = new Texpr1CstNode(new MpqScalar(c.value));
						}
						Texpr1Node t = new Texpr1BinNode(o, lAr, rAr);
						Texpr1Intern xp = new Texpr1Intern(env, t);
						o_fallout.assign(man, lhs_varName, xp, null);
						/* o_branchout does not matter for Definition statements */
					}
				}

			} else if (stmt instanceof JIfStmt) {
				o_fallout = new Abstract1(man, inWrapper.get());
				o_branchout = new Abstract1(man, inWrapper.get());

				JIfStmt ifst = (JIfStmt) stmt;
				Value cond = ifst.getCondition();
				Stmt target = ifst.getTarget();
				System.out.println("debug: lhs = IfStmt: cond = " + cond + ", target = " + target + " ");


				Value op1 = ((AbstractBinopExpr) cond).getOp1();
				Value op2 = ((AbstractBinopExpr) cond).getOp2();
				Texpr1Node op1Node = null;
				Texpr1Node op2Node = null;

				System.out.println("debug: lhs = IfStmt: op1 = " + op1 + ", op2 = " + op2 + " ");

				int op1case = 0;
				int op2case = 0;

				if (op1 instanceof JimpleLocal) {
					JimpleLocal var = (JimpleLocal) op1;
					String rhs_varName = var.getName();
					op1Node = new Texpr1VarNode(rhs_varName);
				} else if (op1 instanceof IntConstant) {
					IntConstant c = ((IntConstant) op1);
					op1Node = new Texpr1CstNode(new MpqScalar(c.value));
					if (c.value == Integer.MAX_VALUE) {
						op1case = 1;
					}
					if (c.value == Integer.MIN_VALUE) {
						op1case = -1;
					}
				}
				if (op2 instanceof JimpleLocal) {
					JimpleLocal var = (JimpleLocal) op2;
					String rhs_varName = var.getName();
					op2Node = new Texpr1VarNode(rhs_varName);
				} else if (op2 instanceof IntConstant) {
					IntConstant c = ((IntConstant) op2);
					op2Node = new Texpr1CstNode(new MpqScalar(c.value));
					if (c.value == Integer.MAX_VALUE) {
						op2case = 1;
					}
					if (c.value == Integer.MIN_VALUE) {
						op2case = -1;
					}
				}

				Texpr1Node binop12 = new Texpr1BinNode(Texpr1BinNode.OP_SUB, op1Node, op2Node);
				Texpr1Node binop21 = new Texpr1BinNode(Texpr1BinNode.OP_SUB, op2Node, op1Node);

				//handling edge cases for != conditions
				int k = Tcons1.DISEQ;
				Texpr1Node node = binop12;
				if (op1case != 0 || op2case != 0) {
					//op1case ==  1: maxint != op2 <=> maxint > op2 <=> binop12 > 0
					//op2case == -1: op1 != minint <=> op1 > minint <=> binop12 > 0
					k = Tcons1.SUP;
				}
				if (op1case == -1 || op2case == 1) {
					//op1case == -1: minint != op2 <=> minint < op2 <=> binop21 > 0
					//op2case ==  1: op1 != maxint <=> op1 < maxint <=> binop21 > 0
					node = binop21;
				}

				Tcons1 tconsBR = null;
				Tcons1 tconsFT = null;

				if (cond instanceof JEqExpr) {
					tconsBR = new Tcons1(env,Tcons1.EQ,binop12);
					tconsFT = new Tcons1(env,k,node);
				} else if (cond instanceof JNeExpr) {
					tconsBR = new Tcons1(env,k,node);
					tconsFT = new Tcons1(env,Tcons1.EQ,binop12);
				} else if (cond instanceof JGtExpr) {
					tconsBR = new Tcons1(env,Tcons1.SUP,binop12);
					tconsFT = new Tcons1(env,Tcons1.SUPEQ,binop21);
				} else if (cond instanceof JGeExpr) {
					tconsBR = new Tcons1(env,Tcons1.SUPEQ,binop12);
					tconsFT = new Tcons1(env,Tcons1.SUP,binop21);
				} else if (cond instanceof JLtExpr) {
					tconsBR = new Tcons1(env,Tcons1.SUP,binop21);
					tconsFT = new Tcons1(env,Tcons1.SUPEQ,binop12);
				} else if (cond instanceof JLeExpr) {
					tconsBR = new Tcons1(env,Tcons1.SUPEQ,binop21);
					tconsFT = new Tcons1(env,Tcons1.SUP,binop12);
				}

				o_fallout.meet(man, tconsFT);
				o_branchout.meet(man, tconsBR);

			}


			/* - code below propagates the o_fallout and o_branchout computed before (as side effect);
			 * - no need to change anything below;
			 * - just make sure that o_fallout and o_branchout vars are computed before;
			*/
			for (Iterator<AWrapper> it = fallOutWrappers.iterator(); it.hasNext();) {
				AWrapper op1 = it.next();
				if (o_fallout != null) {
					op1.set(o_fallout);
				}
			}

			for (Iterator<AWrapper> it = branchOutWrappers.iterator(); it.hasNext();) {
				AWrapper op1 = it.next();
				if (o_branchout != null) {
					op1.set(o_branchout);
				}
			}

		} catch (ApronException e) {
			e.printStackTrace();
			System.exit(1);
		}
	}

	/* ======================================================== */

	/* no need to use or change the variables and methods below */

	/* collect all the local int variables in the class */
	private void recordIntLocalVars() {

		Chain<Local> locals = g.getBody().getLocals();

		int count = 0;
		Iterator<Local> it = locals.iterator();
		while (it.hasNext()) {
			JimpleLocal next = (JimpleLocal) it.next();
			if (next.getType() instanceof IntegerType)
				count += 1;
		}

		local_ints = new String[count];

		int i = 0;
		it = locals.iterator();
		while (it.hasNext()) {
			JimpleLocal next = (JimpleLocal) it.next();
			String name = next.getName();
			if (next.getType() instanceof IntegerType)
				local_ints[i++] = name;
		}
	}

	/* collect all the int fields of the class */
	private void recordIntClassVars() {

		Chain<SootField> ifields = jclass.getFields();

		int count = 0;
		Iterator<SootField> it = ifields.iterator();
		while (it.hasNext()) {
			SootField next = it.next();
			if (next.getType() instanceof IntegerType)
				count += 1;
		}

		class_ints = new String[count];

		int i = 0;
		it = ifields.iterator();
		while (it.hasNext()) {
			SootField next = it.next();
			String name = next.getName();
			if (next.getType() instanceof IntegerType)
				class_ints[i++] = name;
		}
	}

	/* build an environment with integer variables */
	public void buildEnvironment() {

		recordIntLocalVars();
		recordIntClassVars();

		String ints[] = new String[local_ints.length + class_ints.length];

		/* add local ints */
		for (int i = 0; i < local_ints.length; i++) {
			ints[i] = local_ints[i];
		}

		/* add class ints */
		for (int i = 0; i < class_ints.length; i++) {
			ints[local_ints.length + i] = class_ints[i];
		}

		env = new Environment(ints, reals);
	}

	/* instantiate the polyhedra domain */
	private void instantiateDomain() {
		man = new Polka(true);
	}

	/* constructor */
	public Analysis(UnitGraph g, SootClass jc) {
		super(g);

		this.g = g;
		this.jclass = jc;

		buildEnvironment();
		instantiateDomain();

		loopHeads = new HashMap<Unit, Counter>();
		backJumps = new HashMap<Unit, Counter>();
		for (Loop l : new LoopNestTree(g.getBody())) {
			loopHeads.put(l.getHead(), new Counter(0));
			backJumps.put(l.getBackJumpStmt(), new Counter(0));
		}
	}

	@Override
	protected AWrapper entryInitialFlow() {
		Abstract1 top = null;
		try {
			top = new Abstract1(man, env);
		} catch (ApronException e) {
			e.printStackTrace();
			System.exit(1);
		}
		return new AWrapper(top);
	}

	private static class Counter {
		int value;

		Counter(int v) {
			value = v;
		}
	}

	@Override
	protected void merge(Unit succNode, AWrapper in_w1, AWrapper in_w2, AWrapper out_w) {
		Counter count = loopHeads.get(succNode);

		Abstract1 in1 = in_w1.get();
		Abstract1 in2 = in_w2.get();
		Abstract1 out = null;

		try {
			if (count != null) {
				++count.value;
				if (count.value < WIDENING_THRESHOLD) {
					out = in1.joinCopy(man, in2);
				} else {
					out = in1.widening(man, in1.joinCopy(man, in2));
				}
			} else {
				out = in1.joinCopy(man, in2);
			}
			out_w.set(out);
		} catch (Exception e) {
			e.printStackTrace();
			System.exit(1);
		}
	}

	@Override
	protected void merge(AWrapper in_w1, AWrapper in_w2, AWrapper out_w) {

		Abstract1 in1 = in_w1.get();
		Abstract1 in2 = in_w2.get();
		Abstract1 out = null;

		try {
			out = in1.joinCopy(man, in2);
		} catch (ApronException e) {
			e.printStackTrace();
		}

		out_w.set(out);
	}

	@Override
	protected AWrapper newInitialFlow() {
		Abstract1 bottom = null;

		try {
			bottom = new Abstract1(man, env, true);
		} catch (ApronException e) {
		}
		AWrapper a = new AWrapper(bottom);
		a.man = man;
		return a;

	}

	@Override
	protected void copy(AWrapper source, AWrapper dest) {
		try {
			dest.set(new Abstract1(man, source.get()));
		} catch (ApronException e) {
			e.printStackTrace();
		}
	}

	/* widening threshold and widening points */
	private static final int WIDENING_THRESHOLD = 6;
	private HashMap<Unit, Counter> loopHeads, backJumps;

	/* Soot unit graph */
	public UnitGraph g;
	public SootClass jclass;

	/* integer local variables of the method; */
	public String local_ints[];
	/* integer class variables where the method is defined */
	private String class_ints[];

	/* real variables */
	public static String reals[] = { "x" };
}