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

import heuristics.search.StateValue;

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.Random;
import java.util.Set;

import explicit.ModelExplorer;
import explicit.SMG;
import parser.State;
import parser.Values;
import parser.ast.ModulesFile;
import prism.ModelType;
import prism.PrismException;
import prism.PrismLog;
import simulator.Choice;
import simulator.SimulatorEngine;

/*
* A wrapper to ModelExplorer that caches collapsed states.
* 
*/

public class CollapseModelExplorer implements ModelExplorer {
		
	private class Transition {
		public final State state;
		public final double probability;
		public Transition(State s, double p) {
			state = s;
			probability = p;
		}
	}
	
	private class Action {
		public final List<Transition> transitions;
		public final String name;
		public final String module;
		public Action(List<Transition> transitions, String name, String module) {
			this.transitions = transitions;
			this.name = name;
			this.module = module;
		}
	}
	
//	public ModelExplorer modelExplorer = null;
	public SimulatorEngine simEngine = null;
	private ModulesFile modulesFile = null;
	private State currentState = null;
	private int ecCount = 0;
	private ECContainer eCs;
	private Map<Integer,List<Action>> ecActions = new HashMap<Integer,List<Action>>();
	private Map<Integer,ArrayList<String>> actionsPerModule = new HashMap<Integer,ArrayList<String>>();
	private List<String> moduleIndex = new ArrayList<String>();
	private int maxAction;
	
	public CollapseModelExplorer(SimulatorEngine simEngine, ModulesFile modulesFile, ECContainer eCC) throws PrismException {
//		System.out.println("Collapse MEC used");
		this.eCs = eCC;
//		modelExplorer = me;
		this.simEngine = simEngine; 
		simEngine.createNewOnTheFlyPath(modulesFile);
		this.modulesFile = modulesFile;
		this.maxAction = 0;
	}
	
	public void collapse(SMG smg, BitSet mec, List<Map<State, Double>> actions) throws PrismException {
		List<State> states = smg.getStatesList();
//		if (actions.isEmpty()) System.out.println("collapsing with no actions!");
//		System.out.println("collapsing " + mec + " with actions " + actions);
		List<Action> ecAction = new ArrayList<Action>();
//		System.out.println("One MEC:");
		for (int j = mec.nextSetBit(0); j >= 0; j = mec.nextSetBit(j+1)) {
			State s = states.get(j);
//			System.out.println("putting state: " + s + " into EC");
			queryState(s);
//			System.out.println(s + " ~ numC: " + getNumChoices());
			for (int i = 0; i < getNumChoices(); i++) {
//				boolean all = smg.allSuccessorsInSet(j, i, mec);
				boolean all = true;
				List<Transition> currChoice = new ArrayList<Transition>();
				for (int k = 0; k < getNumTransitions(i); k++) {
					State target = computeTransitionTarget(i,k);		
					double probability = getTransitionProbability(i,k);
					currChoice.add(new Transition(target,probability));
					if (!mec.get(states.indexOf(target))) {
						all = false;
					}
				}
				if (!all) {
					String actionName = getActionName(i);
					String module = getModuleName(i);
//					System.out.println("leaving: " + actionName);
					ecAction.add(new Action(currChoice,actionName,module));
				}
			}
			eCs.insert(s, ecCount);
//			System.out.println("ECS get after insert: " + eCs.getEC(s));
		}
//		System.out.println("actions: " + actions.size() + " produced actions: " + ecAction.size());
		ecActions.put(ecCount, ecAction);
		ecCount++;
//		System.out.println("ECS::::::::::::::::::::::"+eCs);
	}
	
	public String getActionName(int i) throws PrismException {
		Integer ec = eCs.getEC(currentState);
		if (ec != null) {
			List<Action> actions = ecActions.get(ec);
			if (actions != null) {
				return actions.get(i).name;
			}
			else throw new PrismException("Inconsistent EC and actions maps in getActionName: " + currentState + " is supposed to be in EC; ecActions: " + ecActions);
		}
		else {
			Choice choice = getChoice(i);
//			String synchAction = simEngine.getTransitionAction(i);
//			if (synchAction == null) {
//				synchAction = choice.getModuleOrAction() + " --> " + choice;
//			}
//			System.out.println("choice class: "  + choice.getClass());
//			String synchAction = choice.getModuleOrAction();
//			System.out.println("ModuleOrAction: " + choice.getCommandString());
//			choice.getUpdateStringFull(i)
//			if (choice.getModuleOrActionIndex() < 0) {
//				synchAction = choice.getModuleOrAction() + " -> " + choice;
//			}
			return choice.getCommandString();
		}
	}
	
	public int getModuleIndex(int i) throws PrismException {
		return -1;
	}
	
//	public String getModuleOrAction(int i) throws PrismException {
//		Choice choice = getChoice(i);
//		return choice.getModuleOrAction();
//	}
	
	public int getGlobalActionIndex(int i) throws PrismException {
		return getGlobalActionIndex(getActionName(i));
	}
	
	private int moduleIndex(String action) {
		String module = action2Module(action);
		if (moduleIndex.contains(module)) {
			return moduleIndex.indexOf(module);
		}
		else {
//			moduleIndex.add(module);
//			return moduleIndex.size() - 1;
			return -1;
		}
	}
	
	private String action2Module(String action) {
		String moduleString = "";
		if (action.charAt(0) == '[') {
			moduleString = "synch";
		}
		else {
			moduleString = action.split("-")[0];
		}
		return moduleString;
	}
	
	private int insertModule(String action) {
		String module = action2Module(action);
		if (moduleIndex.contains(module)) {
			return moduleIndex.indexOf(module);
		}
		else {
			moduleIndex.add(module);
			return moduleIndex.size() - 1;
		}
	}
	
	public int getNumAllocatedActions() {
		return maxAction + 1; //allocated actions + undef
	}
	
	public void insertAction(int i) throws PrismException {
		insertAction(getActionName(i));
	}
	
	public void insertAction(String name) {
		int module = insertModule(name);
		ArrayList<String> actionsOfModule = actionsPerModule.get(module);
		if (actionsOfModule == null) {
			ArrayList<String> newList = new ArrayList<String>();
			newList.add(name);
			actionsPerModule.put(module, newList);
		}
		else if (!actionsOfModule.contains(name)) {
			actionsOfModule.add(name);
		}
	}
	
	public int getGlobalActionIndex(String name) throws PrismException {
		int module = moduleIndex(name);
		int index = 0;
		ArrayList<String> actionsOfModule = actionsPerModule.get(module);
		ArrayList<String> tmp;
		if (actionsOfModule == null) {
			return StateValue.ACTION_UNDEF;
		}
		else {
			index = actionsOfModule.indexOf(name);
		}
		int out = 0;
		for (int j = 0; j < module; j++) {
			tmp = actionsPerModule.get(j);
			if (tmp != null) {
				out += tmp.size();
			}
		}
		out = out + index;
		if (out > maxAction) {
			maxAction = out;
		}
		return out;
	}
	
	public void printActionInfo(PrismLog log) throws PrismException {
		ArrayList<String> actions;
		for (Integer module: actionsPerModule.keySet()) {
			log.println("Actions for module " + moduleIndex.get(module) + ": ");
			actions = actionsPerModule.get(module);
			for (String action: actions) {
				log.println(getGlobalActionIndex(action) + " " + action);
			}
		}
	}
	
	public weka.core.FastVector getActionVector() {
		weka.core.FastVector vector = new weka.core.FastVector(maxAction);
		ArrayList<String> actions;
		for (Integer module: actionsPerModule.keySet()) {
			actions = actionsPerModule.get(module);
			for (String action: actions) {
				vector.addElement(action);
			}
		}
		return vector;
	}
	
	public String getModuleName(int i) throws PrismException {
		Integer ec = eCs.getEC(currentState);
		if (ec != null) {
			List<Action> actions = ecActions.get(ec);
			if (actions != null) {
				return actions.get(i).module;
			}
			else throw new PrismException("Inconsistent EC and actions maps in getModuleIndex: " + currentState + " is supposed to be in EC; ecActions: " + ecActions);
		}
		else {
			Choice choice = getChoice(i);
			if (choice.getModuleOrActionIndex() > 0) return "synchronised";
			else return choice.getModuleOrAction();
		}
	}
	
	public SimulatorEngine getSimEngine() {
		return simEngine;
	}
	
	public ModulesFile getModulesFile() {
		return modulesFile;
	}
	
	@Override
	public State getDefaultInitialState() throws PrismException{
		return modulesFile.getDefaultInitialState();
	}
	
	@Override
	public void queryState(State state) throws PrismException {
		currentState = state;
		simEngine.initialisePath(state);
	}
	
	@Override
	public int getNumChoices() throws PrismException{
		Integer ec = eCs.getEC(currentState);
		if (ec != null) {
			List<Action> actions = ecActions.get(ec);
			if (actions != null) {
//				System.out.println("state: " + currentState + " is in EC, has: " + actions.size() + " many choices");
				return actions.size();
			}
			else throw new PrismException("Inconsistent EC and transition maps in getNumChoices: " + currentState + " is supposed to be in EC; ecActions: " + ecActions);
		}
		else {
//			System.out.println("state: " + currentState + " is !NOT! in EC, has: " + simEngine.getNumChoices() + " many choices");
			return simEngine.getNumChoices();
		}
	}
	
	public int getNumTransitionsNoEC(int i) throws PrismException {
		return simEngine.getNumTransitions(i);
	}
	
	@Override
	public int getNumTransitions(int i) throws PrismException {
		Integer ec = eCs.getEC(currentState);
		if (ec != null) {
			List<Action> actions = ecActions.get(ec);
			if (actions != null) {
				try {
					List<Transition> distr = actions.get(i).transitions;
					if (distr != null) {
						return distr.size();
					}
					else throw new PrismException("No distribution for some transition");
				}
				catch (PrismException e) {
//					System.out.println("Error when gettint number of transitions for: " + currentState + "\neC: " + eCs.get(currentState) + "\nchoices: " + getNumChoices());
					throw(e);
				}
				
			}
			else throw new PrismException("Inconsistent EC and transition maps");
		}
		else {
			return simEngine.getNumTransitions(i);
		}
	}
	
//	public List<String> getActionIndizes() {
//		ArrayList<String> out = new ArrayList<String>();
//		for (int i = 0; i < )
//	}
	
	public int getECActionIndex(int a) throws PrismException {
		Integer ec = eCs.getEC(currentState);
		if (ec != null) {
			String originalName = getChoice(a).getCommandString();
//			System.out.println("Original name: " + originalName);
			List<Action> actions = ecActions.get(ec);
			if (actions != null) {
				for (int i = 0; i < actions.size(); i++) {
//					System.out.println("Candidate: " + actions.get(i).name);
					if (actions.get(i).name.equals(originalName)) {
//						System.out.println("action: " + a + " corresponds to EC action: " + i);
						return i;
					}
				}
				return -1; //the action is an EC internal action
			}
			else throw new PrismException("Inconsistent EC and transition maps in getNumChoices: " + currentState + " is supposed to be in EC; ecActions: " + ecActions);
		}
		else {
//			String originalName = getChoice(a).getCommandString();
//			System.out.println(originalName + " is an action of a non EC state! " + currentState + " ~ " + eCs.getEC(currentState));
			return a;
		}
	}
	
	public int getNumChoicesNoEC() throws PrismException {
		return simEngine.getNumChoices();
	}
	
	public boolean isInEC() {
		Integer ec = eCs.getEC(currentState);
		if (ec != null) {
			return true;
		}
		else return false;
	}
	
	public double getTransitionProbabilityNoEC(int i, int offset) throws PrismException {
		return simEngine.getTransitionProbability(i, offset);
	}
	
	@Override
	public double getTransitionProbability(int i, int offset) throws PrismException{
		Integer ec = eCs.getEC(currentState);
		if (ec != null) {
			List<Action> actions = ecActions.get(ec);
			if (actions != null) {
				List<Transition> distr = actions.get(i).transitions;
				if (distr != null) {
					Double prob = distr.get(offset).probability;
					if (prob != null) {
						return prob;
					}
					else throw new PrismException("Distribution does not contain offset");
				}
				else throw new PrismException("No distribution for some transition");
			}
			else throw new PrismException("Inconsistent EC and transition maps");
		}
		else return simEngine.getTransitionProbability(i, offset);
	}
	
	@Override
	public State computeTransitionTarget(int i, int offset) throws PrismException{
		Integer ec = eCs.getEC(currentState);
		if (ec != null) {
			List<Action> actions = ecActions.get(ec);
			if (actions != null) {
				List<Transition> distr = actions.get(i).transitions;
				if (distr != null) {
					State s = distr.get(offset).state;
					if (s != null) {
						return s;
					}
					else throw new PrismException("Distribution does not contain offset");
				}
				else throw new PrismException("No distribution for some transition");
			}
			else throw new PrismException("Inconsistent EC and transition maps");
		}
		else return simEngine.computeTransitionTarget(i, offset);
	}
	
	@Override
	public int getNumTransitions() throws PrismException{
		Integer ec = eCs.getEC(currentState);
		if (ec != null) {
			List<Action> actions = ecActions.get(ec);
			if (actions != null) {
				int sum = 0;
				for (int i = 0; i < actions.size(); i++) {
					sum += actions.get(i).transitions.size();
				}
				return sum;
			}
			else throw new PrismException("Inconsistent EC and transition maps");
		}
		else return simEngine.getNumTransitions();
	}
	
	@Override
	public int getPlayerForState() throws PrismException{
		if (getModelType() == ModelType.SMG) {
			int player = -1;
			int nc = getNumChoices();
			for (int i = 0; i < nc; i++) {
				int iPlayer = determinePlayerForChoice(modulesFile, i);
				if (player != -1 && iPlayer != player) {
					throw new PrismException("Choices for both player " + player + " and " + iPlayer);
				}
				player = iPlayer;
			}
			return player;
		}
		else return SMG.PLAYER_1;
	}
	
	@Override
	public Map<String, Integer> getPlayerMapping() throws PrismException{
		if (getModelType() == ModelType.SMG) {
			Map<String, Integer> players = new HashMap<String, Integer>();
	        for (int i = 0; i < modulesFile.getNumPlayers(); i++) {
	            players.put(modulesFile.getPlayer(i).getName(), i + 1);
		    }
		    return players;
		}
		else {
			Map<String, Integer> playerMapping = new HashMap<String, Integer>();
            playerMapping.put(String.valueOf(SMG.PLAYER_1), SMG.PLAYER_1);
			return playerMapping;
		}
	}
	
	private int determinePlayerForChoice(ModulesFile modulesFile, int i) throws PrismException
	{
		int modAct, player;

		modAct = simEngine.getTransitionModuleOrActionIndex(i, 0);
		// Synchronous action
		if (modAct > 0) {
			player = modulesFile.getPlayerForAction(modulesFile.getSynch(modAct - 1));
			if (player == -1) {
				throw new PrismException("Action \"" + modulesFile.getSynch(modAct - 1) + "\" is not assigned to any player");
			}
			// 0-indexed to 1-indexed
			player++;
		}
		// Asynchronous action
		else {
			player = modulesFile.getPlayerForModule(simEngine.getTransitionModuleOrAction(i, 0));
			if (player == -1) {

				// for backwards compatibility trying to parse player from the
				// module name (e.g., playerX)
				try {
					player = Integer.parseInt(simEngine.getTransitionModuleOrAction(i, 0).substring(6));
				} catch (Exception e) {
					throw new PrismException("Module \"" + simEngine.getTransitionModuleOrAction(i, 0) + "\" is not assigned to any player");
				}
			}
			// 0-indexed to 1-indexed
			player++;
		}

		return player;
	}
	
	public State computeChoiceTargetDistr(int action) throws PrismException {
		double[] distr = new double[getNumTransitions(action)];
		for (int i = 0; i < getNumTransitions(action); i++) {
			distr[i] = getTransitionProbability(action, i);
		}
		return sampleFromDist(action,distr);
	}
	
	public State computeChoiceTargetDistrNoEC(int action) throws PrismException {
		double[] distr = new double[getNumTransitionsNoEC(action)];
		for (int i = 0; i < getNumTransitionsNoEC(action); i++) {
			distr[i] = getTransitionProbabilityNoEC(action, i);
		}
		return sampleFromDistNoEC(action,distr);
	}
	
	protected State sampleFromDistNoEC(int action, double[] dist) throws PrismException {
		for(int i=1;i<dist.length;i++) {
			dist[i] = dist[i-1] + dist[i];
		}
		
		double rand = Math.random();
		for(int i=0;i<dist.length;i++) {
			if(rand < dist[i]) {
				return simEngine.computeTransitionTarget(action, i);
			}
		}
		throw new PrismException("Not sampling any state, this should never happen");
	}
	
	protected State sampleFromDist(int action, double[] dist) throws PrismException {
		for(int i=1;i<dist.length;i++) {
			dist[i] = dist[i-1] + dist[i];
		}
		
		double rand = Math.random();
		for(int i=0;i<dist.length;i++) {
			if(rand < dist[i]) {
				return computeTransitionTarget(action, i);
			}
		}
		throw new PrismException("Not sampling any state, this should never happen");
	}
	
	@Override
	public ModelType getModelType() throws PrismException{
		return modulesFile.getModelType();
	}
	
	@Override
	public void queryState(State state, double time) throws PrismException{
		queryState(state);
	}
	
	@Override
	public double getTransitionProbability(int i) throws PrismException{
		throw new UnsupportedOperationException();
	}
	
	@Override
	public State computeTransitionTarget(int i) throws PrismException{
		throw new UnsupportedOperationException();
	}
	
	@Override
	public String getTransitionAction(int i, int offset) throws PrismException{
		throw new UnsupportedOperationException();
	}
	
	@Override
	public String getTransitionAction(int i) throws PrismException{
		return simEngine.getTransitionAction(i);
	}
	
	public String getTransitionUpdateStringFull(int i) throws PrismException {
		return simEngine.getTransitionUpdateStringFull(i);
	}
	
	public Choice getChoice(int i) throws PrismException {
		return simEngine.getTransitionList().getChoice(i);
	}
	
	public int getTransitionModuleOrActionIndex(int i) throws PrismException  {
		return simEngine.getTransitionModuleOrActionIndex(i);
	}
	
	public String getTransitionModuleOrAction(int i) throws PrismException {
		return simEngine.getTransitionModuleOrAction(i);
	}
	
//	public Values bla() {
//		return simEngine.getCurrentState();
//	}
}
