package ec.gp.semantic.statistics;

import ec.EvolutionState;
import ec.Individual;
import ec.Population;
import ec.Subpopulation;
import ec.app.regression.testcases.FunctionFactory;
import ec.app.regression.testcases.TestCaseFactory;
import ec.gp.GPIndividual;
import ec.gp.GPNode;
import ec.gp.GPTree;
import ec.gp.koza.KozaFitness;
import ec.gp.semantic.EvoState;
import ec.gp.semantic.ISemantics;
import ec.gp.semantic.evaluation.EvaluationMode;
import ec.gp.semantic.evaluation.IChangeableEvaluationMode;
import ec.gp.semantic.func.SimpleNodeBase;
import ec.gp.semantic.library.ILibrary;
import ec.gp.semantic.utils.Pair;
import ec.gp.syntactic.equivalence.IEquivalenceChecker;
import ec.simple.SimpleProblemForm;
import ec.simple.SimpleStatistics;
import ec.util.ParamClassLoadException;
import ec.util.Parameter;
import ec.util.ParameterDatabase;
import experimentdatabase.DataSet;
import experimentdatabase.Database;
import experimentdatabase.Experiment;
import gnu.trove.map.custom_hash.TObjectLongCustomHashMap;
import gnu.trove.strategy.IdentityHashingStrategy;
import library.INamedElement;

import java.util.*;

public class ExtendedStatistics extends SimpleStatistics {

	private static final String OUT = "out";
	private static final String PRINT_INDIVIDUAL = "printIndividual";
	private static final String METHOD = "method";

	private static final String EQUIVALENCE_CHECKER = "equivalenceChecker";
	private static final String ADDITIONAL_STATISTICS = "additionalStatistics";
	private static final String CALCULATE_UNIQUE_SEMANTICS = "calculateUniqueSemantics";
	private static final String CALCULATE_UNIQUE_INDIVIDUALS = "calculateUniqueIndividuals";

	private static final String EVALUATION_LIMIT = "evaluationLimit";

	private static final Pair<Double, Integer> WRONG_TEST_SET_RESULT = new Pair<Double, Integer>(Double.NaN, 0);

	protected Database database;
	protected Experiment experiment;
	protected DataSet initialization;
	protected DataSet iteration;

	protected boolean printIndividual = false;
	protected boolean calculateUniqueSemantics = false;
	protected boolean calculateUniqueIndividuals = true;

	protected long evaluationLimit = -1;

	protected String targetSymbolicRepresentation;
	protected IEquivalenceChecker equivalenceChecker;

	protected List<IAdditionalStatistics> additionalStatistics = new ArrayList<IAdditionalStatistics>();

	protected final Stopwatch totalTime = new Stopwatch();
	protected final Stopwatch totalTimeNoStats = new Stopwatch();
	protected final Stopwatch initializationTime = new Stopwatch();
	protected final Stopwatch breedingTime = new Stopwatch();
	protected final Stopwatch evaluationTime = new Stopwatch();
	private int timersToResume = 1 << 31;

	/**
	 * Gets result database object.
	 *
	 * @return
	 */
	protected Database getDatabase() {
		return this.database;
	}

	/**
	 * Gets object that stores the statistics for entire experiment.
	 *
	 * @return
	 */
	public Experiment getExperiment() {
		return this.experiment;
	}

	/**
	 * Gets object that stores the statistics for current iteration.
	 *
	 * @return
	 */
	public DataSet getIteration() {
		return this.iteration;
	}

	public void SuspendTimers() {
		assert this.timersToResume == (1 << 31);

		int timers = 0;

		if (this.initializationTime.isRunning()) {
			this.initializationTime.stop();
			timers |= 0b00000010;
		}

		if (this.breedingTime.isRunning()) {
			this.breedingTime.stop();
			timers |= 0b00000100;
		}

		if (this.evaluationTime.isRunning()) {
			this.evaluationTime.stop();
			timers |= 0b00001000;
		}

		if (this.totalTimeNoStats.isRunning()) {
			this.totalTimeNoStats.stop();
			timers |= 0b00000001;
		}

		this.timersToResume = timers;
	}

	public void ResumeTimers() {
		assert (this.timersToResume & (1 << 31)) == 0;

		if ((this.timersToResume & 0b00000010) == 0b00000010) {
			this.initializationTime.start();
		}

		if ((this.timersToResume & 0b00000100) == 0b00000100) {
			this.breedingTime.start();
		}

		if ((this.timersToResume & 0b00001000) == 0b00001000) {
			this.evaluationTime.start();
		}

		if ((this.timersToResume & 0b00000001) == 0b00000001) {
			this.totalTimeNoStats.start();
		}

		this.timersToResume = 1 << 31;
	}

	@Override
	public void setup(final EvolutionState state, final Parameter base) {
		super.setup(state, base);

		Locale.setDefault(Locale.ROOT);

		ParameterDatabase db = state.parameters;

		this.printIndividual = db.getBoolean(base.push(PRINT_INDIVIDUAL), null, printIndividual);
		this.calculateUniqueSemantics = db.getBoolean(base.push(CALCULATE_UNIQUE_SEMANTICS), null, calculateUniqueSemantics);
		this.calculateUniqueIndividuals = db.getBoolean(base.push(CALCULATE_UNIQUE_INDIVIDUALS), null, calculateUniqueIndividuals);
		this.evaluationLimit = db.getLongWithDefault(new Parameter(EVALUATION_LIMIT), null, -1);

		TestCaseFactory tcFactory = TestCaseFactory.getFactory(state);
		if (tcFactory instanceof FunctionFactory) {
			this.targetSymbolicRepresentation = ((FunctionFactory) tcFactory).getSymbolicRepresentation();
		}

		try {
			try {
				this.equivalenceChecker = (IEquivalenceChecker) db.getInstanceForParameter(base.push(EQUIVALENCE_CHECKER), null, IEquivalenceChecker.class);
			} catch (ParamClassLoadException ex) {
				// just ignore
			}

			this.database = new Database(db.getFile(base.push(OUT), null));
			this.experiment = this.database.newExperiment();

			// load additional statistics classes
			for (int i = 0; ; ++i) {
				Parameter param = base.push(String.format("%s.%d", ADDITIONAL_STATISTICS, i));
				if (!db.exists(param))
					break;

				IAdditionalStatistics stats = (IAdditionalStatistics) db.getInstanceForParameter(param, null, IAdditionalStatistics.class);
				this.additionalStatistics.add(stats);
				stats.setup(state, param);
			}
		} catch (Exception ex) {
			throw new RuntimeException(ex);
		}

		this.totalTime.start();
	}

	@Override
	public void preInitializationStatistics(EvolutionState state) {
		super.preInitializationStatistics(state);

		this.iteration = this.experiment.newChildDataSet("generations");
		this.iteration.put("generation", 0); // for initial generation

		this.initializationTime.start();
		this.totalTimeNoStats.start();
	}

	@Override
	public void postInitializationStatistics(final EvolutionState state) {
		this.initializationTime.stop();
		this.totalTimeNoStats.stop();
		super.postInitializationStatistics(state);

		ParameterDatabase params = state.parameters;
		String method = params.getString(new Parameter(METHOD), null);
		String problem = ((INamedElement) state.evaluator.p_problem).getName();
		int seed = params.getInt(new Parameter("seed.0"), null) + (int) state.job[0];

		this.experiment.put("in_problem", problem);
		this.experiment.put("in_method", method);
		this.experiment.put("in_popSize", state.population.subpops[0].individuals.length);
		this.experiment.put("in_generations", state.numGenerations);
		this.experiment.put("in_seed", seed);

		ILibrary<?> library = ((EvoState) state).getLibrary();
		if (library != null) {
			this.experiment.put("libSize", library.size());
		}

		this.experiment.put("initTime", this.initializationTime.getTimeElapsed() * 1.0E-9);

		//this.experiment.save();

		state.output.message("Problem:\t" + problem);
		state.output.message("Method:\t" + method);
	}

	@Override
	public void preEvaluationStatistics(EvolutionState state) {
		super.preEvaluationStatistics(state);

		this.totalTimeNoStats.start();
		this.evaluationTime.reset();
		this.evaluationTime.start();
	}

	@Override
	public void postEvaluationStatistics(final EvolutionState state) {
		this.evaluationTime.stop();
		this.totalTimeNoStats.stop();
		this.iteration.put("evalTime", this.evaluationTime.getTimeElapsed() * 1.0E-9);

		super.postEvaluationStatistics(state);

		this.populationStatistics(state);

		int uniqueSemanticsCount = 0;
		int semanticallyUniqueIndCount = 0;
		if (this.calculateUniqueSemantics) {
			uniqueSemanticsCount = this.getUniqueSemanticsCount(state);
			this.iteration.put("uniqueSemantics", uniqueSemanticsCount);
		}
		if (this.calculateUniqueIndividuals) {
			semanticallyUniqueIndCount = this.getSemanticallyUniqueIndividualCount(state);
			this.iteration.put("semanticallyUniqueIndividuals", semanticallyUniqueIndCount);
		}

		Pair<Double, Integer> testSetResults = this.getTestSetResults((GPIndividual) this.getBestSoFar()[0], state);
		if (testSetResults != WRONG_TEST_SET_RESULT) {
			this.iteration.put("testSetFitness", testSetResults.value1);
			this.iteration.put("testSetHits", testSetResults.value2);
		}

		for (IAdditionalStatistics stats : this.additionalStatistics) {
			stats.run(state);
		}
	}

	private void populationStatistics(final EvolutionState state) {
		final Subpopulation pop = state.population.subpops[0];
		final Individual[] individuals = pop.individuals;
		//final SimpleProblemForm problem = (SimpleProblemForm) state.evaluator.p_problem;

		double fitnessAvg = 0.0;
		double fitnessStddev = 0.0;
		double fitnessSumSq = 0.0;
		double treeDepthAvg = 0.0;
		double treeDepthStddev = 0.0;
		double treeDepthSumSq = 0.0;
		double nodeCountAvg = 0.0;
		double nodeCountStddev = 0.0;
		double nodeCountSumSq = 0.0;
		double nodeInMemoryAvg = 0.0;
		//double nodeInMemoryStddev = 0.0;
		double bestFitness = Double.MAX_VALUE;
		long bestNodeCount = 0;
		long bestDepth = 0;
		int bestHits = 0;
		int equivalentTotal = 0;

		double currentFitness;
		long currentNodeCount;
		int currentDepth;

		GPIndividual ind;
		GPNode root;
		KozaFitness fitness;
		Pair<Integer, Long> size;

		// #ifndef NO_SIZE_CACHE
		this.initSizeCache(individuals.length);
		// #endif

		for (int i = 0; i < individuals.length; ++i) {
			ind = (GPIndividual) individuals[i];
			fitness = (KozaFitness) ind.fitness;
			root = ind.trees[0].child;

			assert ind.evaluated;

			// compute fitness statistics
			currentFitness = fitness.standardizedFitness();
			/*
			if (problem instanceof IChangeableEvaluationMode) {
				assert ((IChangeableEvaluationMode) problem).getEvaluationMode() == EvaluationMode.TrainingSet;
			}
			((SimpleNodeBase<?>) root).resetSemanticsRecursive();
			ind.evaluated = false;
			problem.evaluate(state, ind, 0, 0);

			assert currentFitness == fitness.standardizedFitness();
			*/
			fitnessAvg += currentFitness;
			fitnessSumSq += currentFitness * currentFitness;
			if (currentFitness < bestFitness)
				bestFitness = currentFitness;

			if (fitness.hits > bestHits)
				bestHits = fitness.hits;

			size = getSizeFast(root);
			//assert size.value1 == root.depth();
			//assert size.value2 == root.numNodes(GPNode.NODESEARCH_ALL);

			// compute tree depth statistics
			currentDepth = size.value1;//root.depth();//
			treeDepthAvg += currentDepth;
			treeDepthSumSq += currentDepth * currentDepth;

			// compute node count statistics
			currentNodeCount = size.value2;//root.numNodes(GPNode.NODESEARCH_ALL);//
			nodeCountAvg += currentNodeCount;
			nodeCountSumSq += currentNodeCount * currentNodeCount;

			if (currentFitness == bestFitness) {
				bestDepth = currentDepth;
				bestNodeCount = currentNodeCount;
			}

			if (this.equivalenceChecker != null
					&& fitness.standardizedFitness() <= 1E-100 //syntactically equivalent program must be also semantically equivalent
					&& this.equivalenceChecker.areEquivalent(root.makeCTree(true, true, true),
					this.targetSymbolicRepresentation)) {
				++equivalentTotal;
			}
		}

		// #ifndef NO_SIZE_CACHE
		//state.output.message(String.format("Cache size: %d", __getSizePrivate_cache.size()));
		nodeInMemoryAvg = __getSizePrivate_cache.size();
		this.resetSizeCache();
		// #endif

		final double countF = (double) individuals.length;
		final double invCountF = 1.0 / countF;
		fitnessAvg *= invCountF;
		fitnessStddev = Math.sqrt(fitnessSumSq * invCountF - fitnessAvg * fitnessAvg);
		treeDepthAvg *= invCountF;
		treeDepthStddev = Math.sqrt(treeDepthSumSq * invCountF - treeDepthAvg * treeDepthAvg);
		nodeCountAvg *= invCountF;
		nodeCountStddev = Math.sqrt(nodeCountSumSq * invCountF - nodeCountAvg * nodeCountAvg);
		nodeInMemoryAvg *= invCountF;
		//nodeInMemoryStddev = Math.sqrt(nodeInMemoryStddev * invCountF - nodeInMemoryAvg * nodeInMemoryAvg);

		this.iteration.put("bestFitness", bestFitness);
		this.iteration.put("bestSoFarFitness", ((KozaFitness) this.getBestSoFar()[0].fitness).standardizedFitness());
		this.iteration.put("bestHits", bestHits);
		this.iteration.put("bestNodeCount", bestNodeCount);
		this.iteration.put("bestDepth", bestDepth);
		this.iteration.put("fitnessAvg", fitnessAvg);
		this.iteration.put("fitnessStddev", fitnessStddev);
		this.iteration.put("fitnessSumSq", fitnessSumSq);
		this.iteration.put("treeDepthAvg", treeDepthAvg);
		this.iteration.put("treeDepthStddev", treeDepthStddev);
		this.iteration.put("treeDepthSumSq", treeDepthSumSq);
		this.iteration.put("nodeCountAvg", nodeCountAvg);
		this.iteration.put("nodeCountStddev", nodeCountStddev);
		this.iteration.put("nodeCountSumSq", nodeCountSumSq);
		// #ifndef NO_SIZE_CACHE
		this.iteration.put("nodeInMemoryAvg", nodeInMemoryAvg);
		//this.iteration.put("nodeInMemoryStddev", nodeInMemoryStddev);
		// #endif

		if (this.equivalenceChecker != null)
			this.iteration.put("equivalentTotal", equivalentTotal);
	}

	private long __getSizePrivate_maxSize;

	// #ifndef NO_SIZE_CACHE
	//private IdentityHashMap<GPNode, Long> __getSizePrivate_cache;
	private TObjectLongCustomHashMap<GPNode> __getSizePrivate_cache;
	// #endif

	/**
	 * @param subtree
	 * @return Value1 = depth, Value2 = node count
	 */
	private Pair<Integer, Long> getSizeFast(GPNode subtree) {
		__getSizePrivate_maxSize = 0;

		int depth = getSizeInternal(subtree);

		return new Pair<Integer, Long>(depth, __getSizePrivate_maxSize);
	}

	/**
	 * @param subtree
	 * @return Value1 = depth, Value2 = node count
	 */
	public Pair<Integer, Long> getSize(GPNode subtree) {
		// #ifndef NO_SIZE_CACHE
		if (__getSizePrivate_cache == null) {
			this.initSizeCache(3);
		}
		// #endif

		return getSizeFast(subtree);
	}

	private void initSizeCache(int individuals) {
		// #ifndef NO_SIZE_CACHE
		__getSizePrivate_cache = new TObjectLongCustomHashMap<GPNode>(IdentityHashingStrategy.INSTANCE, 0x00000200 * individuals, 0.75f, Long.MAX_VALUE);
		// #endif
	}

	public void resetSizeCache() {
		__getSizePrivate_cache = null;
	}

	private int getSizeInternal(GPNode subtree) {
		// #ifndef NO_SIZE_CACHE
		long cached = __getSizePrivate_cache.get(subtree);
		if (cached != Long.MAX_VALUE) {
			__getSizePrivate_maxSize += cached & 0x00FFFFFFFFFFFFFFL;
			return (int) ((cached & 0xFF00000000000000L) >>> 0x38L);
		}

		long prevSize = __getSizePrivate_maxSize;
		//#endif

		int maxDepth = 0;
		int curDepth;

		for (int ch = 0; ch < subtree.children.length; ++ch) {
			curDepth = this.getSizeInternal(subtree.children[ch]);
			if (curDepth > maxDepth)
				maxDepth = curDepth;
		}

		++maxDepth; //add me
		++__getSizePrivate_maxSize;

		//#ifndef NO_SIZE_CACHE

		assert maxDepth < (1 << 0x08);
		assert (long) (__getSizePrivate_maxSize - prevSize) <= 0x00FFFFFFFFFFFFFFL;

		__getSizePrivate_cache.put(subtree,
				((((long) maxDepth) << 0x38L) | ((long) (__getSizePrivate_maxSize - prevSize) & 0x00FFFFFFFFFFFFFFL)));
		//#endif

		return maxDepth;
	}

	/**
	 * @param ind
	 * @param state
	 * @return The method returns NaN as fitness if problem does not have test set.
	 */
	protected Pair<Double, Integer> getTestSetResults(final GPIndividual ind, final EvolutionState state) {
		if (!(state.evaluator.p_problem instanceof IChangeableEvaluationMode)) {
			return WRONG_TEST_SET_RESULT;
		}

		// get handles
		IChangeableEvaluationMode pMode = (IChangeableEvaluationMode) state.evaluator.p_problem;
		SimpleProblemForm problem = (SimpleProblemForm) state.evaluator.p_problem;
		KozaFitness trainingFitness = (KozaFitness) ind.fitness.clone();

		// evaluate on test set
		pMode.setEvaluationMode(EvaluationMode.TestSet);
		((SimpleNodeBase<?>) ind.trees[0].child).resetSemanticsRecursive();
		ind.evaluated = false;
		problem.evaluate(state, ind, 0, 0);

		KozaFitness testFitness = (KozaFitness) ind.fitness;
		ISemantics testSemantics = ((SimpleNodeBase<?>) ind.trees[0].child).getSemantics();

		for (int i = 0; i < testSemantics.size(); ++i) {
			DataSet db = this.getIteration().newChildDataSet("semantics");
			db.put("type", "test");
			db.put("i", i);
			db.put("s", testSemantics.getValue(i));
		}

		state.output.message(String.format("Test set fitness %.3f, semantics: %s", testFitness.standardizedFitness(), testSemantics));

		// revert mode to training set
		pMode.setEvaluationMode(EvaluationMode.TrainingSet);
		ind.fitness = trainingFitness;


		return new Pair<Double, Integer>((double) testFitness.standardizedFitness(), testFitness.hits);
	}

	/**
	 * Calculates number of unique semantics in population, (including all subtrees).
	 *
	 * @param state
	 * @return
	 */
	protected int getUniqueSemanticsCount(final EvolutionState state) {
		Population population = state.population;
		assert population.subpops.length > 0;

		HashSet<ISemantics> unique = new HashSet<ISemantics>(population.subpops[0].individuals.length << 3);

		for (int sp = 0; sp < population.subpops.length; ++sp) {
			Subpopulation subpop = population.subpops[sp];
			for (int i = 0; i < subpop.individuals.length; ++i) {
				GPIndividual ind = (GPIndividual) subpop.individuals[i];
				for (int t = 0; t < ind.trees.length; ++t) {
					GPTree tree = ind.trees[t];
					addSemanticsRecursive((SimpleNodeBase<?>) tree.child, unique);
				}
			}
		}

		return unique.size();
	}

	protected int getSemanticallyUniqueIndividualCount(final EvolutionState state) {
		Population population = state.population;
		assert population.subpops.length > 0;

		HashSet<ISemantics> unique = new HashSet<ISemantics>(population.subpops[0].individuals.length);

		for (int sp = 0; sp < population.subpops.length; ++sp) {
			Subpopulation subpop = population.subpops[sp];
			for (int i = 0; i < subpop.individuals.length; ++i) {
				GPIndividual ind = (GPIndividual) subpop.individuals[i];
				for (int t = 0; t < ind.trees.length; ++t) {
					GPTree tree = ind.trees[t];
					unique.add(((SimpleNodeBase<?>) tree.child).getSemantics());
				}
			}
		}

		return unique.size();
	}

	private static void addSemanticsRecursive(final SimpleNodeBase<?> node, HashSet<ISemantics> set) {
		set.add(node.getSemantics());
		for (int i = 0; i < node.children.length; ++i) {
			addSemanticsRecursive((SimpleNodeBase<?>) node.children[i], set);
		}
	}

	@Override
	public void preBreedingStatistics(EvolutionState state) {
		super.preBreedingStatistics(state);

		this.totalTimeNoStats.start();
		this.breedingTime.reset();
		this.breedingTime.start();
	}

	@Override
	public void postBreedingStatistics(EvolutionState state) {
		this.breedingTime.stop();
		this.totalTimeNoStats.stop();

		super.postBreedingStatistics(state);

		Runtime runtime = Runtime.getRuntime();
		long memory = runtime.totalMemory() - runtime.freeMemory();
		this.iteration.put("totalMemory", memory);

		// timers
		double generationTime = (this.breedingTime.getTimeElapsed() + this.evaluationTime.getTimeElapsed()) * 1.0E-9;
		double totalTime = this.totalTime.getTimeElapsed() * 1.0E-9;
		double totalTimeNoStats = this.totalTimeNoStats.getTimeElapsed() * 1.0E-9;
		this.iteration.put("breedTime", this.breedingTime.getTimeElapsed() * 1.0E-9);
		this.iteration.put("genTime", generationTime);
		this.iteration.put("totalTime", totalTime);
		this.iteration.put("totalTimeNoStats", totalTimeNoStats);


		// Do not save old iteration, wait for entire evolution to be saved.
		// In this way, we will require less locks on database file and therefore increase
		// level of parallelization of db access.
		//this.iteration.close();

		// prepare new iteration object
		if (!((EvoState) state).terminateCondition()) {
			this.iteration = this.experiment.newChildDataSet("generations");
			iteration.put("generation", state.generation + 1); // for next generation
		}

		state.output.message(String.format("Generation time: %10.3fs Total time: %10.3fs (w/o stats: %10.3fs) memory %dMB", generationTime, totalTime, totalTimeNoStats, memory >> 20));
	}

	@Override
	public void finalStatistics(final EvolutionState state, final int result) {
		super.finalStatistics(state, result);

		int initialGeneration = state.generation;
		for (; state.generation < state.numGenerations; ++state.generation) {
			// gather statistics for next generations (equal to statistics from last generation)
			this.preEvaluationStatistics(state);
			this.postEvaluationStatistics(state);

			if (state.generation == state.numGenerations - 1) {
				// make sure that the last tuple contains an evaluation count that is greater or equal the limit
				this.iteration.put("totalEvaluations",
						Math.max(SimpleNodeBase.getTotalEvaluationCount(), this.evaluationLimit));
			}

			this.preBreedingStatistics(state);
			this.postBreedingStatistics(state);
		}
		state.generation = initialGeneration;

		// since this is only place where experiment statistics are saved, we do not need to mark the experiment as finished
		//this.experiment.put("finished", true);
		this.experiment.close();

		for (IAdditionalStatistics stats : this.additionalStatistics) {
			stats.destroy(state);
		}
		this.additionalStatistics.clear();
	}
}
