/**
 * Copyright (c) 2012-2013, Daniele Codecasa <codecasa.job@gmail.com>,
 * Models and Algorithms for Data & Text Mining (MAD) laboratory of
 * Milano-Bicocca University, and all the CTBNCToolkit contributors
 * that will follow.
 * All rights reserved.
 *
 * @author Daniele Codecasa and all the CTBNCToolkit contributors that will follow.
 * @copyright 2012-2013 Daniele Codecasa, MAD laboratory, and all the CTBNCToolkit contributors that will follow
 */
package CTBNCToolkit.tests.validators;

import java.util.*;

import CTBNCToolkit.*;
import CTBNCToolkit.clustering.ClusteringResults;
import CTBNCToolkit.performances.*;

/**
 * @author Daniele Codecasa <codecasa.job@gmail.com>
 * 
 * K-fold cross validation method.
 * No random sorting of the dataset
 * is done. This guarantee the same
 * division if the cross validation
 * is used to test different algorithms
 * with the same k and the same dataset.
 * 
 * This algorithm return micro and
 * macro averaging performances.
 * 
 * @param <TimeType> type of the time interval in the trajectories (Integer = discrete time, Double = continuous time)
 * @param <NodeType> type of node in the model
 * @param <TM> type of the model used in the test
 * @param <PType> type of returned performances
 * @param <Stat> statistics generated by the learning process
 * @param <AggPType> type of aggregate performances returned
 * @param <SingPType> type of single run performances returned
 */
public class CrossValidation<TimeType extends Number & Comparable<TimeType>,NodeType extends INode,TM extends ICTClassifier<TimeType,NodeType>, AggPType extends IAggregatePerformances<TimeType,SingPType>, SingPType extends ISingleRunPerformances<TimeType>>
		extends
		BaseValidationMethod<TimeType, NodeType, TM, AggPType>
		implements
		IValidationMethod<TimeType, NodeType, TM, AggPType> {

	private IAggregatePerformancesFactory<TimeType, AggPType, SingPType> aggregatePerformancesFactory;
	
	private int k;
	private boolean leaveOneOut;
	private boolean memorizeFolding;
	private List<List<ITrajectory<TimeType>>> foldedDataset;
	private int nTimeDatasetUse;
	
	
	/**
	 * Base constructor.
	 * k=10 fold cross validation without
	 * folding memorization.
	 * 
	 * @param aggregatePerformancesFactory factory of aggregate performances to use in the tests.
	 */
	public CrossValidation(IAggregatePerformancesFactory<TimeType, AggPType, SingPType> aggregatePerformancesFactory) {
		
		this(aggregatePerformancesFactory, 10, false);
	}
	
	/**
	 * Constructor.
	 * 
	 * @param aggregatePerformancesFactory factory of aggregate performances to use in the tests.
	 * @param k number of folds in the cross validation (if k <= 0, leave one out modality is enabled)
	 */
	public CrossValidation(IAggregatePerformancesFactory<TimeType, AggPType, SingPType> aggregatePerformancesFactory, int k) {
		
		this(aggregatePerformancesFactory, k, false);
	}
	
	/**
	 * Constructor.
	 * 
	 * @param aggregatePerformancesFactory factory of aggregate performances to use in the tests.
	 * @param k number of folds in the cross validation (if k <= 0, leave one out modality is enabled)
	 * @param memorizeFolding flag that if activate force in the validation to use the last stored dataset partitioning
	 */
	public CrossValidation(IAggregatePerformancesFactory<TimeType, AggPType, SingPType> aggregatePerformancesFactory, int k, boolean memorizeFolding) {
		
		if( aggregatePerformancesFactory == null)
			throw new IllegalArgumentException("Error: null aggregate performances factory argument");
		
		this.aggregatePerformancesFactory = aggregatePerformancesFactory;
		
		this.k = k;
		this.leaveOneOut = (k <= 0);
		this.memorizeFolding = memorizeFolding;
		this.foldedDataset = null;
		this.nTimeDatasetUse = 0;
		
		super.setVerbose( false);
		super.setPrintSSFlag( false);
		super.setClusterInSampleFlag( false);
	}
	
	
	/**
	 * Set the aggregate performances
	 * factory.
	 * 
	 * @param aggregatePerformancesFactory factory of aggregate performances to use in the tests.
	 * @exception IllegalArgumentException in case of illegal argument.
	 */
	public void setAggregatePerformancesFactory(IAggregatePerformancesFactory<TimeType, AggPType, SingPType> aggregatePerformancesFactory) throws IllegalArgumentException {
		
		if( aggregatePerformancesFactory == null)
			throw new IllegalArgumentException("Error: null aggregate performances factory argument");
		
		this.aggregatePerformancesFactory = aggregatePerformancesFactory;
	}
	
	/**
	 * Return the number of folds used
	 * in the test.
	 * 
	 * @return number of folds
	 */
	public int getKFolds() {
		
		if( this.foldedDataset != null)
			return this.foldedDataset.size();
		else
			return this.k;
	}
	
	/**
	 * Activate or disable the modality
	 * that force the validator to use a
	 * previously loaded (and folded) dataset.
	 * 
	 * @param memorizeFolding flag that if activate force in the validation to use the last stored dataset partitioning
	 */
	public void setMemorizeFolding(boolean memorizeFolding) {
		
		this.memorizeFolding = memorizeFolding;
	}
	
	/**
	 * Reset the last folded dataset in order
	 * to store a new one.
	 */
	public void resetFoldedDataset() {
		
		this.foldedDataset = null;
		super.verbosePrint("Folded dataset reseted after being used " + this.nTimeDatasetUse + " times!\n");
	}
	
	/**
	 * True if the last folded dataset
 	 * memorization is enabled. False
 	 * otherwise.
	 * 
	 * @return true if the last folded dataset memorization is enable, false otherwise.
	 */
	public boolean isFoldsMemorizationEnabled() {
		
		return this.memorizeFolding;
	}
	
	/**
	 * Set a folded dataset to use in
	 * the tests.
	 * 
	 * @param foldedDataset folded dataset to load
	 * @throws RuntimeException in case of illegal arguments or state of the class
	 */
	public void setFoldedDataset(List<List<ITrajectory<TimeType>>> foldedDataset) throws RuntimeException {
		
		if( !this.memorizeFolding)
			throw new RuntimeException("Not enabled use of pre-loaded folded dataset (memorizeFolding flag must be true)");
		if( foldedDataset == null)
			throw new IllegalArgumentException("Null foldedDataset argument");
		if( foldedDataset.isEmpty())
			throw new IllegalArgumentException("Empty foldedDataset in input");
		
		super.verbosePrint("Load a new folded dataset!\n");		
		this.foldedDataset = foldedDataset;
		this.nTimeDatasetUse = 0;
	}
	
	/* (non-Javadoc)
	 * @see CTBNToolkit.tests.validators.IValidationMethod#validate(CTBNToolkit.IModel, CTBNToolkit.ILearningAlgorithm, CTBNToolkit.IClassifyAlgorithm, java.util.List)
	 */
	@SuppressWarnings({ "unchecked"})
	@Override
	public AggPType validate(TM mdl,
			ILearningAlgorithm<TimeType, NodeType> learnAlgo,
			IClassifyAlgorithm<TimeType, NodeType> infAlgo,
			List<ITrajectory<TimeType>> dataset) {

		// Initialization
		if( this.leaveOneOut)
			this.k = dataset.size();
		NodeIndexing nodeIndexing = mdl.getNodeIndexing();
		DiscreteNode classNode = mdl.getClassNode();
		Map<Integer,String> classIndexToValue = new TreeMap<Integer,String>();
		for(int i = 0; i < classNode.getStatesNumber(); ++i)
			classIndexToValue.put(i, classNode.getStateName(i));
		// Performance declaration
		AggPType aggOutOfSamplePerformances = this.aggregatePerformancesFactory.newAggregateInstance(classIndexToValue);
		AggPType aggInSamplePerformances = null;
		

		// K-fold division
		if( (!this.memorizeFolding) || this.foldedDataset == null) {
			super.verbosePrint("\n\nK-Folds Cross Validation schema using " + this.k + " folds (leav one out = " + this.leaveOneOut + ").\n");
			this.foldedDataset = this.makeKFolds(dataset, this.k);
		} else
			super.verbosePrint("\n\nK-Folds Cross Validation schema using " + this.foldedDataset.size() + " folds.\n");
		// Textual warning added because of the possible errors caused by accidental loading of pre-folded datasets
		if( this.memorizeFolding) {	
			this.nTimeDatasetUse++;
			super.verbosePrint("FOLDED DATASET USED FOR THE " + this.nTimeDatasetUse + " TIMES !!!\n");
		}

		
		// Cross Validation
		for ( int f = 0; f < foldedDataset.size(); ++f)
		{
			// Performances generation
			SingPType performances = this.aggregatePerformancesFactory.newSingleRunInstance(classIndexToValue);
			
			// Training set and test set generation
			List<ITrajectory<TimeType>> testSet  = this.foldedDataset.get( f);
			List<ITrajectory<TimeType>> trainSet = this.joinOtherFolds(this.foldedDataset, f);
			

			// Model learning
			super.verbosePrint( "\nCycle " + (f+1) + "\n");
			super.verbosePrint( "Learning: from " + trainSet.get(0).getName() + " To " + trainSet.get(trainSet.size()-1).getName() + "\n");
			super.verbosePrint( "          from " + testSet.get(0).getName() + " To " +  testSet.get(testSet.size()-1 ).getName() + " excluded\n");

			long startTime = System.currentTimeMillis();
			ILearningResults classificationResult = learnAlgo.learn( mdl, trainSet);
			double learningTime = (System.currentTimeMillis() - startTime) / 1000.0;
			// In sample performances
			if( super.getClusterInSampleFlag() && ClusteringResults.class.isAssignableFrom(classificationResult.getClass())) {
			
				if( aggInSamplePerformances == null)	// generate in sample aggregate container
					aggInSamplePerformances = this.aggregatePerformancesFactory.newAggregateInstance(classIndexToValue);
				
				// Create the in sample test performances
				SingPType inSampleSinglePerformance = this.aggregatePerformancesFactory.newSingleRunInstance(classIndexToValue);
				inSampleSinglePerformance.addResults((( ClusteringResults<TimeType>)classificationResult).getClusterizedTrajectories() );
				// Add the in sample test performances
				aggInSamplePerformances.addPerformances(inSampleSinglePerformance);
				inSampleSinglePerformance.calculateFinalResults(true);
			}
			// Sufficient statistics
			printSufficientStatistics(classificationResult, mdl);
			classificationResult = null;
			performances.setLearnedModel( (TM) mdl.clone());
			performances.setLearningTime( learningTime);
			
			super.verbosePrint("Model learnt in " + learningTime + " seconds\n\n");

			
			// Testing
			double prob;
			String classPrediction;
			String trueClass;
			super.verbosePrint( "Testing : from " +  testSet.get(0).getName() + " To " + testSet.get(testSet.size()-1).getName() + "\n");
			for( int iTrj = 0; iTrj < testSet.size(); ++iTrj)
			{				
				try
				{
					super.verbosePrint("Inference on " + testSet.get(iTrj).getName() + "\n");
					startTime = System.currentTimeMillis();
					IClassificationResult<TimeType> res = mdl.classify( infAlgo, testSet.get(iTrj));
					double inferenceTime =  (System.currentTimeMillis() - startTime) / 1000.0;
					
					classPrediction = res.getClassification();
					trueClass = res.getNodeValue(0, nodeIndexing.getClassIndex());
					if( infAlgo.probabilityFlag()) {
						prob = res.getProbability(classPrediction);
						super.verbosePrint( testSet.get( iTrj).getName() + ": True Class: " + trueClass + ", Predicted: " + classPrediction + ", Probability: " + prob + ", Inference time: " + inferenceTime + "\n");
					} else
						super.verbosePrint( testSet.get( iTrj).getName() + ": True Class: " + trueClass + ", Predicted: " + classPrediction + ", Inference time: " + inferenceTime + "\n");
					
					performances.addResult( res, inferenceTime);
				}
				catch(Exception e)
				{
					System.out.println(e.getMessage());
				}
			}
			
			aggOutOfSamplePerformances.addPerformances(performances);
		}
		
		if( super.getClusterInSampleFlag() && IExternalClusteringPerformances.class.isAssignableFrom( aggOutOfSamplePerformances.getClass()))
			(( IExternalClusteringPerformances) aggOutOfSamplePerformances).setInSamplePerformances(
					(IExternalClusteringPerformances) aggInSamplePerformances);

		return aggOutOfSamplePerformances;
	}

	/**
	 * Divide the dataset in folds.
	 * 
	 * @param dataset dataset to divide
	 * @param k number of folds
	 * @return k-folds
	 */
	private List<List<ITrajectory<TimeType>>> makeKFolds(List<ITrajectory<TimeType>> dataset, int k)
	{
		List<List<ITrajectory<TimeType>>> folds = new ArrayList<List<ITrajectory<TimeType>>>(k);

		int foldDim, surplus, index2;
		foldDim  = dataset.size() / k;
		surplus  = dataset.size() % k;
		for (int index1 = 0; index1 < dataset.size(); index1 = index2)
		{			
			index2 = index1 + foldDim;
			if( surplus != 0) {
				++index2;
				--surplus;
			}
			
			folds.add( dataset.subList(index1, index2));
		}
		
		return folds;
	}
	
	
	/**
	 * Merge all the folds in one list
	 * excluding the one indicated by
	 * the index.
	 * 
	 * @param folds list of folds
	 * @param index index of the fold to exclude
	 * @return merged list
	 */
	private ArrayList<ITrajectory<TimeType>> joinOtherFolds(List<List<ITrajectory<TimeType>>> folds, int index)
	{	
		ArrayList<ITrajectory<TimeType>> joinedList = new ArrayList<ITrajectory<TimeType>>(folds.get(0).size() * (folds.size()-1));

		for ( int i = 0; i < folds.size(); ++i)
			if (i != index)
				joinedList.addAll(folds.get( i));

		return joinedList;
	}

}
