/* ========================================================================== *
 *  Evolutionary cascade-correlation.
 *
 *  Author: Firas Safadi (fsafadi@ulg.ac.be)
 * ========================================================================== */

#pragma once

#include "EvolutionaryAlgorithm.h"

class CascadeCorrelation;

class CCSimulator : public Simulator
{
public:
  typedef void(CascadeCorrelation::*CCSimFunc)(EvolutionaryAlgorithm*);

  CCSimulator(CascadeCorrelation* cc, CCSimFunc sim);
  ~CCSimulator();

  void sim(EvolutionaryAlgorithm* ea);

protected:
  CascadeCorrelation* m_pCascadeCorrelation;
  CCSimFunc m_fSim;
};

class CascadeCorrelation
{
public:
  typedef enum { GeneticAlgorithmEAID, CrossEntropyEAID } EAID;

  //Constructor
  CascadeCorrelation();
  //Destructor
  ~CascadeCorrelation();

  //Loads dataset from file
  void loadDataSet(char* path);
  //Loads dataset from matrix
  void loadDataSet(Matrix& dataset, bool* inTestSet = NULL);

  //Builds the network according to set parameters
  void buildNetwork();
  //Trains the network for the loaded dataset
  void trainNetwork();

  //Returns the current MSE observed on the training set
  real trainingError() const { return m_dTrainingError; }
  //Returns the current MSE observed on the test set
  real testError() const { return m_dTestError; }

  //Returns the hidden layer optimizer
  EvolutionaryAlgorithm* getUnitEA() { return m_pUnitEA; }
  //Returns the output layer optimizer
  EvolutionaryAlgorithm* getOutputEA() { return m_pOutputEA; }

  //Returns the maximum cache width in columns
  nat maxCacheWidth() const { return 1 + m_iInputDim + m_iMaxHiddenLayers; }
  //Returns the current cache width in columns
  nat cacheWidth() const { return 1 + m_iInputDim + m_iHiddenLayers; }
  //Returns the current number of hidden layers
  nat hiddenLayers() const { return m_iHiddenLayers; }

  //Sets the network parameters
  void setNetworkParameters(nat samples, real testProp, nat inputDim, nat outputDim, real bias, real errorGoal, nat maxHiddenLayers, nat overfitStop, bool randomTestSamples);

  //Sets the parameters of the hidden layer optimizer
  void setUnitEAParameters(nat popCount, nat popSize, EAID id);
  //Sets the parameters of the output layer optimizer
  void setOutputEAParameters(nat popCount, nat popSize, EAID id);

  //Disables training progress display
  void disableSilentMode() { m_bSilentMode = false; }
  //Enables training progress display
  void enableSilentMode() { m_bSilentMode = true; }

  //Writes network to file
  void writeNetwork(char* path) const;
  //Tests network on the test set
  void testNetwork() const;

  //Retrieves top performance record
  void topPerf(nat& hiddenLayers, real& trainingError, real& testError);

protected:
  //Returns the top performer for the output layer
  nat topOutputInd() const { return m_pOutputEA->popSize()*m_iTopPopulation + m_iTopIndividual; }

  Matrix m_pTrainingCache;
  Matrix m_pTestCache;

  Matrix m_pTrainingTargets;
  Matrix m_pTestTargets;

  Matrix m_pTrainingError;

  Matrix m_pUnitWeights;
  Matrix m_pOutputWeights;

  Matrix m_pTrainingSimulation;
  Matrix m_pTestSimulation;

  Matrix m_pOutputWeightsHistory;
  Matrix m_pErrorHistory;

  Natrix m_pTestSetIDs;

  EvolutionaryAlgorithm* m_pUnitEA;
  EvolutionaryAlgorithm* m_pOutputEA;

  CCSimulator* m_pUnitSimulator;
  CCSimulator* m_pOutputSimulator;

  nat m_iSamples;
  nat m_iInputDim;
  nat m_iOutputDim;

  real m_dPropTestSamples;
  nat m_iTrainingSamples;
  nat m_iTestSamples;

  nat m_iMaxHiddenLayers;
  nat m_iHiddenLayers;

  real m_dBias;

  real m_dErrorGoal;
  real m_dTrainingError;
  real m_dTestError;

  nat m_iOptimalLayerCount;

  nat m_iOverfitCounter;
  nat m_iOverfitStop;

  bool m_bNetworkBuilt;
  bool m_bNetworkTrained;
  bool m_bDataLoaded;

  bool m_bRandomizeTestSamples;
  bool m_bSilentMode;

  cublasHandle_t m_pHandle;

  nat m_iTopPopulation;
  nat m_iTopIndividual;
  real m_dTopFitness;

  nat m_iUnitEAPopCount;
  nat m_iUnitEAPopSize;

  nat m_iOutputEAPopCount;
  nat m_iOutputEAPopSize;

  EAID m_eUnitEAID;
  EAID m_eOutputEAID;

private:
  //Creates an evolutionary algorithm of the specified kind
  EvolutionaryAlgorithm* makeEA(nat popCount, nat popSize, EAID id);

  //Adds a new hidden layer
  void addHiddenUnit();
  //Trains output layer
  void trainOutputWeights();

  //Computes the network error
  void computeNetworkError();
  //Computes the training error
  void computeTrainingError();
  //Computes the test error
  void computeTestError();
  
  //Evaluates hidden layer optimizer populations
  void simUnit(EvolutionaryAlgorithm* ea);
  //Evaluates output layer optimizer populations
  void simOutput(EvolutionaryAlgorithm* ea);

  //Validates network parameters
  void validateNetworkParameters();
};
