//==============================================================================
//	
//	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.search.Heuristic;
import heuristics.update.StateUpdate;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.io.PrintWriter;
import java.util.ArrayList;
import java.util.Locale;

import explicit.MDP;
import explicit.Model;
import explicit.NondetModel;
import explicit.PrismExplicit;
import parser.State;
import parser.ast.Expression;
import parser.ast.PropertiesFile;
import prism.PrismException;
import prism.PrismLog;
import prism.Result;
import strat.MDStrategy;

/**
 * Experiment class is done to perform repeated experiments, find best M and log data
 */
public class Experiment {

	private static double initConfidence = 0.0001;
	private static double maxConfidence = 0.5;
	
	public static double totalEvaluationTime = 0.0;
	
	
	
	/**
	 * Performs a binary search for best M parameter.
	 * For every M parameter, it calls the run method to perform evaluation and output the results
	 */
	public static DecisionTree search(PrismLog mainLog, Result heuristicRes, CollapseModelExplorer pme, StateUpdate stateUpdate, String dtFile, DecisionTree tree, int trains, int tests,Integer searchAround, double delta, int compareScheme, int learningScheme, double conf, NondetModel model, PrismExplicit modelChecker, PropertiesFile propertiesFile, Expression expression, String optString, MDStrategy strat, Model stratModel, boolean onlyLernTarget, boolean countOnlyOnce, boolean empiricEval, boolean bdd) throws Exception {
		PrintWriter writer=writeHeader(dtFile,mainLog);
		DeltaStar ds = new DeltaStar(pme.getSimEngine(),pme,stateUpdate,compareScheme,learningScheme, onlyLernTarget, countOnlyOnce);
			for (int j = 0; j < trains; j++) {
			if (strat != null && stratModel != null) {
				ds.deltaStarSimulation(pme.getDefaultInitialState(), delta,strat,stratModel,false,false);
			}
			else if (bdd) {
				ds.deltaStarSimulation(pme.getDefaultInitialState(), delta,null,null,true,false);
			}
			else {
				ds.deltaStarSimulation(pme.getDefaultInitialState(), delta);
			}
		}
		if (tree == null) tree = new DecisionTreeBDD(pme.getModulesFile());
		ds.learn(tree);
		double error  = 100;
		int lowerM = 1;
		int upperM = tree.getMinGoodBad();
		
		if (searchAround != null) {
			upperM = searchAround*2;
			lowerM = searchAround/2;
		}
		
		int initUp = upperM;
		int currentM = 0;
		int currentTreeSize = Integer.MAX_VALUE;
		while (error > 0.01 && upperM != lowerM && !(error < 0.01 && currentTreeSize == 1)) {
			currentM = (upperM + lowerM)/2;
			String[] wekaOpts = {"-C",conf+"","-M",currentM+""};
			String optionString = optString;
			error = doDTRun(mainLog, pme,stateUpdate,tree,wekaOpts,optionString,delta,conf,currentM,tests,trains, heuristicRes,writer, model, modelChecker, propertiesFile, expression, ds, learningScheme,empiricEval);
			if (error > 0.01) upperM = currentM;
			else lowerM = currentM;
			currentTreeSize = tree.getSize();
		}
		return tree;
	}
	
	/**
	 * Performs experiments on sampled M, C and delta parameters or on single input values
	 */
	public static DecisionTree run(PrismLog mainLog, Result heuristicRes, CollapseModelExplorer pme, StateUpdate stateUpdate, String dtFile, DecisionTree tree, int trains, int tests, double delta, boolean deltaIt, int min, boolean minIt, double conf, boolean confIt, int compareScheme, int learningScheme, NondetModel model, PrismExplicit modelChecker, PropertiesFile propertiesFile, Expression prop, String optString, MDStrategy strat, Model stratModel, boolean onlyLernTarget, boolean countOnlyOnce, boolean empiricEval, boolean bdd) throws Exception {
			PrintWriter writer=writeHeader(dtFile,mainLog);
			
			String optionString = "";
			pme.queryState(pme.getDefaultInitialState());
			int confidenceDivisor = 1;
			if (conf > 1.0) confidenceDivisor = (int)(conf-1.0);
			
			int minDivisor = 1;
			if (min > 1) minDivisor = (int)(min-1);
			
			int deltaIterations = (int) delta;
			if (!deltaIt) deltaIterations = 1;
			
			for (int i = 0; i < deltaIterations; i++) {
				double currentDelta = 0.0;
				if (!deltaIt) currentDelta = delta;
				if (i != 0) delta = Math.pow(10, -i);
				DeltaStar ds = new DeltaStar(pme.getSimEngine(),pme,stateUpdate,compareScheme,learningScheme, onlyLernTarget, countOnlyOnce);
				long trainBefore = System.currentTimeMillis();
				for (int j = 0; j < trains; j++) {
					if (strat != null && stratModel != null) {
						ds.deltaStarSimulation(pme.getDefaultInitialState(), delta,strat,stratModel,false,false);
					}
					else if (bdd) {
						ds.deltaStarSimulation(pme.getDefaultInitialState(), delta,null,null,false,false);
					}
					else {
						ds.deltaStarSimulation(pme.getDefaultInitialState(), delta);
					}
				}
				
				if (tree == null) tree = new DecisionTreeWeka(pme.getModulesFile(),1);
				ds.learn(tree);

				long trainAfter = System.currentTimeMillis();
				mainLog.println("\nTraining " + trains + " times with delta: " + delta + " took: " + (double) (trainAfter-trainBefore)/1000 + " seconds\n");

				int instances = tree.getMinGoodBad() + 1;
				mainLog.println("Decision tree learned instances: " + instances);

				int confIterations = (int) conf;
				if (!confIt) confIterations = 1;
				
				/***
				 * Confidence Factor and min instances iterations
				 */
				for (int k = 0; k < confIterations; k++) {
					
					int currentMin = 1;
					double currentConf = initConfidence + k*(maxConfidence-initConfidence) / confidenceDivisor;
					if (!confIt) currentConf = conf;
					
					int minIterations = (int) min;
					if (!minIt) minIterations = 1;
					
					double b = Math.log(instances) / minIterations;
					double factor = (instances - 1) / (Math.exp(b*(min-1)));
					
					for (int l = 0; l < minIterations; l++) {
						/**
						 * exponential growth of M
						 * minInstances = (int) Math.round(factor * Math.exp(l*b) - factor + 1);
						 * 
						 * linear growth of M
						 * minInstances = 1 + i*(instances-1)/min;
						 */
						currentMin = (int) Math.round(factor * Math.exp(l*b) - factor + 1);
						if (!minIt) currentMin = min;
						mainLog.println("M: " + currentMin);
						String[] wekaOpts = {"-C",currentConf+"","-M",currentMin+""};
						optionString = optString;
						doDTRun(mainLog, pme,stateUpdate,tree,wekaOpts,optionString,currentDelta,currentConf,currentMin,tests,trains, heuristicRes,writer,model,modelChecker,propertiesFile, prop, ds, learningScheme,empiricEval);
					}
				}
			}
			if (dtFile != null) writer.close();
		return tree;
	}
	
	/**
	 * Performs decision tree evaluation and outputs the results the a csv file.
	 * Evaluation can be based on simulations or model checking the product of the strategy and the constructed model. 
	 */
	public static double doDTRun(PrismLog mainLog, CollapseModelExplorer pme, StateUpdate stateUpdate, DecisionTree tree, String[] wekaOpts, String optionString, double delta, double conf, int min, int tests, int trains, Result heuristicRes, PrintWriter writer, NondetModel model, PrismExplicit modelChecker, PropertiesFile propertiesFile, Expression prop, DeltaStar ds, int learningScheme, boolean empiricEval) throws Exception {
		long start = System.currentTimeMillis();
		String[] wekaOpts_overwrite = {"-M",""+min,"-C",""+conf,"-B"};
		tree.setOptions(wekaOpts_overwrite);
		tree.buildTree();
		double referenceResult = (Double)heuristicRes.getResult();
		double error = 0.05;
		double leftSuccess = 0.0;
		double rightSuccess = Double.MAX_VALUE;
		double leftError = 0.0;
		double rightError = 1.0;
		double success = 0.0;
		int alltests = 0;
		PlayDT dtSim = new PlayDT(pme.getSimEngine(),pme, stateUpdate, tree);
		int leftGoods = 0;
		int rightGoods = 0;
		Result exactResult = new Result(0);
		if (!(referenceResult > 0.0)) referenceResult = 1.0;
		if (empiricEval) {
			//errors more than 10% should not be investigated further
			//no more than 4 000 000 simulations
			//left and right error agree by 1%
			while (!(error > 0.1 || alltests > 4*1e6 || Math.abs(leftError - rightError) < 0.01)) {
				for (int j = 0; j < tests; j++) {
					if (dtSim.doSimulation()) leftGoods++;
					if (dtSim.doSimulation()) rightGoods++;
				}
				alltests += tests;
				leftSuccess = ((double)leftGoods)/alltests;
				rightSuccess = ((double)rightGoods)/alltests;
				leftError = (1 - leftSuccess/referenceResult);
				rightError = (1 - rightSuccess/referenceResult);
				success = (leftSuccess + rightSuccess) / 2;
				error = Math.abs((1-success/referenceResult));				
			}
			
		}
		else {
			int[][] testStrat = tree.toStrategy(pme.getModulesFile(), (MDP) model, new java.util.HashMap<State,Integer>());
			explicit.DTMCFromMDPAndDTStrategy resultModel = new explicit.DTMCFromMDPAndDTStrategy((MDP) model,testStrat);
			exactResult = modelChecker.modelCheck(resultModel, pme.getModulesFile(), propertiesFile, prop);
			error = (1-(Double)exactResult.getResult()/referenceResult);
			success = (Double) exactResult.getResult();
		}
		tree.dumpTree(mainLog);
		int size = tree.getSize();
		long end = System.currentTimeMillis();
		double time = (double)(end-start) /1000;
		totalEvaluationTime += time;
		String row = size + ", " +  String.format(Locale.US, "%.5f", success) + " " + String.format(Locale.US, "%.5f", referenceResult) + ", " + String.format(Locale.US, "%.5f", error)  + ", " + conf + ", " + min + ", "+ optionString  + ", " + String.format(Locale.US, "%.3f", time) + ", " + tree.getTypeString() + ", " + learningScheme + ", " + trains + ", " + alltests*2 + " , " + String.format(Locale.US, "%.3f", ds.getAvgGoodPerState()) + " , " + String.format(Locale.US, "%.3f", ds.getAvgActionPerState()) + ", " + stateUpdate.getNumExplored();
		if (writer != null){
			writer.println(row);
			writer.flush();
		}
		else {
			mainLog.println(row);
		}
		return error;
	}
	
	/**
	 * Writes the header of the csv file
	 */
	private static PrintWriter writeHeader(String dtFile, PrismLog mainLog) throws IOException {
		PrintWriter writer = null;
		String head = "treeSize, value, referenceValue, relError, C, M, Options, evalTime(s), stratType, learnScheme, trains, verifies, gps, aps, explored";
		if (dtFile != null){
			writer = new PrintWriter(new BufferedWriter(new FileWriter(dtFile+".csv", true)));
			BufferedReader br = new BufferedReader(new FileReader(dtFile+".csv"));     
			if (br.readLine() == null) {
				writer.println(head);
				writer.flush();
			}
		}
		else {
			mainLog.println(head);
		}
		return writer;
	}
}
