//==============================================================================
//	
//	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.search;

import java.util.ArrayList;
import java.util.BitSet;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.Stack;

import heuristics.CollapseModelExplorer;
import heuristics.ECContainer;
import heuristics.HeuristicsSMGModelChecker;
import heuristics.nextstate.HeuristicNextState;
import heuristics.update.StateUpdate;
import explicit.CachedModelExplorer;
import explicit.ConstructModel;
import explicit.Distribution;
import explicit.MDP;
import explicit.MDPSimple;
import explicit.Model;
import explicit.SMG;
import explicit.SMGModelChecker;
import explicit.ModelExplorer;
import parser.State;
import parser.ast.Expression;
import parser.ast.ExpressionPATL;
import prism.ModelType;
import prism.PrismException;

public class HeuristicBRTDP extends Heuristic{

	private boolean precomp = false;
	
	public HeuristicBRTDP(HeuristicsSMGModelChecker mc, StateUpdate su, ECContainer eCC, HeuristicNextState nextState, ModelExplorer pme, boolean min, boolean collapseMECs) throws PrismException{
		super(mc, su, eCC, nextState, pme, min, collapseMECs);
		if(mc.getProperty().getExpression() instanceof ExpressionPATL) {
            ExpressionPATL exp = (ExpressionPATL)mc.getProperty().getExpression();
            pme.queryState(pme.getDefaultInitialState());
            smg.setPlayerMapping(pme.getPlayerMapping());
            smg.setCoalition(exp.getCoalition());
            System.out.println("COALITION " + exp.getCoalition() + "\n mapping: " + pme.getPlayerMapping());
		} else {
		    Set<String> coalition = new HashSet<String>();
		    coalition.add(String.valueOf(SMG.PLAYER_1));
            smg.setCoalition(coalition);
            System.out.println("COALITION 2nd case " + coalition);
		}
	}
	
	@Override
	protected void heuristicStart() throws PrismException{
	    super.heuristicStart();
		if(precomp) {
            preCompFullModel();
	    }
	}
	
	private Set<State> seen = new HashSet<State>();
	private int buildPartialEvery = 100000;//default 100000
	private int precompCount = 0;
	private int precompThreshold = 15;
	
	@Override
	public void heuristicStep(State s) throws PrismException {
		Stack<State> visited = new Stack<State>();
		while(checkValidTrial(stateUpdate, s)) {
			seen.add(s);
			visited.add(s);
			if (pme.getNumChoices() > 1) {
				if (!nonDetStates.contains(s)) nondettotal += pme.getNumChoices();
				nonDetStates.add(s);
			}
			int bestAction = stateUpdate.update(s);
			s = nextState.sample(s,bestAction);
			if(trialSteps % buildPartialEvery == 0) {
				updatePrecomp();
				precompCount++;
				if(precompCount > precompThreshold) {
					buildPartialEvery = buildPartialEvery*10;
					precompCount = 0;
				}
			}
			trialSteps++;
		}
		seen.add(s);
		visited.add(s);
		stateUpdate.setTargetOrZeroValue(s);
		
		while(!visited.empty()) {
			State vs = visited.pop();
			stateUpdate.update(vs);
		}
		
		if(trials % reportEvery == 0) {
			reportProgress(trials, trialSteps, false);
		}
		
		trials++;
	}

	@Override
	public boolean isDone() throws PrismException {
		return isDone(initialState);
	}
	
	@Override
	public boolean isDone(State s) throws PrismException {
		StateValue sv = stateUpdate.getQValue(s);
		if(sv != null) {
			double lowerBoundInitialState = sv.getLowerBound();
			double upperBoundInitialState = sv.getUpperBound();
			return stateUpdate.same(lowerBoundInitialState, upperBoundInitialState);
		}
		return false;
	}
	
	private boolean checkValidTrial(StateUpdate stateUpdate, State s) throws PrismException{
		boolean b1 = !stateUpdate.isTarget(s) && !stateUpdate.isZero(s) && !stateUpdate.isOnlySelfLoop(s);
		if(stateUpdate.getQValue(s) != null) {
			return b1 && !stateUpdate.same(stateUpdate.getQValue(s).getLowerBound(), stateUpdate.getQValue(s).getUpperBound());
		} else {
			return b1;
		}
	}
	
	private void updatePrecomp() throws PrismException{
		SMGModelChecker smgModelChecker = getMC();
		SMG smg = buildPartialModel();
		smg.findDeadlocks(true);
		BitSet target = computeTarget(smgModelChecker, smg);
		
		BitSet prob0Target = new BitSet();
		prob0Target.or(target);
		prob0Target.or(unvisited);
		
		BitSet prob0 = computeProb0(smgModelChecker, smg, prob0Target);
		for (int i = prob0.nextSetBit(0); i >= 0; i = prob0.nextSetBit(i+1)) {
			State s = smg.getStatesList().get(i);
			stateUpdate.setZero(s, true);
		}
		/*BitSet prob1 = computeProb1(smgModelChecker, smg, target);
		for (int i = prob1.nextSetBit(0); i >= 0; i = prob1.nextSetBit(i+1)) {
			State s = smg.getStatesList().get(i);
			if(seen.contains(s)) {
				stateUpdate.setTarget(s, true);
			}
		}*/
		if(pme.getModelType() == ModelType.SMG) {
            collapseMECs(smg, unvisited);
        }
        if(pme.getModelType() == ModelType.MDP && !min) {
            collapseMECs(smg, unvisited);
        }
        
        if(precomp) {
            //compareWithPreComp(smg, prob0,prob1);
        }
	}
	
	private void compareWithPreComp(SMG smg, BitSet prob0, BitSet prob1) throws PrismException{
		  for(int i=0;i<smg.getNumStates();i++) {
		      State s1 = smg.getStatesList().get(i);
		      int p1 = smg.getPlayer(i);
		      int index = preCompSMGMapping.get(s1);
		      int p2 = preCompSMG.getPlayer(index);
		      if(p1 != p2) {
		          throw new PrismException("Player not the same i=" + i + " index=" + index);
		      }
		  }
		  
		  
		  for (int j = prob0.nextSetBit(0); j >= 0; j = prob0.nextSetBit(j+1)) {
		      State s = smg.getStatesList().get(j);
		      int index = preCompSMGMapping.get(s);
		      if(preCompZero.get(index) == false) {
		          throw new PrismException("Prob0 not the same");
		      }
	      }
		  for (int j = prob1.nextSetBit(0); j >= 0; j = prob1.nextSetBit(j+1)) {
	          State s = smg.getStatesList().get(j);
	          int index = preCompSMGMapping.get(s);
	          if(preCompOne.get(index) == false) {
	              throw new PrismException("Prob1 not the same");
	          }
		  }
	}
	
	
	private SMG smg = new SMG();
	private Map<State, Integer> state2Index = new HashMap<State, Integer>();
	private BitSet unvisited = new BitSet();
	
	private SMG buildPartialModel() throws PrismException{
	    List<State> statesList = smg.getStatesList();
		if(statesList == null) {
			statesList = new ArrayList<State>();
		}
		unvisited.clear();
		if(smg.getFirstInitialState() == -1) {
		    //Empty model
		    int index = smg.addState();
			smg.addInitialState(index);
			state2Index.put(initialState, index);
			statesList.add(initialState);
			
			pme.queryState(initialState);
			smg.setPlayer(index, pme.getPlayerForState());
		}
		Iterator<State> it = seen.iterator();
		while(it.hasNext()) {
			State s = it.next();
			Integer index = state2Index.get(s);
			if(index == null) {
				index = smg.addState();
				state2Index.put(s, index);
				statesList.add(s);
				
				pme.queryState(s);
				smg.setPlayer(index, pme.getPlayerForState());
			}
			List<Distribution> dists = buildAllDistributionsIn(smg, s, seen, statesList);
			smg.clearState(index);
			for(int i=0;i<dists.size();i++) {
				smg.addChoice(index, dists.get(i));
			}
		}
		smg.setStatesList(statesList);
		mc.getLog().println("Model built " + smg.getNumStates());
		long duration = System.currentTimeMillis() - modelCheckingTime;
		mc.getLog().println("Global timer: " + ((double)duration/1000)  + " secs.");
		mc.getLog().flush();
		return smg;
	}
	
	private List<Distribution> buildAllDistributionsIn(SMG smg, State s, Set<State> seen, List<State> statesList) throws PrismException{
		List<Distribution> dists = new ArrayList<Distribution>();
		pme.queryState(s);
		int choices = pme.getNumChoices();
		for(int i=0;i<choices;i++) {
		    pme.queryState(s);
		    Distribution d = new Distribution();
			int trans = pme.getNumTransitions(i);
			boolean inSeen = false;
			for(int j=0;j<trans;j++) {
			    pme.queryState(s);
			    double prob = pme.getTransitionProbability(i, j);
				State t = pme.computeTransitionTarget(i,j);
				inSeen = inSeen || seen.contains(t);
				Integer index = state2Index.get(t);
				if(index == null) {
					index = smg.addState();
					state2Index.put(t, index);
					statesList.add(t);
					
					pme.queryState(t);
					smg.setPlayer(index, pme.getPlayerForState());
				}
				d.add(index, prob);
				if(!seen.contains(t)) {
					unvisited.set(index);
				}
			}
			dists.add(d);
		}
		return dists;
	}
	
	private BitSet computeProb0(SMGModelChecker smgModelChecker, SMG smg, BitSet target) throws PrismException{
		return smgModelChecker.prob0(smg, null, target, min, !min);
	}
	
	private BitSet computeProb1(SMGModelChecker smgModelChecker, SMG smg, BitSet target) throws PrismException{
		return smgModelChecker.prob1(smg, null, target, min, !min);
	}
	
	private BitSet computeTarget(SMGModelChecker smgModelChecker, SMG smg) throws PrismException {
		Expression targetExp = null;
		if(mc.getExpression().getOperand1() != null) {
			targetExp = mc.getExpression();
		} else {
			targetExp = mc.getExpression().getOperand2();
		}
		return smgModelChecker.checkExpression(smg, targetExp).getBitSet();
	}
	
	private SMGModelChecker getMC() throws PrismException{
		SMGModelChecker smgModelChecker = new SMGModelChecker(mc);
		smgModelChecker.setModulesFileAndPropertiesFile(mc.getModulesFile(), mc.getPropertiesFile());
		return smgModelChecker;
	}
	
	private void collapseMECs(SMG smg, BitSet unvisited) throws PrismException {
		if(!collapseMECs) {
			mc.getLog().println("MEC collapsing disabled");
			mc.getLog().flush();
			return;
		}
		mc.getLog().println("Starting MEC collapsing");
		mc.getLog().flush();
		long start = System.currentTimeMillis();
		BitSet all = new BitSet();
		BitSet notOne = new BitSet();
		for(int i=0;i<smg.getNumStates();i++) {
			all.set(i);
			StateValue sv = stateUpdate.getQValue(smg.getStatesList().get(i));
			if(sv != null) {
				if(sv.getUpperBound() < 1.0) {
					notOne.set(i);
				}
			}
		}
		//all.xor(notOne);
		all.xor(unvisited);
		explicit.ECComputer ec = explicit.ECComputer.createECComputer(mc, smg);
		ec.computeMECStates(all);
		
		List<BitSet> mecs = ec.getMECStates();
		for(int i=0;i<mecs.size();i++) {
			BitSet mec = mecs.get(i);
			int player = checkAllSamePlayers(smg, mec);
			List<Map<State, Double>> actions = getAllLeavingMEC(smg, mec);
			if(actions.size() > 0) {
				if(!min && player == SMG.PLAYER_1 && containsTarget(smg, mec)) {
					for (int j = mec.nextSetBit(0); j >= 0; j = mec.nextSetBit(j+1)) {
						State s = smg.getStatesList().get(j);
						stateUpdate.setTarget(s, true);
					}
				} else {
					collapse(smg,mec,actions);
				}
			} else {
				if(containsTarget(smg, mec)) {
					for (int j = mec.nextSetBit(0); j >= 0; j = mec.nextSetBit(j+1)) {
						State s = smg.getStatesList().get(j);
						stateUpdate.setTarget(s, true);
					}
				} else {
				    for (int j = mec.nextSetBit(0); j >= 0; j = mec.nextSetBit(j+1)) {
                        State s = smg.getStatesList().get(j);
                        stateUpdate.setZero(s, true);
                    }
				}
			}
		}
		long duration = System.currentTimeMillis() - start;
		mc.getLog().println("MEC collapsing done " + (double)duration/1000 + " secs.");
		mc.getLog().flush();
	}
	
	private int checkAllSamePlayers(SMG smg, BitSet mec) throws PrismException {
	    int player = -1;
	    for (int i = mec.nextSetBit(0); i >= 0; i = mec.nextSetBit(i+1)) {
            if(player == -1) {
              player = smg.getPlayer(i);
            } else {
                State s = smg.getStatesList().get(i);
                
                if(!stateUpdate.isTarget(s) && smg.getNumChoices(i) > 1) {
                    int currentPlayer = smg.getPlayer(i);
                    if(currentPlayer != player) {
                        throw new PrismException("Only MECs of the same player are allowed");
                    }
                }
            }
        }
	    return player;
	}
	
	private boolean containsTarget(SMG smg, BitSet mec) throws PrismException{
		for (int i = mec.nextSetBit(0); i >= 0; i = mec.nextSetBit(i+1)) {
			State s = smg.getStatesList().get(i);
			if(stateUpdate.isTarget(s)) {
				return true;
			}
		}
		return false;
	}
	
	private void collapse(SMG smg, BitSet mec, List<Map<State, Double>> actions) throws PrismException {
		if (pme instanceof CachedModelExplorer) {
			for (int j = mec.nextSetBit(0); j >= 0; j = mec.nextSetBit(j+1)) {
				collapse(smg, j, actions);
			}
		}
		else if (pme instanceof CollapseModelExplorer) {
			((CollapseModelExplorer) pme).collapse(smg, mec, actions);
		}
	}
	
	private void collapse(SMG smg, int s, List<Map<State, Double>> actions) throws PrismException {
		State state = smg.getStatesList().get(s);
		if (pme instanceof CachedModelExplorer) {
			CachedModelExplorer cme = (CachedModelExplorer)pme;
			cme.updateCache(state, actions);
		}
		
	}
	
	private List<Map<State, Double>> getAllLeavingMEC(SMG smg, BitSet mec) {
		List<Map<State, Double>> actions = new ArrayList<Map<State, Double>>();
		for (int s = mec.nextSetBit(0); s >= 0; s = mec.nextSetBit(s+1)) {
			for(int i = 0; i < smg.getNumChoices(s);i++) {
				boolean all = smg.allSuccessorsInSet(s, i, mec);
				if(!all) {
					actions.add(getDistribution(smg, s, i));
				}
			}
		}
		return actions;
	}
	
	private Map<State, Double> getDistribution(SMG smg, int s, int i) {
		Map<State, Double> d = new HashMap<State, Double>();
		Iterator<Entry<Integer,Double>> it =  smg.getTransitionsIterator(s, i);
		while(it.hasNext()) {
			Entry<Integer, Double> e = it.next();
			d.put(smg.getStatesList().get(e.getKey()), e.getValue());
		}
		return d;
	}
	
	private BitSet preCompTarget = new BitSet();
	private BitSet preCompZero = new BitSet();
	private BitSet preCompOne = new BitSet();
	private SMG preCompSMG;
	private Map<State,Integer> preCompSMGMapping = new HashMap<State,Integer>();
	
	private void preCompFullModel() throws PrismException{
		ConstructModel constructModel = new ConstructModel(mc, mc.getEngine(), mc.getProperty(), mc.getPropertiesFile());
		Model model = (Model)constructModel.constructModel(mc.getModulesFile(), false, false);
			
		for(int i=0;i<model.getNumStates();i++) {
		    preCompSMGMapping.put(model.getStatesList().get(i), i);
		}			
		if(model instanceof SMG) {
		    preCompSMG = (SMG)model;
			ExpressionPATL exp = (ExpressionPATL)mc.getProperty().getExpression();
			preCompSMG.setCoalition(exp.getCoalition());
		} else {
			if(model instanceof MDP) {
			    preCompSMG = new SMG((MDPSimple)model);
			} else {
				throw new UnsupportedOperationException();
			}
		}

		SMGModelChecker smgModelChecker = getMC();
		preCompTarget = computeTarget(smgModelChecker, preCompSMG);
		preCompZero = computeProb0(smgModelChecker, preCompSMG, preCompTarget);
		preCompOne = computeProb1(smgModelChecker, preCompSMG, preCompTarget);
			
		List<State> statesList = preCompSMG.getStatesList();
		for(int i=0;i<statesList.size();i++) {
			State s = statesList.get(i);
			if(preCompTarget.get(i) || preCompOne.get(i)) {
				stateUpdate.setTarget(s, true);
			} else {
				stateUpdate.setTarget(s, false);
			}
			if(preCompZero.get(i)) {
				stateUpdate.setZero(s, true);
			}
			pme.queryState(s);
		}
		if(pme.getModelType() == ModelType.SMG) {
		    collapseMECs(preCompSMG, new BitSet());
		}
		if(pme.getModelType() == ModelType.MDP && !min) {
		    collapseMECs(preCompSMG, new BitSet());
		}
	}
}