//==============================================================================
//	
//	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 jdd.DebugJDD;
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_BSCC extends Heuristic{

	private Set<String> coalition;
	private Map<String,Integer> playerMapping;
	
	
	public HeuristicBRTDP_BSCC(HeuristicsSMGModelChecker mc, StateUpdate su, ECContainer eCC, HeuristicNextState nextState, ModelExplorer pme, boolean min, boolean collapseMECs) throws PrismException{
		super(mc, su, eCC, nextState, pme, min, collapseMECs);
		nonDetStates = new HashSet<State>();
		nondettotal = 0;
		if(mc.getProperty().getExpression() instanceof ExpressionPATL) {
            ExpressionPATL exp = (ExpressionPATL)mc.getProperty().getExpression();
            pme.queryState(pme.getDefaultInitialState());
            playerMapping = pme.getPlayerMapping();
            coalition = exp.getCoalition();
//            System.out.println("COALITION " + exp.getCoalition() + "\n mapping: " + pme.getPlayerMapping());
		} else {
			playerMapping = new HashMap<String,Integer>();
		    coalition = new HashSet<String>();
		    coalition.add(String.valueOf(SMG.PLAYER_1));
//		    System.out.println("COALITION 2nd case " + coalition);
		}
	}
	
	private int checkEC = 20;
	private int goBack = 19;
	
	@Override
	public void heuristicStep(State s) throws PrismException {
		Stack<State> visited = new Stack<State>();
		
		Map<State,Integer> visitedTimes = new HashMap<State,Integer>();
		int steps = 0;
		while(checkValidTrial(stateUpdate, s)) {
			visited.push(s);
//			int newVisited = visitedTimes.getOrDefault(s, 0)+1;
			pme.queryState(s);
//			if (pme.getNumChoices() != pme.getNumTransitions()) System.out.println("choices: " + pme.getNumChoices() + " != transitions" + pme.getNumTransitions());
			if (pme.getNumChoices() > 1) {
				if (!nonDetStates.contains(s)) nondettotal += pme.getNumChoices();
				nonDetStates.add(s);
			}
			Integer newVisited = visitedTimes.get(s);
			if (newVisited == null) {
				newVisited = 0;
			}
			newVisited++;
			visitedTimes.put(s, newVisited);
			if (newVisited >= checkEC) {
				int occurances = 0;
				Set<State> bscc = new HashSet<State>();
				while (occurances < goBack) { //Visited should not run out of elements if goBack <= checkEC
					State currentState = visited.pop(); 
					if (currentState.equals(s)) occurances++;
					bscc.add(currentState);
				}
				SMG partialModel = buildPartialModel(bscc);
				collapseMECs(partialModel);
				break;
			}
			visitedTimes.put(s,newVisited);
			int bestAction = stateUpdate.update(s);
			s = nextState.sample(s,bestAction);
			trialSteps++;
		}
		stateUpdate.setTargetOrZeroValue(s);
		
		while(!visited.empty()) {
			
			State vs = visited.pop();
			stateUpdate.update(vs);
		}
		
		if(trials % reportEvery == 0) {
			reportProgress(trials, trialSteps, false);
			checkEC++;
			goBack++;
		}
		if (steps > 10000) {
			System.out.println("Being in one step for 10000 times now");
		}
//		stateUpdate.deRefAll();
//		System.out.println("finishing a step");
//		DebugJDD.printNegativeTraces( );
		steps++;
		trials++;
	}
	
	private SMG buildPartialModel(Set<State> states) throws PrismException {
//		System.out.println("Building model for " + states);
		SMG partialModel = new SMG();
		List<State> stateList = new ArrayList<State>(states);
		
	    List<State> extraStates = new ArrayList<State>();
	    Map<State, Integer> state2Index = new HashMap<State, Integer>();
	    
	    Set<Integer> players = new HashSet<Integer>();
	    
//	    Iterator<State> it = states.iterator();
	    partialModel.addInitialState(0);
//	    for (int i = 0; i < stateList.size(); i++) {
//	    	State currentState = stateList.get(i);
//			state2Index.put(currentState, i);
//		}
	    int stlength = stateList.size();
	    for (int k = 0; k < stlength; k++) {
	    	State currentState = stateList.get(k);
	    	state2Index.put(currentState, k);
	    	pme.queryState(currentState);
	    	int player = pme.getPlayerForState();
	    	players.add(player);
	    	int choices = pme.getNumChoices();
	    	partialModel.addState(player);
	    	for (int i = 0; i < choices; i++) {
	    		int transitions = pme.getNumTransitions(i);
	    		Distribution d = new Distribution();
	    		for (int j = 0; j < transitions; j++) {
	    			Double prob = pme.getTransitionProbability(i, j);
	    			State target = pme.computeTransitionTarget(i, j);
	    			if (!stateList.contains(target)) {
//	    				System.out.println(target + " is extra state from: " + currentState);
	    				stateList.add(target);
	    				extraStates.add(target);
	    				state2Index.put(target, stateList.size() - 1);
	    			}
	    			Integer ind = state2Index.get(target);
	    			if (ind == null) {
	    				ind = stateList.indexOf(target);
	    			}
	    			d.add(ind, prob);
	    		}
	    		partialModel.addChoice(k, d);
	    	}
	    }
//	    System.out.println("extra states: " + extraStates);
//	    for(int k = 0; k < extraStates.size(); k++) {
//	    	State currentState = extraStates.get(k);
//	    	int stateIndex = state2Index.get(currentState);
//	    	pme.queryState(currentState);
//	    	int player = pme.getPlayerForState();
//	    	players.add(player);
//	    	partialModel.addState(player);
//	    	Distribution d = new Distribution();
//	    	d.add(stateIndex, 1.0);
//	    	partialModel.addChoice(stateIndex, d);
//	    }
	    partialModel.setStatesList(stateList);
//	    System.out.println("coalition: " + coalition + "\n" + partialModel.getPlayerMapping()+"\n" +players);
	    partialModel.setPlayerMapping(playerMapping);
	    partialModel.setCoalition(coalition);
		mc.getLog().println("Model built " + partialModel.getNumStates());
//		System.out.println(partialModel);
		long duration = System.currentTimeMillis() - modelCheckingTime;
		mc.getLog().println("Global timer: " + ((double)duration/1000)  + " secs.");
		mc.getLog().flush();
		return partialModel;
	}
	
	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(),1e-6);
		} else {
			return b1;
		}
	}
	
	private void collapseMECs(SMG smg) 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);
//				}
//			}
		}
		explicit.ECComputer ec = explicit.ECComputer.createECComputer(mc, smg);
		ec.computeMECStates();
		
		List<BitSet> mecs = ec.getMECStates();
		List<List<State>> mecStates = new ArrayList<List<State>>();
//		for (int i = 0; i < mecs.size(); i++) {
//			List<State> onemec = new ArrayList<State>();
//			for (int j = mecs.get(i).nextSetBit(0); j >= 0; j = mecs.get(i).nextSetBit(j+1)) {
//				pme.
//			}
//		}
//		System.out.println("mecs found: " + mecs);
		for(int i=0;i<mecs.size();i++) {
//			System.out.println("Having a MEC!!!!!!!!!!");
			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);
//						System.out.println("state s " + s + " is target");
					}
				} else {
					collapse(smg,mec,actions);
					for (int j = mec.nextSetBit(0); j >= 0; j = mec.nextSetBit(j+1)) {
						State s = smg.getStatesList().get(j);
						stateUpdate.update(s);
//						System.out.println("Best action for " + s + ": " + stateUpdate.getBestAction(s));
					}
				}
			} 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);
//						System.out.println("state s " + s + " is target");
					}
				} else {
				    for (int j = mec.nextSetBit(0); j >= 0; j = mec.nextSetBit(j+1)) {
                        State s = smg.getStatesList().get(j);
//                        System.out.println("state s " + s + " is zero");
                        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();
			int key = e.getKey();
			State tmp = smg.getStatesList().get(key);
			d.put(tmp, e.getValue());
		}
		return d;
	}
}