//==============================================================================
//	
//	Copyright (c) 2013-
//	Authors:
//	* Dave Parker <david.parker@comlab.ox.ac.uk> (University of Oxford)
//	* Mateusz Ujma <mateusz.ujma@cs.ox.ac.uk> (University of Oxford)
//	
//------------------------------------------------------------------------------
//	
//	This file is part of PRISM.
//	
//	PRISM is free software; you can redistribute it and/or modify
//	it under the terms of the GNU General Public License as published by
//	the Free Software Foundation; either version 2 of the License, or
//	(at your option) any later version.
//	
//	PRISM is distributed in the hope that it will be useful,
//	but WITHOUT ANY WARRANTY; without even the implied warranty of
//	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//	GNU General Public License for more details.
//	
//	You should have received a copy of the GNU General Public License
//	along with PRISM; if not, write to the Free Software Foundation,
//	Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
//	
//==============================================================================

package heuristics.update;

import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Set;

import jdd.JDD;
import explicit.CachedModelExplorer;
import explicit.ModelExplorer;
import explicit.STPGExplicit;
import parser.State;
import parser.Values;
import parser.ast.Expression;
import parser.ast.ExpressionIdent;
import parser.ast.LabelList;
import prism.PrismException;
import prism.PrismUtils;
import simulator.PrismModelExplorer;
import strat.MDStrategy;
import heuristics.CollapseModelExplorer;
import heuristics.init.ValueInit;
import heuristics.search.StateValue;

/*
 * Class StateUpdate updates the value in the state based on the optimal action at given point and values of its successors.
 * The subclasses (primarily StateUpdateSMG) will implement model specific update procedures.
 * 
 */

public abstract class StateUpdate {
	
	public static int NO_BOUND = -1;
	
	protected ModelExplorer pme;
	private Values constantValues;
	private Expression target;
	private State initialState;
	private StateValueContainer stateValueContainer;
	private ValueInit valueInit = new ValueInit(this);
	protected int bound = NO_BOUND;
	private LabelList specialLabels;
	private Set<Integer> coalition;
	protected boolean min;
	protected double epsilon;
	protected boolean absolute;
	
	public StateUpdate(ModelExplorer me, StateValueContainer container, int bound, boolean min, double epsilon, boolean absolute) throws PrismException{
		this.pme = me;
		this.bound = bound;
		this.specialLabels = new LabelList();
		this.specialLabels.addLabel(new ExpressionIdent("deadlock"), new ExpressionIdent("deadlock"));
		this.specialLabels.addLabel(new ExpressionIdent("init"), new ExpressionIdent("init"));
		this.stateValueContainer = container;
		if(me != null) {
			this.initialState = pme.getDefaultInitialState();
		}
		this.min = min;
		this.epsilon = epsilon;
//		this.epsilon = 1e-1;
		this.absolute = absolute;
	}
	
	public abstract int update(State s, int depth) throws PrismException;
	
	public abstract int update(State s) throws PrismException;
	
	public abstract int getBestAction(State s, int depth) throws PrismException;
	
	public abstract int getBestAction(State s) throws PrismException;
	
	/*
	 * We cache data that is not going to change during the exploration process.
	*/
	
	private Map<State, Boolean> cacheState2Target = new HashMap<State,Boolean>();
	private Map<State, Boolean> cacheState2SelfLoop = new HashMap<State,Boolean>();
	private Map<State, Boolean> cacheState2Zero = new HashMap<State,Boolean>();
	
	public MDStrategy strat = null;
	public explicit.Model model = null;
	
	public void mimicMDStrat(MDStrategy strat, explicit.Model model) {
		List<State> states = model.getStatesList();
		this.strat = strat;
		this.model = model;
		for (int i = 0; i < states.size(); i++) {
			State s = states.get(i);
			int act = strat.getChoiceIndex(i);
			setQValue(s,new StateValue(1, 1, act));
		}
	}
	
	public boolean isTarget(State s) throws PrismException{
		Boolean cachedValue = cacheState2Target.get(s) ;
		if(cachedValue == null) {
			pme.queryState(s);
			Expression evTarget = target.deepCopy();
			specialLabels.setLabel(0, pme.getNumTransitions() == 0 ? Expression.True() : Expression.False());
			specialLabels.setLabel(1, initialState.equals(s) ? Expression.True() : Expression.False());
			cachedValue = evTarget.evaluateBoolean(constantValues, s);
			cacheState2Target.put(s, cachedValue);
		} 
		return cachedValue;
	}
	
	public boolean isOnlySelfLoop(State s) throws PrismException{
		Boolean cachedValue = cacheState2SelfLoop.get(s);
		if(cachedValue == null) {
			pme.queryState(s);
			int numChoices = pme.getNumChoices();
			for(int i=0;i<numChoices;i++) {
				int trans = pme.getNumTransitions(i);
				for(int j=0;j<trans;j++) {
					State succ = pme.computeTransitionTarget(i,j);
					if(!succ.equals(s)) {
						cacheState2SelfLoop.put(s,false);
						return false;
					}
				}
			}
			if (!isTarget(s)) setZero(s, true);
			cachedValue = true;
			cacheState2SelfLoop.put(s, cachedValue);
		}
		return cachedValue;
	}
	
	public boolean isZero(State s) throws PrismException{
		Boolean cachedValue = cacheState2Zero.get(s);
		if(cachedValue == null) {
			return false;
		} 
		return cachedValue;
	}
	
	public void setTarget(State s, boolean t) {
		cacheState2Target.put(s, t);
		if(t) {
			setQValue(s, new StateValue(1, 1, StateValue.ACTION_TERMINAL));
		}
	}
	
	public void setZero(State s, boolean t) {
		cacheState2Zero.put(s, t);
//		System.out.println("Setting " + s + " zero in setZero");
		if(t) {
			setQValue(s, new StateValue(0, 0, StateValue.ACTION_TERMINAL));
		}
	}
	
	public int getBound() {
		return bound;
	}
	
	public void setQValue(State s, StateValue val, int depth) {
		stateValueContainer.setQValue(s, val, depth);
	}
	
	public StateValue getQValue(State s, int depth) {
		return stateValueContainer.getQValue(s, depth);
	}
	
	/*
	 * This method sets the current value in the given state
	 * 
	 */
	
	public void setQValue(State s, StateValue val) {
		stateValueContainer.setQValue(s, val, StateValueContainer.UNBOUNDED);
	}
	
	/*
	 * This method returns the current value in the given state
	 * 
	 */
	
	public StateValue getQValue(State s) {
		return stateValueContainer.getQValue(s, StateValueContainer.UNBOUNDED);
	}
	
	public void setTargetOrBoundValue(State s, int depth) throws PrismException{
		if(isTarget(s)) {
			setQValue(s, new StateValue(1.0, 1.0, StateValue.ACTION_TERMINAL), depth);
			return;
		}
		if(depth >= bound) {
			setQValue(s, new StateValue(0, 0, StateValue.ACTION_TERMINAL), depth);
			return;
		}
	}
	
	public void setTargetOrZeroValue(State s) throws PrismException{
		if(isTarget(s)) {
			setQValue(s, new StateValue(1.0, 1.0, StateValue.ACTION_TERMINAL));
			return;
		}
		if(isZero(s)) {
//			System.out.println("setting " + s + " zero");
			setQValue(s, new StateValue(0, 0, StateValue.ACTION_TERMINAL));
			return;
		}
	}
	
	public int getStateValueNumber() {
		int stateValuesNumber = 0;
		/*Iterator<State> it = qValues.keySet().iterator();
		while(it.hasNext()) {
			State s = it.next();
			stateValuesNumber += qValues.get(s).size();
		}*/
		return stateValuesNumber;
	}
	
	public int getNumExplored() {
		return stateValueContainer.getSize();
	}
	
	public int getNumReachable() throws PrismException{
		LinkedList<StateAndDepth> explore = new LinkedList<StateAndDepth>();
		Set<StateAndDepth> visited = new HashSet<StateAndDepth>();
		int depth = bound == NO_BOUND ? StateValueContainer.UNBOUNDED : 0;
		State s = initialState;
		explore.add(new StateAndDepth(s, depth));
		while(!explore.isEmpty()) {
			StateAndDepth sD = explore.remove(0);
			s = sD.state;
			depth = sD.depth;
			StateValue sV = stateValueContainer.getQValue(s, depth);
//			System.out.println("state: " + s + " depth: " + depth + " action: " + sV.getAction());
			if(sV != null) {
				int action = sV.getAction();
				if(action != StateValue.ACTION_TERMINAL && action != StateValue.ACTION_UNDEF) {
					pme.queryState(sD.state);
//					CollapseModelExplorer mex = (CollapseModelExplorer)pme;
//					System.out.println("Current state " + sD.state + " in EC?: " + mex.isInEC());
					int trans = 0;
//					try {
						trans = pme.getNumTransitions(action);
//					}
//					catch (Exception e) {
////					finally {
////						System.out.println(JDD.ONE);
////						System.out.println(JDD.ZERO);
//						System.out.println("Error when acessing action " + action + " | " + stateValueContainer.getQValue(s, depth).getAction() + " , depth: " + depth + " of " + sD.state +"\n#choices: " + pme.getNumChoices());
//						ModelExplorer testME = mex.modelExplorer;
//						testME.queryState(s);
//						System.out.println("choices according to cachedME: " + testME.getNumChoices());
//						e.printStackTrace();
//					}
					for(int i=0;i<trans;i++) {
						s = pme.computeTransitionTarget(action, i);
						depth = bound == NO_BOUND ? StateValueContainer.UNBOUNDED : depth++;
						StateAndDepth newSD = new StateAndDepth(s, depth);
						if(!visited.contains(newSD) && !explore.contains(newSD)) {
							explore.add(newSD);
						}
					}
				} 
			} 
			visited.add(sD);
		}
		
		return visited.size();
	}
	
	private class StateAndDepth {
		protected State state;
		protected int depth;
		
		public StateAndDepth(State s, int d) {
			state = s;
			depth = d;
		}
		
		public boolean equals(Object o) {
			if(o == null) throw new NullPointerException();
			StateAndDepth sD = (StateAndDepth)o;
			return state.equals(sD.state) && depth == sD.depth;
		}
		
		public int hashCode() {
			return state.hashCode();
		}
		
		public String toString() {
			return "State " + state + " depth " + depth;
		}
	}
	
	public int getNumAllValues() {
		return stateValueContainer.getNumAllValues();
	}
	
	/*
	 * Method same() is used to compare state values during the exploration process. 
	 * Please note that the values of epsilon can have a critical influence on the speed of convergence.
	 * 
	 */
	
	public boolean same(double d1, double d2) {
		return PrismUtils.doublesAreClose(d1, d2, epsilon, absolute);
	}
	
	public boolean same(double d1, double d2,double inpEps) {
		return PrismUtils.doublesAreClose(d1, d2, inpEps, absolute);
	}
	
	public void setConstantValues(Values cV) {
		constantValues = cV;
	}
	
	public void setTarget(Expression e) {
		this.target = e;
	}
	
	public Expression getTarget() {
		return this.target;
	}
	
	public void setCoalition(Set<Integer> coalition) {
		this.coalition = coalition;
	}
	
	public Set<Integer> getCoalition() {
		return this.coalition;
	}
	
	public int getPlayer(State s) throws PrismException{
		pme.queryState(s);
		return coalition.contains(pme.getPlayerForState()) ? STPGExplicit.PLAYER_1 : STPGExplicit.PLAYER_2;
	}
	
	public void setModelExplorer(ModelExplorer me) throws PrismException{
		this.pme = me;
		this.initialState = pme.getDefaultInitialState();
	}
	
	public void setStateValueContainer(StateValueContainer container) {
		this.stateValueContainer = container;
	}
	
	public ModelExplorer getModelExplorer() {
		return this.pme;
	}
	
	public StateValueContainer getStateValueContainer() {
		return this.stateValueContainer;
	}
	
	/*
	 * Computes the lower bound of a given action (at given depth)
	 * 
	 */
	
	public StateValue getLowerBoundActionValue(State s, int a, int depth) throws PrismException {
		pme.queryState(s);
		int trans = pme.getNumTransitions(a);
		double bestVal = 0;
		for(int i=0;i<trans;i++) {
			pme.queryState(s);
			double prob = pme.getTransitionProbability(a, i);
			State succ = pme.computeTransitionTarget(a, i);
			StateValue tVal = null;
			if(depth == -1) {
				tVal = getQValue(succ);
			} else {
				tVal = getQValue(succ, depth);
			}
			
			if(tVal != null) {
				bestVal += prob * tVal.getLowerBound();
			} else {
				double succBestValue = valueInit.getLowerBoundValue(pme, succ, depth);
				bestVal += prob * succBestValue;
			}
		}
		return new StateValue(bestVal, -1, a);
	}
	
	/*
	 * Computes the upper bound of a given action (at given depth)
	 * 
	 */
	
	public StateValue getUpperBoundActionValue(State s, int a, int depth) throws PrismException {
		pme.queryState(s);
		int trans = pme.getNumTransitions(a);
		double worseVal = 0;
		for(int i=0;i<trans;i++) {
			pme.queryState(s);
			double prob = pme.getTransitionProbability(a, i);
			State succ = pme.computeTransitionTarget(a, i);
			StateValue tVal = null;
			if(depth == -1) {
				tVal = getQValue(succ);
			} else {
				tVal = getQValue(succ, depth);
			}
			if(tVal != null) {
				worseVal += prob * tVal.getUpperBound();
			} else {
				double succWorseValue = valueInit.getUpperBoundValue(pme, succ, depth);
				worseVal += prob * succWorseValue;
			}
		}
		return new StateValue(-1, worseVal, a);
	}
	
	public StateValue getLowerBoundActionValue(State s, int a) throws PrismException {
		return getLowerBoundActionValue(s, a, NO_BOUND);
	}
	
	public StateValue getUpperBoundActionValue(State s, int a) throws PrismException {
		return getUpperBoundActionValue(s, a, NO_BOUND);
	}
}
