//==============================================================================
//	
//	Copyright (c) 2002-
//	Authors:
//	* Dave Parker <david.parker@comlab.ox.ac.uk> (University of Oxford)
//	* Vojtech Forejt <vojtech.forejt@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 explicit;

import java.io.File;
import java.io.FileNotFoundException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import parser.State;
import parser.Values;
import parser.VarList;
import parser.ast.Expression;
import parser.ast.ExpressionPATL;
import parser.ast.ExpressionTemporal;
import parser.ast.ModulesFile;
import parser.ast.PropertiesFile;
import parser.ast.Property;
import prism.ModelType;
import prism.Prism;
import prism.PrismComponent;
import prism.PrismException;
import prism.PrismLangException;
import prism.PrismLog;
import prism.PrismPrintStreamLog;
import prism.ProgressDisplay;
import prism.UndefinedConstants;
import simulator.SimulatorEngine;

public class ConstructModel extends PrismComponent
{
	// The simulator engine
	protected SimulatorEngine engine;

	// Options:
	// Find deadlocks during model construction?
	protected boolean findDeadlocks = true;
	// Automatically fix deadlocks?
	protected boolean fixDeadlocks = true;

	// Details of built model
	protected List<State> statesList;
	
	protected Property property;
	protected PropertiesFile propertiesFile;

	public ConstructModel(PrismComponent parent, SimulatorEngine engine, Property property, PropertiesFile propertiesFile) throws PrismException
	{
		super(parent);
		this.engine = engine;
		this.property = property;
		this.propertiesFile = propertiesFile;
	}

	public List<State> getStatesList()
	{
		return statesList;
	}

	public void setFixDeadlocks(boolean b)
	{
		fixDeadlocks = b;
	}

	/**
	 * Build the set of reachable states for a PRISM model language description and return.
	 * @param modulesFile The PRISM model
	 */
	public List<State> computeReachableStates(ModulesFile modulesFile) throws PrismException
	{
		constructModel(modulesFile, true, false);
		return statesList;
	}

	/**
	 * Construct an explicit-state model from a PRISM model language description and return.
	 * @param modulesFile The PRISM model
	 */
	public Model constructModel(ModulesFile modulesFile) throws PrismException
	{
		return constructModel(modulesFile, false, true);
	}

	/**
	 * Construct an explicit-state model from a PRISM model language description and return.
	 * If {@code justReach} is true, no model is built and null is returned;
	 * the set of reachable states can be obtained with {@link #getStatesList()}.
	 * @param modulesFile The PRISM model
	 * @param justReach If true, just build the reachable state set, not the model
	 * @param buildSparse Build a sparse version of the model (if possible)?
	 */
	public Model constructModel(ModulesFile modulesFile, boolean justReach, boolean buildSparse) throws PrismException
	{
		return constructModel(modulesFile, justReach, buildSparse, true);
	}

	/**
	 * Construct an explicit-state model from a PRISM model language description and return.
	 * If {@code justReach} is true, no model is built and null is returned;
	 * the set of reachable states can be obtained with {@link #getStatesList()}.
	 * @param modulesFile The PRISM model
	 * @param justReach If true, just build the reachable state set, not the model
	 * @param buildSparse Build a sparse version of the model (if possible)?
	 * @param distinguishActions True if actions should be attached to distributions (and used to distinguish them)
	 */
	public Model constructModel(ModulesFile modulesFile, boolean justReach, boolean buildSparse, boolean distinguishActions) throws PrismException
	{
		// Model info
		ModelType modelType;
		// State storage
		StateStorage<State> states;
		LinkedList<State> explore;
		State state, stateNew;
		// Explicit model storage
		ModelSimple modelSimple = null;
		DTMCSimple dtmc = null;
		CTMCSimple ctmc = null;
		MDPSimple mdp = null;
		STPGExplicit stpg = null;
		SMG smg = null;
		ModelExplicit model = null;
		Distribution distr = null;
		// Misc
		int i, j, nc, nt, src, dest, player;
		long timer;
		// int id;
		int nPlayers = 0;
		State tempState;

		// SMG vars
		Map<Integer, List<Integer>> playerChoices = new HashMap<Integer, List<Integer>>();
		List<Integer> originalStates = new ArrayList<Integer>();
		List<Integer> choices;
		boolean lazy = false;
		int schedIndx = 0;

		// Don't support multiple initial states
		if (modulesFile.getInitialStates() != null) {
			throw new PrismException("Cannot do explicit-state reachability if there are multiple initial states");
		}

		// Display a warning if there are unbounded vars
		VarList varList = modulesFile.createVarList();
		if (varList.containsUnboundedVariables())
			mainLog.printWarning("Model contains one or more unbounded variables: model construction may not terminate");
		
		// Starting reachability...
		mainLog.print("\nComputing reachable states...\n");
		mainLog.flush();
		ProgressDisplay progress = new ProgressDisplay(mainLog);
		progress.start();
		timer = System.currentTimeMillis();

		// Initialise simulator for this model
		modelType = modulesFile.getModelType();

		engine.createNewOnTheFlyPath(modulesFile);

		// Create model storage
		if (!justReach) {
			// Create a (simple, mutable) model of the appropriate type
			switch (modelType) {
			case DTMC:
				modelSimple = dtmc = new DTMCSimple();
				break;
			case CTMC:
				modelSimple = ctmc = new CTMCSimple();
				break;
			case MDP:
				modelSimple = mdp = new MDPSimple();
				break;
			case STPG:
				modelSimple = stpg = new STPGExplicit();
				break;
			case SMG:
				modelSimple = smg = new SMG();
				break;
			}
		}

		// Initialise states storage
		states = new IndexedSet<State>(true);
		explore = new LinkedList<State>();
		// Add initial state(s) to 'explore'
		// Easy (normal) case: just one initial state
		Map<State, Integer> state2Depth = new HashMap<State,Integer>();
		if (modulesFile.getInitialStates() == null) {
			state = modulesFile.getDefaultInitialState();
			explore.add(state);
			state2Depth.put(state,0);
		}
		// Otherwise, there may be multiple initial states
		// For now, we handle this is in a very inefficient way
		else {
			Expression init = modulesFile.getInitialStates();
			List<State> allPossStates = varList.getAllStates();
			for (State possState : allPossStates) {
				if (init.evaluateBoolean(modulesFile.getConstantValues(), possState)) {
					explore.add(possState);
					state2Depth.put(possState,0);
				}
			}
		}
		// Copy initial state(s) to 'states' and to the model
		for (State initState : explore) {
			states.add(initState);
			if (!justReach) {
				modelSimple.addState();
				modelSimple.addInitialState(modelSimple.getNumStates() - 1);
			}
		}
		
		int maxDepth = getMaxModelExplorationDepth();
		
		src = -1;
		while (!explore.isEmpty()) {
			// Pick next state to explore
			// (they are stored in order found so know index is src+1)
			state = explore.removeFirst();
			int depth = state2Depth.get(state);
			if(depth > maxDepth) continue;
			depth++;
			src++;

			// Use simulator to explore all choices/transitions from this state
			engine.initialisePath(state);
			nc = engine.getNumChoices();
			// For an STPG, first determine which player owns the state
			if (modelType == ModelType.STPG || modelType == ModelType.SMG) {
				player = -1;
				
				for (i = 0; i < nc; i++) {
					
					int iPlayer = determinePlayerForChoice(modulesFile, modelType, i);
					if (player != -1 && iPlayer != player) {
						throw new PrismException("Choices for both player " + player + " and " + iPlayer + " in state " + state);
					}
					player = iPlayer;
				}
				if (modelType == ModelType.STPG)
					stpg.setPlayer(src, player);
				else
					smg.setPlayer(src, player);
			}
			simulator.TransitionList trans;
			// Look at each outgoing choice in turn
			for (i = 0; i < nc; i++) {
				trans = engine.getTransitionList();
				if (!justReach && (modelType == ModelType.MDP || modelType == ModelType.STPG || modelType == ModelType.SMG)) {
					distr = new Distribution();
				}
				// Look at each transition in the choice
				nt = engine.getNumTransitions(i);
				for (j = 0; j < nt; j++) {
					stateNew = engine.computeTransitionTarget(i, j);
					// Is this a new state?
					if (states.add(stateNew)) {
						// If so, add to the explore list
						explore.add(stateNew);
						// And to model
						if (!justReach) {
							modelSimple.addState();
						}
						
						state2Depth.put(stateNew, depth);
					}
					// Get index of state in state set
					dest = states.getIndexOfLastAdd();
					// Add transitions to model
					if (!justReach) {
						switch (modelType) {
						case DTMC:
							dtmc.addToProbability(src, dest, engine.getTransitionProbability(i, j));
							break;
						case CTMC:
							ctmc.addToProbability(src, dest, engine.getTransitionProbability(i, j));
							break;
						case MDP:
						case STPG:
						case SMG:
							distr.add(dest, engine.getTransitionProbability(i, j));
							break;
						}
					}
				}
				if (!justReach) {
					if (modelType == ModelType.MDP) {
						if (distinguishActions) {
							mdp.addActionLabelledChoice(src, distr, trans.getChoice(i));
						} else {
							mdp.addChoice(src, distr);
						}
					} else if (modelType == ModelType.STPG) {
						if (distinguishActions) {
							stpg.addActionLabelledChoice(src, distr, engine.getTransitionAction(i, 0));
						} else {
							stpg.addChoice(src, distr);
						}
					} else if (modelType == ModelType.SMG) {
						if (distinguishActions) {
							smg.addActionLabelledChoice(src, distr, engine.getTransitionAction(i, 0));
						} else {
							smg.addChoice(src, distr);
						}
					}
				}
			}
//			if (mdp.getNumChoices(src) != engine.getNumChoices()) System.out.println("Already different during construction");
			// Print some progress info occasionally
			progress.updateIfReady(src + 1);
		}

		// Finish progress display
		progress.update(src + 1);
		progress.end(" states");

		// Reachability complete
		mainLog.print("Reachable states exploration" + (justReach ? "" : " and model construction"));
		mainLog.println(" done in " + ((System.currentTimeMillis() - timer) / 1000.0) + " secs.");
		//mainLog.println(states);

		// Find/fix deadlocks (if required)
		if (!justReach && findDeadlocks) {
			modelSimple.findDeadlocks(fixDeadlocks);
		}

		boolean sort = true;
		int permut[] = null;

		if (sort) {
			// Sort states and convert set to list
			mainLog.println("Sorting reachable states list...");
			permut = states.buildSortingPermutation();
			statesList = states.toPermutedArrayList(permut);
			//mainLog.println(permut);
		} else {
			statesList = states.toArrayList();
		}
		states.clear();
		states = null;
		//mainLog.println(permut);
		//mainLog.println(statesList);

		// Construct new explicit-state model (with correct state ordering)
		if (!justReach) {
			switch (modelType) {
			case DTMC:
				model = new DTMCSimple(dtmc, permut);
				break;
			case CTMC:
				model = new CTMCSimple(ctmc, permut);
				break;
			case MDP:
				//				buildSparse = false;
				if (buildSparse) {
					model = new MDPSparse(mdp, true, permut);
				} else {
					model = new MDPSimple(mdp, permut);
				}
				break;
			case STPG:
				model = new STPGExplicit(stpg, permut);
				break;
			case SMG:
				HashMap<String, Integer> players = new HashMap<String, Integer>();
				for (i = 0; i < modulesFile.getNumPlayers(); i++)
					players.put(modulesFile.getPlayer(i).getName(), i + 1);
				model = new SMG(smg, permut, players);
				break;
			}
			model.setStatesList(statesList);
			model.setConstantValues(new Values(modulesFile.getConstantValues()));
			// System.out.println(model);
			//mainLog.println("Model: " + model);
		}

		// Discard permutation
		permut = null;

		return model;
	}
	
	private int getMaxModelExplorationDepth() throws PrismLangException, PrismException{
		int maxDepth = Integer.MAX_VALUE;
		if(property != null) {
			if(property.getExpression() instanceof ExpressionPATL) {
				ExpressionTemporal et = null;
				int type = ((ExpressionPATL)property.getExpression()).getExpressionType();
				switch(type) {
					case ExpressionPATL.PRB:
						et = (ExpressionTemporal)((ExpressionPATL)property.getExpression()).getExpressionProb().getExpression();
					break;
					case ExpressionPATL.REW:
						et = (ExpressionTemporal)((ExpressionPATL)property.getExpression()).getExpressionRew().getExpression();
				}
				if(et.isSimplePathFormula() && et.hasBounds()) {
					maxDepth = et.getUpperBound().evaluateInt(this.propertiesFile.getConstantValues());
					if (et.upperBoundIsStrict())
						maxDepth--;
					if (maxDepth < 0) {
						String bound = et.upperBoundIsStrict() ? "<" + (maxDepth + 1) : "<=" + maxDepth;
						throw new PrismException("Invalid bound " + bound + " in bounded until formula");
					}
				}
			}
		}
		return Integer.MAX_VALUE;
	}

	/**
	 * For game models, determine the player who owns the {@code i}th choice in
	 * the state currently being explored by the simulator. Returns the index
	 * (starting from 1) of the player.
	 */
	private int determinePlayerForChoice(ModulesFile modulesFile, ModelType modelType, int i) throws PrismException
	{
		int modAct, player;

		modAct = engine.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(engine.getTransitionModuleOrAction(i, 0));
			if (player == -1) {

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

		return player;
	}

	/**
	 * Test method.
	 */
	public static void main(String[] args)
	{
		try {
			// Simple example: parse a PRISM file from a file, construct the model and export to a .tra file
			PrismLog mainLog = new PrismPrintStreamLog(System.out);
			Prism prism = new Prism(mainLog, mainLog);
			ModulesFile modulesFile = prism.parseModelFile(new File(args[0]));
			UndefinedConstants undefinedConstants = new UndefinedConstants(modulesFile, null);
			if (args.length > 2)
				undefinedConstants.defineUsingConstSwitch(args[2]);
			modulesFile.setUndefinedConstants(undefinedConstants.getMFConstantValues());
			ConstructModel constructModel = new ConstructModel(prism, prism.getSimulator(), null, null);
			Model model = constructModel.constructModel(modulesFile);
			model.exportToPrismExplicitTra(args[1]);
		} catch (FileNotFoundException e) {
			System.out.println("Error: " + e.getMessage());
			System.exit(1);
		} catch (PrismException e) {
			System.out.println("Error: " + e.getMessage());
			System.exit(1);
		}
	}
}
