package heuristics.update;

import heuristics.VarInfo;
import heuristics.search.StateValue;

import java.util.BitSet;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import parser.VarList;
import parser.ast.ModulesFile;
import prism.PrismException;
import jdd.*;
import parser.State;

public class StateValueContainerMTBDD implements StateValueContainer
{
	private Map<Integer,JDDNode> low;
	private Map<Integer,JDDNode> high;
	private Map<Integer,JDDNode> action;
	private VarInfo varInfo;
	private Map<State,Integer> inserted;
	private JDDNode actiontest;
	
	public StateValueContainerMTBDD(VarInfo varInfo) throws PrismException {
		this.varInfo = varInfo;
		low =  new HashMap<Integer,JDDNode>();
		high = new HashMap<Integer,JDDNode>();
		action = new HashMap<Integer,JDDNode>();
		inserted = new HashMap<State,Integer>();
		actiontest = JDD.Constant(StateValue.ACTION_UNDEF);
	}
	
	@Override
	public int getSize() { //TODO: return the number of explored states
		return -1;
	}
	
	@Override
	public void setQValue(State s, StateValue val, int depth) {
		JDDNode dd = varInfo.state2JDD(s);
		JDDNode lowDepth = low.get(depth);
		JDDNode highDepth = high.get(depth);
		JDDNode actionDepth = action.get(depth);
		double l = val.getLowerBound();
		double h = val.getUpperBound();
		JDDNode newLow;
		JDDNode newHigh;
		JDDNode newAction;
		int a = val.getAction();
		if (lowDepth == null) {
			JDD.Ref(dd);
			newLow = JDD.ITE(dd, JDD.Constant(l), JDD.ZERO);
		}
		else {
			JDD.Ref(dd);
			newLow = JDD.ITE(dd, JDD.Constant(l), lowDepth);
		}
		if (highDepth == null) {
			JDD.Ref(dd);
			newHigh = JDD.ITE(dd, JDD.Constant(h), JDD.ONE);
		}
		else {
			JDD.Ref(dd);
			newHigh = JDD.ITE(dd, JDD.Constant(h), highDepth);
		}
		if (actionDepth == null) {
			newAction = JDD.ITE(dd, JDD.Constant(a), JDD.Constant(StateValue.ACTION_UNDEF));
		}
		else {
			newAction = JDD.ITE(dd, JDD.Constant(a), actionDepth);
		}
		low.put(depth,newLow);
		high.put(depth,newHigh);
		action.put(depth,newAction);
	}
	
	@Override
	public StateValue getQValue(State s, int depth)
	{
		JDDNode dd = varInfo.state2JDD(s);
		JDDNode lowDepth = low.get(depth);
		JDDNode highDepth = high.get(depth);
		JDDNode actionDepth = action.get(depth);
		double l = 0.0;
		double h = 1.0;
		int a = StateValue.ACTION_UNDEF;
		if (lowDepth == null) {
			return null;
		}
		else {
			JDD.Ref(lowDepth);
			JDD.Ref(dd);
			JDDNode restrictLow = JDD.Restrict(lowDepth, dd);
			l = restrictLow.getValue();
		}
		if (highDepth == null) {
			return null;
		}
		else {
			JDD.Ref(highDepth);
			JDD.Ref(dd);
			JDDNode restrictHigh = JDD.Restrict(highDepth, dd);
			h = restrictHigh.getValue();
		}
		if (actionDepth == null) {
			return null;
		}
		else {
			JDD.Ref(actionDepth);
			JDDNode restrictAction = JDD.Restrict(actionDepth, dd);
			a = ((Double)restrictAction.getValue()).intValue();
		}
		return new StateValue(l,h,a);
	}

	@Override
	public int getNumAllValues()
	{
		// TODO Auto-generated method stub
		return 0;
	}

	@Override
	public void derefAll()
	{
		for (JDDNode node : low.values()) {
//			System.out.println("derefing low: " + node);
			JDD.Deref(node);
		}
		for (JDDNode node : high.values()) {
//			System.out.println("derefing high: " + node);
			JDD.Deref(node);
		}
		for (JDDNode node : action.values()) {
//			System.out.println("derefing action: " + node);
			JDD.Deref(node);
		}
	}
}
