//==============================================================================
//	
//	Copyright (c) 2015-
//	Author:
//  * Andreas Fellner <andreas.fellner@ist.ac.at (Institute of Science and Technology Austria)
//	
//------------------------------------------------------------------------------
//
//  This file was created as part of the implementation for 
//  Counterexample Explanation by Learning
//  Small Strategies in Markov Decision Processes
//  Brázdil et al. 2015  
//
//	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.decisiontree;

import java.util.*;


import parser.State;
import prism.PrismException;
import simulator.Choice;
import simulator.SimulatorEngine;
import simulator.TransitionList;
import strat.MDStrategy;
import heuristics.search.Heuristic;
import heuristics.search.StateValue;
import heuristics.update.StateUpdate;
import explicit.Model;
import heuristics.CollapseModelExplorer;

/**
 * DeltaStar class is used to gather data for decision tree learning from a BRTDP or PRISM strategy.
 * Data about simulations is gathered and learned in a separate procedure.
 * The model to explore is given as a model explorer.
 */
public class DeltaStar
{
	private SimulatorEngine uncollapsed;
	private CollapseModelExplorer pme;
	private StateUpdate stateUpdate;
	private HashMap<State,Integer> sampled;
	private HashMap<State,Integer> target;
	private HashMap<State,ArrayList<Integer>> stateActionGood;
	private HashMap<State,ArrayList<Integer>> internalActions;
	private int averageLength;
	private HashMap<State,Integer> numActions;
	private int globalGoods;
	private int globalBads;
	private int compareScheme;
	private int learningScheme;
	private int reachedTarget;
	private int reachedZero;
	private int reachedSelfLoop;
	private int reachedBigLoop;
	private boolean onlyLearnTarget;
	private boolean countOnlyOnce;
	private double avgGoodPerState;
	private double avgActionPerState;
	private Map<State,Integer> chosenAction;
	
	public DeltaStar(SimulatorEngine uncollapsed, CollapseModelExplorer pme, StateUpdate stateUpdate, int compareScheme, int learningScheme, boolean onlyLearnTarget, boolean countOnlyOnce) throws PrismException {
		this.uncollapsed = uncollapsed;
		this.pme = pme;
		this.stateUpdate = stateUpdate;
		sampled = new HashMap<State,Integer>();
		target = new HashMap<State,Integer>();
		stateActionGood = new HashMap<State,ArrayList<Integer>>();
		internalActions = new HashMap<State,ArrayList<Integer>>();
		averageLength = 0;
		numActions = new HashMap<State,Integer>();
		globalGoods = 0;
		globalBads = 0;
		this.compareScheme = compareScheme;
		this.learningScheme = learningScheme;
		this.setReachedTarget(0);
		this.setReachedZero(0);
		this.setReachedSelfLoop(0);
		this.setReachedBigLoop(0);
		this.onlyLearnTarget = onlyLearnTarget;
		this.countOnlyOnce = countOnlyOnce;
		this.avgGoodPerState = 0;
		this.avgActionPerState = 0;
		this.chosenAction = new HashMap<State,Integer>();
	}
	
	public void deltaStarSimulation(State initState, double delta) throws PrismException {
		deltaStarSimulation(initState,delta,null,null,false,false);
	}
	
	/**
	 * Perform one simulation on the specified model.
	 * 
	 * @param initState state to start the simulation from
	 * @param delta take actions with value in delta range from optimal one
	 * @param strat if specified, use the strat instead of brtdp strat given by stateupdate
	 * @param model the model for strat
	 * @param onlyBest if true play only best action
	 * @param stickToAction if true remember decisions and play them again later
	 * @throws PrismException
	 */
	public void deltaStarSimulation(State initState, double delta, MDStrategy strat, Model model, boolean onlyBest, boolean stickToAction) throws PrismException {
		stickToAction = false; //Disabled for the moment
		State s = initState;
		Collection<State> states;
		if (countOnlyOnce) states = new HashSet<State>();
		else states = new ArrayList<State>();
		HashMap<State,Integer> seen = new HashMap<State,Integer>();
		Integer currentSeen = 0;
		while (!stateUpdate.isTarget(s) && !stateUpdate.isZero(s) && !stateUpdate.isOnlySelfLoop(s) && !(currentSeen > 1000)) {
			currentSeen = seen.get(s);
			if (currentSeen == null) {
				currentSeen = 0;
			}
			currentSeen++;
			seen.put(s, currentSeen);
			states.add(s);
			pme.queryState(s);
			uncollapsed.initialisePath(s);
			ArrayList<Integer> goodActions = stateActionGood.get(s);
			ArrayList<Integer> intern = new ArrayList<Integer>();
			boolean stratSaysRandom = false;
			if (strat != null && model != null) { //Following strategy from PRISM
				goodActions = new ArrayList<Integer>();
				int choice = strat.getChoiceIndex(model.getStatesList().indexOf(s));
				if (choice != -2) {
					goodActions.add(choice);
					globalGoods++;
					uncollapsed.initialisePath(s);
					TransitionList trans = uncollapsed.getTransitionList();
					globalBads+= trans.getNumChoices() -1;
				}
				else {
					TransitionList trans = uncollapsed.getTransitionList();
					for (int i = 0;i < trans.getNumChoices(); i++) intern.add(i);
					globalGoods+= trans.getNumChoices() ;
					stratSaysRandom = true;
				}
				stateActionGood.put(s,goodActions);
				internalActions.put(s,intern);
			}
			else { //Following strategy of BRTDP computation
				StateValue stv = stateUpdate.getQValue(s);
				if (stv == null) {
					stateUpdate.update(s);
					stv = stateUpdate.getQValue(s);
					if (stv == null) System.out.println("Still null!");
				}
				int best = stv.getAction();
				double upperBest = Math.abs(stateUpdate.getUpperBoundActionValue(s, best).getUpperBound());
				double lowerBest = Math.abs(stateUpdate.getLowerBoundActionValue(s, best).getLowerBound());
				double v = 0.0;
				if (compareScheme == 0) v = upperBest *  (1 - (upperBest - lowerBest));
				else v = lowerBest;
				goodActions = new ArrayList<Integer>();
				TransitionList trans = uncollapsed.getTransitionList();
				if (onlyBest) { //only best action is good
					goodActions.add(best);
					globalGoods++;
					globalBads+=trans.getNumChoices()-1;
				}
				else { //iterate over the states and find all good ones
					for (int i = 0; i < trans.getNumChoices(); i++) {
						pme.queryState(s);
						int a = pme.getECActionIndex(i); // bad because quadratic (search through actions everytime)
						if (a >= 0) {
							double upperThis = Math.abs(stateUpdate.getUpperBoundActionValue(s, a).getUpperBound());
							double lowerThis = Math.abs(stateUpdate.getLowerBoundActionValue(s, a).getLowerBound());
							double vThis = 0.0;
							if (compareScheme == 0) vThis = upperThis *  (1 - (upperThis - lowerThis));
							else vThis = lowerThis;
							if (Math.abs(vThis - v) <= delta) {
								goodActions.add(i);
								globalGoods++;
							}
							else {
								globalBads++;
							}
						}
						else {
							intern.add(i);
						}
					}
				}
				stateActionGood.put(s,goodActions);
				internalActions.put(s,intern);			
			}
			goodActions = stateActionGood.get(s);
			intern = internalActions.get(s);
			int playAction = 0;
			if (stickToAction && chosenAction.containsKey(s)) { //Action was previously chosen, now we stick to it; the part before is just for counting occurances
				playAction = chosenAction.get(s);
			}
			else {
				if (goodActions.isEmpty() && !stratSaysRandom) { //Have to play uniformly in EC??
					Random r = new Random();
					playAction = intern.get(r.nextInt(intern.size()));
				}
				else if (!stratSaysRandom) {
					Random r = new Random();
					playAction = goodActions.get(r.nextInt(goodActions.size()));
				}
				else {
					Random r = new Random();
					uncollapsed.initialisePath(s);
					TransitionList trans = uncollapsed.getTransitionList();
					playAction = r.nextInt(trans.getNumChoices());
				}
				if (stickToAction) chosenAction.put(s,playAction);
			}
			s = pme.computeChoiceTargetDistrNoEC(playAction);
		}
		if (stateUpdate.isTarget(s) || !onlyLearnTarget) {
			for (State x: states) {
				Integer sam = sampled.get(x);
				if (sam == null) {
					sam = 0;
				}
				sam++;
				sampled.put(x, sam);
			}
		}
	}
	
	/**
	 * Feeds the learnt data into a decision tree.
	 * Has different learning modes, which use the number of times a state has been visited.
	 * 
	 */
	public void learn(DecisionTree tree) throws PrismException {
		int good = 0;
		int bad = 0;
		Iterator<State> i = stateActionGood.keySet().iterator();
		int numStatesSampled = 0;
		while (i.hasNext()) {
			
			State s = i.next();
			
			Integer sam = sampled.get(s);
			if (sam == null) sam = 0;
			if (sam > 0) {
				numStatesSampled++;
				uncollapsed.initialisePath(s);
				ArrayList<Integer> goods = stateActionGood.get(s);
				ArrayList<Integer> interns = internalActions.get(s);
				TransitionList trans = uncollapsed.getTransitionList();
				double goodWeight = 0;
				double badWeight = 0;
				int numChoices = trans.getNumChoices();
				if (learningScheme == 0) {
					badWeight = sam;
					goodWeight = sam;
				}
				else if (learningScheme == 1) {
					if (globalGoods > 0) goodWeight = (globalGoods + globalBads) / globalGoods;
					if (globalBads > 0) badWeight = (globalGoods + globalBads) / globalBads;
				}
				else {
					badWeight = 1;
					goodWeight = 1;
				}
				
				avgGoodPerState += goods.size();
				avgActionPerState += numChoices - interns.size();
				
				if (goods.size() > 0) {
					
					for (int j = 0; j < numChoices; j++) {
						Choice c = trans.getChoice(j);
						String module = "";
						if (c.getModuleOrActionIndex() > 0) module = "synchronised";
						else module = c.getModuleOrAction();
						String name = c.getCommandString();
						if (goods.contains(j)) {
							for (int k = 0; k < goodWeight; k++) {
								good++;
								tree.learnGood(s, module, name);
							}					
						}
						else {
							for (int k = 0; k < badWeight; k++) {
								bad++;
								tree.learnBad(s, module, name);
							}
						}
	 				}
	
				}
				else { //Only EC internal actions are good
					for (int j = 0; j < numChoices; j++) {
						Choice c = trans.getChoice(j);
						String module = "";
						if (c.getModuleOrActionIndex() > 0) module = "synchronised";
						else module = c.getModuleOrAction();
						String name = c.getCommandString();
						if (interns.contains(j)) {
							for (int k = 0; k < goodWeight; k++) {
								good++;
								tree.learnGood(s, module, name);
							}
						}
						else {
							for (int k = 0; k < badWeight; k++) {
								bad++;
								tree.learnBad(s, module, name);
							}
						}
					}
				}
			}
		}
		avgGoodPerState = avgGoodPerState / numStatesSampled;
		avgActionPerState = avgActionPerState / numStatesSampled;
	}
	

	public int getNumStateActions() {
		int out = 0;
		for (int num : numActions.values()) {
			out += num;
		}
		return out;
	}
	
	public int getAverageLength() {
		return averageLength;
	}
	
	public double getSuccessRate(State s) {
		if (target.containsKey(s)) {
			return ((double)target.get(s)) / sampled.get(s);
		}
		else {
			if (!sampled.containsKey(s)) {
				System.out.println("Not even sampled!");
			}
			return 0;
		}
	}

	public int getReachedTarget()
	{
		return reachedTarget;
	}

	public void setReachedTarget(int reachedTarget)
	{
		this.reachedTarget = reachedTarget;
	}

	public int getReachedZero()
	{
		return reachedZero;
	}

	public void setReachedZero(int reachedZero)
	{
		this.reachedZero = reachedZero;
	}

	public int getReachedSelfLoop()
	{
		return reachedSelfLoop;
	}

	public void setReachedSelfLoop(int reachedSelfLoop)
	{
		this.reachedSelfLoop = reachedSelfLoop;
	}

	public int getReachedBigLoop()
	{
		return reachedBigLoop;
	}

	public void setReachedBigLoop(int reachedBigLoop)
	{
		this.reachedBigLoop = reachedBigLoop;
	}
	
	/**
	 * Experimental code performing simulations not from the start, but playing all good actions and explore their space fully.
	 */
	public void fullDeltaExplore(State initState, double delta, Heuristic heur) throws PrismException {
		Set<State> visited = new HashSet<State>();
		Stack<State> pending = new Stack<State>();
		State s = initState;
		pending.push(s);
		int current = 0;
		while (!pending.isEmpty()  && current < 700000) {
			s = pending.pop();
			visited.add(s);
			pme.queryState(s);
			current++;
			while (!heur.isDone(s)) {
				
				heur.heuristicStep(s);
			}
			StateValue stv = stateUpdate.getQValue(s);
			if (stv == null) stv = new StateValue(0,1,0);
			double upperBest = stv.getUpperBound();
			ArrayList<Integer> goodActions = new ArrayList<Integer>();
			pme.queryState(s);
			for (int i = 0; i < pme.getNumChoices(); i++) {
				pme.queryState(s);
//				System.out.println("pme hashCode: " + pme.hashCode() + " " + pme.getNumChoices() + " state: " + s);
				double currentUp = heur.getUpper(s, i);
				pme.queryState(s);
				if (Math.abs(upperBest - currentUp) <= delta) {
					goodActions.add(i);
					for (int j = 0; j < pme.getNumTransitions(i); j++) {
						State suc = pme.computeTransitionTarget(i, j);
						if (!visited.contains(suc)) pending.push(suc);
					}
				}
			}
//			System.out.println("full heuristic explore done with: " + s);
			stateActionGood.put(s, goodActions);
			sampled.put(s, 1);
		}
	}

	public double getAvgGoodPerState()
	{
		return avgGoodPerState;
	}

	public void setAvgGoodPerState(double avgGoodPerState)
	{
		this.avgGoodPerState = avgGoodPerState;
	}

	public double getAvgActionPerState()
	{
		return avgActionPerState;
	}

	public void setAvgActionPerState(double avgActionPerState)
	{
		this.avgActionPerState = avgActionPerState;
	}
}
