//==============================================================================
//	
//	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 heuristics.CollapseModelExplorer;
import heuristics.ECContainerMap;

import java.io.BufferedWriter;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.*;

import explicit.ConstructModel;
import explicit.DTMCExplicit;
import explicit.DTMCSimple;
import explicit.Model;
import explicit.ModelExplicit;
import explicit.NondetModel;
import parser.State;
import parser.ast.Command;
import parser.ast.ModulesFile;
import parser.ast.PropertiesFile;
import parser.ast.Property;
import prism.PrismComponent;
import prism.PrismException;
import prism.PrismLangException;
import prism.PrismLog;
import simulator.PrismModelExplorer;
import simulator.SimulatorEngine;
import simulator.TransitionList;
import strat.MDStrategy;
import strat.MDStrategyArray;
import strat.MemorylessDeterministicStrategy;
import strat.Strategy;
import strat.Strategy.Choice;
import weka.core.FastVector;

/**
 * DecisionTree is an abstract container for representing strategies.
 * The implementations are BDD, PRISM Strat and Weka, which is the only actual decision tree.
 * The abstraction is used to work with those types of strategies in a uniform way during training and evaluation.
 */
public abstract class DecisionTree
{
	/**
	 * Extracts module names from a ModulesFile and stores them in a FastVector 
	 */
	protected FastVector moduleActions(ModulesFile mf) throws PrismLangException {
		FastVector actions = new FastVector();
		Vector<String> labels = mf.getSynchs();
		for (String l: labels) actions.addElement(l.toString());
		for (int i = 0; i < mf.getNumModules(); i++) {
			for (Command c: mf.getModule(i).getCommands()) {
				if (c.getSynch().equals("")) actions.addElement(c.evaluatePartially(mf.getConstantValues(),null).simplify().toString());
			}
		}
		return actions;
	}
	
	/**
	 * Returns the minimum of good and bad instances learnt so far.
	 */
	public int getMinGoodBad() {
		return 0;
	}
	
	/**
	 * Outputs the tree to the specified PrismLog
	 */
	public abstract void dumpTree(PrismLog log);
	
	/**
	 * Prints information about the tree to the specified PrismLog
	 */
	public abstract void printInfo(PrismLog log);
	
	/**
	 * Returns the size of the tree
	 */
	public abstract int getSize();
	
	/**
	 * Builds the tree from learnt data for weka tree.
	 */
	public abstract void buildTree();
	
	/**
	 * Builds the tree from learnt data for weka tree.
	 */
	public abstract void learnGood(State s, String module, String action);
	
	public abstract void learnBad(State s, String module, String action);
	
	/**
	 * Return the type of the tree as a String "BDD","PRISM Strat" or "J48"
	 */
	public abstract String getTypeString();
	
	/**
	 * Trains a memoryless deterministic strategy obtained from PRISM into a decision tree
	 * by iterating over the states of the tree and learning the best action as good, while the rest is learned as bad
	 */
	public void trainStrategy(MDStrategy strat, Model model, ModulesFile mf, SimulatorEngine engine) throws PrismException {
		CollapseModelExplorer me = new CollapseModelExplorer(engine,mf,new ECContainerMap());
		List<State> stateList = model.getStatesList();
		State currentState = null;
		String currentModule = null;
		String currentAction = null;
		int currentIndex = 0;
		Set<State> nonDetStates = new HashSet<State>();
		int nondettotal = 0;
		for (int i = 0; i < stateList.size(); i++) {
			currentState = stateList.get(i);
			currentIndex = strat.getChoiceIndex(i);
			me.queryState(currentState);
			if (me.getNumChoices() > 1) {
				if (!nonDetStates.contains(currentState)) nondettotal += me.getNumChoices();
				nonDetStates.add(currentState);
			}
			for (int j = 0; j < me.getNumChoices(); j++) {
				currentAction = me.getActionName(j);
				currentModule = me.getModuleName(j);
				if (j != currentIndex) learnBad(currentState,currentModule,currentAction);
				else learnGood(currentState,currentModule,currentAction);
			}
		}
	}
	
	/**
	 * Constructs a deterministic markov chain from this decision tree and a corresponding MDP.
	 * This is done by resolving non-determinism according to the strategy.
	 */
	public DTMCExplicit constructDTMC(ModulesFile mf, SimulatorEngine engine, explicit.NondetModel model) throws PrismException {
		CollapseModelExplorer me = new CollapseModelExplorer(engine,mf, new ECContainerMap());
		List<State> states = model.getStatesList();
		DTMCSimple dtmc = new DTMCSimple(states.size());
		dtmc.copyFrom((ModelExplicit) model);
		for (int i = 0; i < states.size(); i++) {
			dtmc.clearState(i);
			ArrayList<Integer> goodActions = new ArrayList<Integer>();
			State s = states.get(i);
			me.queryState(s);
			double bestOfTheRestValue = 0.0;
			int bestOfTheRest = -1;
			
			for (int j = 0; j < me.getNumChoices(); j++) {
				simulator.Choice c = me.getChoice(j);
				String module = "";
				if (c.getModuleOrActionIndex() > 0) module = "synchronised";
				else module = c.getModuleOrAction();
				String name = c.getCommandString();
				boolean classif = false;
				classif = classify(s,module,name);
				if (classif) {
					goodActions.add(j);
				}
				else {
					double chanceGood = 0.0;
					chanceGood = chanceGood(s,module,name);
					if (chanceGood > bestOfTheRestValue) {
						bestOfTheRest = j;
						bestOfTheRestValue = chanceGood;
					}
				}
			}
			
			int act = 0;
			if (goodActions.isEmpty()) {
				if (bestOfTheRest >= 0) {
					for (int j = 0; j < me.getNumTransitions(bestOfTheRest); j++) {
						double prob = me.getTransitionProbability(bestOfTheRest, j);
						State target = me.computeTransitionTarget(bestOfTheRest, j);
						int targetIndex = states.indexOf(target);
						dtmc.addToProbability(i, targetIndex, prob);
					}
				}
				else {
					for (int k = 0; k < me.getNumChoices(); k++) {
						for (int j = 0; j < me.getNumTransitions(k); j++) {
							double prob = me.getTransitionProbability(k, j);
							State target = me.computeTransitionTarget(k, j);
							int targetIndex = states.indexOf(target);
							dtmc.addToProbability(i, targetIndex, prob);
						}
					}
					if (me.getNumChoices() == 0) {
						dtmc.addToProbability(i, i, 1.0);
					}
				}
			}
			else {
				for (Integer good: goodActions) {
					for (int j = 0; j < me.getNumTransitions(good); j++) {
						double prob = me.getTransitionProbability(good, j);
						State target = me.computeTransitionTarget(good, j);
						int targetIndex = states.indexOf(target);
						dtmc.addToProbability(i, targetIndex, prob);
					}
				}
			}
		}
		return dtmc;
	}
	
	/**
	 * Constructs the two dimentional integer array that can be used to create a memoryless deterministic strategy.
	 * Iterates over the state space and stores the action index of the strategy action.
	 */
	public int[][] toStrategy(ModulesFile mf, explicit.NondetModel model, Map<State,Integer> succTrace) throws PrismException {
		List<State> states = model.getStatesList();
		int[][] strat = new int[states.size()][];
		for (int i = 0; i < model.getNumStates(); i++) {
			ArrayList<Integer> goodActions = new ArrayList<Integer>();
			State s = states.get(i);
			double bestOfTheRestValue = 0.0;
			int bestOfTheRest = -1;
			for (int j = 0; j < model.getNumChoices(i); j++) {
				simulator.Choice c = (simulator.Choice) model.getAction(i, j);
				if (c != null) {
					String module = "";
					if (c.getModuleOrActionIndex() > 0) module = "synchronised";
					else module = c.getModuleOrAction();
					String name = c.getCommandString();
					boolean classif = false;
					classif = classify(s,module,name);
					if (classif) {
						goodActions.add(j);
					}
					else {
						double chanceGood = 0.0;
						chanceGood = chanceGood(s,module,name);
						if (chanceGood > bestOfTheRestValue) {
							bestOfTheRest = j;
							bestOfTheRestValue = chanceGood;
						}
					}
				}
			}
			int[] takeActions;
			if (goodActions.isEmpty()) {
				if (bestOfTheRest > 0) {
					takeActions = new int[1];
					takeActions[0] = bestOfTheRest;
				}
				else {
					takeActions = new int[1];
					takeActions[0] = -2;
				}
			}
			else {
				takeActions = new int[goodActions.size()];
				for (int j = 0; j < goodActions.size(); j++) {
					takeActions[j] = goodActions.get(j);
				}
			}
			for (int j = 0; j < takeActions.length; j++) {
				if (takeActions[j] > 10) System.out.println("Taking action: " + takeActions[j]);
			}
			strat[i] = takeActions;
		}
		return strat;	
	}
	
	/**
	 * Classifies a state/module/action triple as good (true) or bad (false).
	 * For BDD and PRISM strategy only one action can be good, for weka decision tree it can be more.
	 */
	public abstract boolean classify(State s, String module, String action);

	/**
	 * Returns the probability of an state/module/action triple being good.
	 * This is the classification score (1.0 ~ true, 0.0 ~ false) for BDD and PRISM strategy.
	 * For weka decision tree the classification score, which is a value in [0,1] is returned.
	 */
	public abstract double chanceGood(State s, String module, String action);

	/**
	 * Set decision tree options like M and C.
	 */
	public abstract void setOptions(String[] wekaOpts) throws Exception;
}
