//==============================================================================
//	
//	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 java.util.*;

import heuristics.VarInfo;
import heuristics.search.StateValue;
import jdd.JDD;
import jdd.JDDNode;
import jdd.JDDVars;
import parser.State;
import parser.ast.ModulesFile;
import prism.PrismException;
import prism.PrismLog;
import weka.core.FastVector;

public class DecisionTreeBDD extends DecisionTree {

	private JDDNode strategy;
	private VarInfo varInfo;
	private int numModules;
	private int numActions;
	private JDDVars vars;
	private int stateVars;
	private int numAllVars;
	private FastVector allActions;
	private FastVector modules;
	private Map<State,Map<String,Map<String,JDDNode>>> storedJDDs;
	
	public DecisionTreeBDD(ModulesFile mf) throws PrismException {
		varInfo = new VarInfo(mf);
		
		numModules = mf.getNumModules();
		modules = new FastVector(mf.getNumModules()+1);
		for (String m: mf.getModuleNames()) modules.addElement(m);
		modules.addElement("synchronised");
		allActions = moduleActions(mf);
		this.numActions = allActions.size();
		vars = varInfo.getJDDVars();
		stateVars = vars.n();
		numAllVars = stateVars;
		for (int i = 0; i < Math.log(numModules); i++) {
			vars.addVar(JDD.Var(numAllVars++));
		}
		for (int i = 0; i < Math.log(numActions); i++) {
			vars.addVar(JDD.Var(numAllVars++));
		}
		strategy = JDD.ZERO;
		storedJDDs = new HashMap<State,Map<String,Map<String,JDDNode>>>();
	}
	
	@Override
	public void dumpTree(PrismLog log)
	{
		JDD.Ref(strategy);
		int num = numNodes(strategy);
		log.println("DT Strat as BDD with " + num + " nodes");
	}

	@Override
	public void buildTree()
	{
		//Tree is already built during learning
	}
	
	private int numNodes(JDDNode node) {
		return allNodes(node).size();
	}
	
	
	private Set<JDDNode> allNodes(JDDNode node) {
		Set<JDDNode> out = new HashSet<JDDNode>();
		out.add(node);
		if (!node.isConstant()) {
			out.addAll(allNodes(node.getElse()));
			out.addAll(allNodes(node.getThen()));
		}
		return out;
	}
	
	private JDDNode stateModuleAction2JDD(State s, String module, String action) {
		JDDNode dd = varInfo.state2JDD(s);
		int currentVar = stateVars;
		int tmpModule = modules.indexOf(module);
		int actionIndex = allActions.indexOf(action);
		int tmpAction = actionIndex;
		for (int i = 0; i < Math.log(numModules); i++) {
			JDDNode tmpVar = vars.getVar(currentVar++);
			JDD.Ref(tmpVar);
			if (tmpModule % 2 == 0) {
				JDD.Ref(tmpVar);
				dd = JDD.And(dd, JDD.Not(tmpVar));
			}
			else {
				dd = JDD.And(dd, tmpVar);
			}
			tmpModule = tmpModule / 2;
		}
		for (int i = 0; i < Math.log(numActions); i++) {
			JDDNode tmpVar = vars.getVar(currentVar++);
			JDD.Ref(tmpVar);
			if (i == 0) {
				if (actionIndex == StateValue.ACTION_UNDEF) {
					dd = JDD.And(dd, tmpVar);
				}
				else {
					JDD.Ref(tmpVar);
					dd = JDD.And(dd, JDD.Not(tmpVar));
				}
			}
			else {
				if (tmpAction % 2 == 0) {
					JDD.Ref(tmpVar);
					dd = JDD.And(dd, JDD.Not(tmpVar));
				}
				else {
					dd = JDD.And(dd, tmpVar);
				}
				tmpAction = tmpAction / 2;
			}
		}
		return dd;
	}
	
	private void learn(State s, String module, String action, boolean good) {
		JDDNode dd = stateModuleAction2JDD(s,module,action);
		JDDNode tmp;
		if (good) {
			JDD.Ref(dd);
			JDD.Ref(JDD.ONE);
			tmp = JDD.ITE(dd, JDD.ONE, strategy);
		}
		else {
			JDD.Ref(dd);
			JDD.Ref(JDD.ZERO);
			tmp = JDD.ITE(dd, JDD.ZERO, strategy);
		}
		strategy = tmp;
	}
	
	@Override
	public void learnGood(State s, String module, String action)
	{
		learn(s,module,action,true);
	}

	@Override
	public void learnBad(State s, String module, String action)
	{
		learn(s,module,action,false);
	}

	@Override
	public boolean classify(State s, String module, String action)
	{
		JDDNode dd = stateModuleAction2JDD(s,module,action);
		JDD.Ref(strategy);
		double value = JDD.Restrict(strategy, dd).getValue();
		return value > 0.0;
	}

	@Override
	public double chanceGood(State s, String module, String action)
	{
		if (classify(s,module,action)) return 1.0;
		else return 0.0;
	}

	@Override
	public void printInfo(PrismLog log)
	{
		log.println("BDD with " + getSize() + " states");
	}

	@Override
	public int getSize()
	{
		return numNodes(strategy);
	}

//	@Override
//	public double getConfidence()
//	{
//		return 1.0;
//	}

	@Override
	public void setOptions(String[] wekaOpts)
	{
	}

	@Override
	public String getTypeString()
	{
		return "BDD";
	}

}
