//==============================================================================
//	
//	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.HashMap;
import java.util.Map;

import parser.State;
import parser.ast.ModulesFile;
import prism.PrismLangException;
import prism.PrismLog;
import weka.classifiers.Classifier;
import weka.classifiers.meta.CVParameterSelection;
import weka.classifiers.trees.*;
import weka.core.Attribute;
import weka.core.FastVector;
import weka.core.Instance;
import weka.core.Instances;
import weka.core.Utils;

public class DecisionTreeWeka extends DecisionTree
{
	
	private Classifier tree;
	private Instances trainingSet;
	
	private int numAttributes;
	private FastVector attributes;
	
	private ModulesFile mf;
	private int type;
	
	private int goodInstances;
	private int badInstances;
	
	private FastVector allActions;
	
	private Map<State,Map<String,Map<String,Instance>>> storedInstances;
	
	/**
	 * Creates an Weka DecisionTree
	 * Attributes are:
	 * 	every module variable as numeric attribute
	 * 	module of the state as numeric attribute
	 * 	played action as a nominal (set of strings) attribute
	 * 
	 * @param mf ModulesFile trained
	 * @param treeType type of tree classifier, see DecisionTree class
	 * @param allActions List of all action strings, for nominal classification
	 */
	public DecisionTreeWeka(ModulesFile mf, int treeType) {
		numAttributes = mf.getNumVars()+3; //1 attr per variable + module id + action id + classification
		attributes = new FastVector(numAttributes);
		for (int i = 0; i < mf.getNumVars(); i++) {
			attributes.addElement(new Attribute(mf.getVarName(i)));
		}
		FastVector modules = new FastVector(mf.getNumModules()+1);
		for (String m: mf.getModuleNames()) modules.addElement(m);
		modules.addElement("synchronised");
		attributes.addElement(new Attribute("module_of_state",modules));
		
		try {
			allActions=moduleActions(mf);
		} catch (PrismLangException e) {
			e.printStackTrace();
		}
		attributes.addElement(new Attribute("played_action",allActions));
		FastVector fvClassVal = new FastVector(2);
		fvClassVal.addElement("good");
		fvClassVal.addElement("bad");
		Attribute classAttribute = new Attribute("classification", fvClassVal);
		attributes.addElement(classAttribute);
		trainingSet = new Instances("Rel", attributes, 100);
		trainingSet.setClassIndex(numAttributes-1);
		this.type = treeType;
		tree = WekaTree.createTree(treeType);
		goodInstances = 0;
		badInstances = 0;
		this.mf = mf;
		storedInstances = new HashMap<State,Map<String,Map<String,Instance>>>();
	}

	public int getNumInstances() {
		return trainingSet.numInstances();
	}
	
	public void setOptions(String[] opts) throws Exception {
		if (type == 1 || type == 2) tree.setOptions(opts);
	}
	
	@Override
	public void dumpTree(PrismLog log)
	{
		log.println(tree);
	}

	@Override
	public void buildTree()
	{
		try {
			tree.buildClassifier(trainingSet);
		} catch (Exception e) {
			e.printStackTrace();
			System.exit(-1);
		}
	}
	
	@Override
	public void printInfo(PrismLog log) {
		if (tree instanceof J48) {
			J48 jtree = (J48)tree;
			log.println("Confidence factor: " + jtree.getConfidenceFactor());
			log.println("Size: " + jtree.measureTreeSize());
			log.println("Rules: " + jtree.measureNumRules());
		}
	}
	
	public void optimiseConf() throws Exception {
	      // setup classifier
	      CVParameterSelection ps = new CVParameterSelection();
	      ps.setClassifier(new J48());
	      ps.setNumFolds(5);  // using 5-fold CV
	      ps.addCVParameter("C 0.1 0.5 5");

	      // build and output best options
	      ps.buildClassifier(trainingSet);
	      System.out.println(Utils.joinOptions(ps.getBestClassifierOptions()));
	}
	
	public Instance stateAction2Instance(State s, String module, String action, boolean good, boolean training) {
		// Instance storing disabled for the moment
//		Map<String,Map<String,Instance>> stateMap = storedInstances.get(s);
//		Map<String,Instance> moduleMap;
//		if (stateMap != null) {
//			moduleMap = stateMap.get(module);
//			if (moduleMap != null) {
//				Instance actionOutput = moduleMap.get(action);
//				if (actionOutput != null) {
//					return actionOutput;
//				}
//			}
//			else {
//				moduleMap = new HashMap<String,Instance>();
//			}
//		}
//		else {
//			moduleMap = new HashMap<String,Instance>();
//			stateMap = new HashMap<String,Map<String,Instance>>();
//		}
		int actualNumAttr = numAttributes - 1; //attributes minus +/- classification
		if (training) actualNumAttr++; //for training; classification is an attribute
		Instance instance = new Instance(actualNumAttr);
		instance.setDataset(trainingSet);
		for (int i = 0; i < mf.getNumVars(); i++) {
			double val = 0.0;
			if (s.varValues[i] instanceof Integer) {
				val = ((Integer)s.varValues[i]).doubleValue();
			}
			else if (s.varValues[i] instanceof Boolean) {
				if ((Boolean)s.varValues[i]) val = 1.0;
				else val = 0.0;
			}
			else {
				System.out.println(s.varValues[i] + " is neither Integer nor Boolean");
			}
			instance.setValue(i, val);
		}
		instance.setValue(numAttributes-3, module);
		try {
			instance.setValue(numAttributes-2, action);
		}
		catch (Exception e) {
			System.out.println("Action: " + action + " is not defined");
			System.out.println("Defined actions: " + allActions);
			for (int i = 0; i < allActions.size(); i++) System.out.println(allActions.elementAt(i));
			System.exit(-1);
		}
		if (training) {
			String classif = "bad";
			if (good) classif = "good";
			instance.setValue((Attribute)attributes.elementAt(numAttributes-1), classif);
		}
//		moduleMap.put(action,instance);
//		stateMap.put(module, moduleMap);
//		storedInstances.put(s,stateMap);
		return instance;
	}
	
	@Override
	public boolean classify(State s, String module, String action)
	{
		Instance testInstance = stateAction2Instance(s,module,action,true,false);
		testInstance.setDataset(trainingSet);
		
		double classif = 1.0;
		try {
			classif = tree.classifyInstance(testInstance);
		} catch (Exception e) {
			e.printStackTrace();
			System.exit(-1);
		}
		return classif < 1.0;
	}
	
	@Override
	public double chanceGood(State s, String module, String action) {
		Instance testInstance = stateAction2Instance(s,module,action,true,false);
		testInstance.setDataset(trainingSet);
		try {
			return tree.distributionForInstance(testInstance)[0];
		} catch (Exception e) {
			e.printStackTrace();
		}
		return 0.0;
	}
	
	@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);
	}
	
	public void learn(State s, String module, String action, boolean good) {
		learn(s,module,action,good,1.0);
	}
	
	public void learn(State s, String module, String action, boolean good, double weight) {
		Instance instance = stateAction2Instance(s,module,action,good,true);
		instance.setWeight(weight);
		trainingSet.add(instance);
		if (good) goodInstances++;
		else badInstances++;
	}

	public int getMinGoodBad() {
		if (badInstances < goodInstances) return badInstances;
		else return goodInstances;
	}
	
	@Override
	public int getSize()
	{
		switch (type) {
			case 1: return ((Double)((J48)tree).measureTreeSize()).intValue();
			case 2: return ((Double)((J48graft)tree).measureTreeSize()).intValue();
			case 3: return ((Double)((ADTree)tree).measureTreeSize()).intValue();
			case 4: return ((Double)((BFTree)tree).measureTreeSize()).intValue();
	//		case 5: return ((Double)((DecisionStump)tree).measureTreeSize()).intValue();
			case 6: return ((Double)((LADTree)tree).measureTreeSize()).intValue();
	//		case 7: tree = return ((Double)((REPTree)tree).get
			case 8: return ((Double)((SimpleCart)tree).measureTreeSize()).intValue();
	//		case 9: return ((RandomForest)tree).getNumTrees();
			default: return -1;
		}
	}

	@Override
	public String getTypeString()
	{
		return WekaTree.treeTypes.get(type);
	}
}
