/* ========================================================================== *
 *  Test program.
 *
 *  Author: Firas Safadi (fsafadi@ulg.ac.be)
 * ========================================================================== */

#include "Matrix.h"
#include "GeneticAlgorithm.h"
#include "CascadeCorrelation.h"
#include "sort.h"

//Uncomment the following line to detect memory leaks
#define DETECT_MEMORY_LEAKS

//Uncomment the following line to enable writing to disk
#define DISABLE_DISK_WRITES

#ifdef DETECT_MEMORY_LEAKS
#define CRTDBG_MAP_ALLOC
#include <stdlib.h>
#include <crtdbg.h>
#endif

#ifndef DISABLE_DISK_WRITES
#define MgtWrite(s) s
#else
#define MgtWrite(s)
#endif

typedef struct perfrecord_t {
  nat popCount;
  nat popSize;
  nat hiddenLayers;
  real trainingError;
  real testError;
  float trainingTime;
} PerfRecord;

void test_matrix_mult()
{
  char buf[BUF_SIZE];

  cublasHandle_t handle;
  cublas_call(cublasCreate(&handle));

  Matrix A("A.txt", Matrix::HostAndDevice);
  Matrix B("B.txt", Matrix::HostAndDevice);
  Matrix C;

  Matrix D(5, 5, Matrix::HostAndDevice);

  C.resize(9, 7);
  C.spawn();
  C.despawn();
  C.resize(7, 9);
  C.spawn();
  C.init(0);

  std::cout << std::endl;
  std::cout << "Total memory reserved on host: " << size_to_string(buf, Matrix::totalBytesOnHost()) << std::endl;
  std::cout << "Total memory reserved on device: " << size_to_string(buf, Matrix::totalBytesOnDevice()) << std::endl;

  A.dispInfo();
  B.dispInfo();

  matrix_mult_h(A, B, C, false, true);
  matrix_mult_d(handle, A, B, C, false, true);

  C.dispInfo();
  C.dispData();

  std::cout << std::endl;
  std::cout << "Data consistency: " << (C.consistent() ? "Yes" : "No") << std::endl;

  cudaEvent_t start, stop;
  float elapsedTime;

  cuda_call(cudaEventCreate(&start));
  cuda_call(cudaEventCreate(&stop));

  cuda_call(cudaEventRecord(start));

  D.init(5.0);

  cuda_call(cudaEventRecord(stop));

  cuda_call(cudaEventSynchronize(stop));

  cuda_call(cudaEventElapsedTime(&elapsedTime, start, stop));

  D.dispInfo();
  D.dispData();

  MgtWrite(D.write("D.txt"));

  std::cout << std::endl;
  std::cout << "Elapsed time: " << elapsedTime << " ms" << std::endl;

  std::cout << std::endl;
  std::cout << "Data synced: " << (D.synced() ? "Yes" : "No") << std::endl;

  cublas_call(cublasDestroy(handle));
}

void test_bitonic_sort()
{
  nat iter = 1;
  real total = 0.0;

  for (nat e = 0; e <= 1; e++) {
    nat batchSize = 64;
    nat arrayLength = 16;

    std::cout << "testing n = " << arrayLength << "... ";

    for (nat i = 0; i < iter; i++) {
      Matrix M(1, batchSize*arrayLength, Matrix::HostAndDevice);
      Natrix N(1, batchSize*arrayLength, Natrix::HostAndDevice);

      for (nat m = 0; m < batchSize; m++) {
        for (nat n = 0; n < arrayLength; n++) {
          M(0, m*arrayLength + n) = (real)rand() / (real)RAND_MAX;
          N(0, m*arrayLength + n) = n & (arrayLength - 1);
        }
      }

      M.syncDeviceWithHost();
      N.syncDeviceWithHost();

      cudaEvent_t start, stop;
      float elapsedTime;

      cuda_call(cudaEventCreate(&start));
      cuda_call(cudaEventCreate(&stop));

      cuda_call(cudaEventRecord(start));

      bitonic_sort(M.dp(), N.dp(), batchSize, arrayLength);

      cuda_call(cudaEventRecord(stop));

      cuda_call(cudaEventSynchronize(stop));

      cuda_call(cudaEventElapsedTime(&elapsedTime, start, stop));

      total += elapsedTime;

      M.syncHostWithDevice();
      N.syncHostWithDevice();

      for (nat m = 0; m < batchSize; m++) {
        std::cout << std::endl;
        for (nat n = 0; n < arrayLength; n++)
          std::cout << M(0, m*arrayLength + n) << " ";
        std::cout << std::endl;
        for (nat n = 0; n < arrayLength; n++)
          std::cout << N(0, m*arrayLength + n) << " ";
        std::cout << std::endl;
      }
    }

    std::cout << "Average time = " << total / iter << std::endl;
  }
}

void test_genetic_algorithm()
{
  char buf[BUF_SIZE];

  nat popCount = 1024;
  nat popSize = 32;
  nat maxGenSize = 32;
  nat genSize = maxGenSize;

  nat epochs = 100;

  nat threadsPerBlock = 128;

  GeneticAlgorithm ga(popCount, popSize, genSize, maxGenSize);

  ga.setSpeciesGenetics(-6.0, -4.0, 4.0, 6.0, -1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.4, 0.6, 0.001, 0.01, GeneticAlgorithm::Random);
  ga.setEvolutionEpochs(1000);
  ga.setThreadsPerBlock(threadsPerBlock);

  ga.hatchSpecies();

  std::cout << std::endl;
  std::cout << "Total memory reserved on host (approx.): " << size_to_string(buf, Matrix::totalBytesOnHost()) << std::endl;
  std::cout << "Total memory reserved on device (approx.): " << size_to_string(buf, Matrix::totalBytesOnDevice()) << std::endl;

  cudaEvent_t start, stop;
  float elapsedTime;

  cuda_call(cudaEventCreate(&start));
  cuda_call(cudaEventCreate(&stop));

  cuda_call(cudaEventRecord(start));

  ga.evolveSpecies();

  cuda_call(cudaEventRecord(stop));

  cuda_call(cudaEventSynchronize(stop));

  cuda_call(cudaEventElapsedTime(&elapsedTime, start, stop));

  std::cout << std::endl;
  std::cout << "Elapsed time: " << elapsedTime << " ms" << std::endl;

  MgtWrite(ga.getPrevPop().write("prevpop.txt"));
  MgtWrite(ga.getPrevFit().write("prevfit.txt"));
  MgtWrite(ga.getPrevBoard().write("prevbrd.txt"));

  MgtWrite(ga.getCurrPop().write("pop.txt"));
  MgtWrite(ga.getCurrFit().write("fit.txt"));
  MgtWrite(ga.getCurrBoard().write("brd.txt"));

  nat topPop;
  nat topInd;
  real topFit;

  ga.topPerf(topPop, topInd, topFit);

  std::cout << std::endl;
  std::cout << "Top performer: <" << topPop << ", #" << topInd << "> *** " << topFit << " ***" << std::endl;

  std::cout << std::endl;
  std::cout << "Genetics:" << std::endl;

  ga.dispGenetics(topPop);
}

void test_cascade_correlation()
{
  char buf[BUF_SIZE];

  nat popCount = 1;
  nat popSize = 1024;
  nat threadsPerBlock = 128;

  CascadeCorrelation cc;

  cc.setNetworkParameters(506, 0.2, 13, 1, 1.0, 0.0, 114, 50, true);
  cc.setUnitEAParameters(popCount, popSize, CascadeCorrelation::GeneticAlgorithmEAID);
  cc.setOutputEAParameters(popCount, popSize, CascadeCorrelation::GeneticAlgorithmEAID);
  cc.buildNetwork();

  GeneticAlgorithm* uga = dynamic_cast<GeneticAlgorithm*>(cc.getUnitEA());
  GeneticAlgorithm* oga = dynamic_cast<GeneticAlgorithm*>(cc.getOutputEA());

  uga->setSpeciesGenetics(-1.0, 0.0, 0.0, 1.0, -3.0, 0.0, 0.0, 3.0, 0.04, 0.06, 0.44, 0.46, 0.44, 0.46, 0.4, 0.6, 0.001, 0.1, GeneticAlgorithm::Random);
  uga->setEvolutionEpochs(50);
  uga->setThreadsPerBlock(threadsPerBlock);

  oga->setSpeciesGenetics(-1.0, 0.0, 0.0, 1.0, -2.0, 0.0, 0.0, 2.0, 0.04, 0.06, 0.44, 0.46, 0.44, 0.46, 0.4, 0.6, 0.001, 0.1, GeneticAlgorithm::Random);
  oga->setEvolutionEpochs(50);
  oga->setThreadsPerBlock(threadsPerBlock);

  cc.loadDataSet("housing.dat");

  std::cout << std::endl;
  std::cout << "Total memory reserved on host (approx.): " << size_to_string(buf, Matrix::totalBytesOnHost()) << std::endl;
  std::cout << "Total memory reserved on device (approx.): " << size_to_string(buf, Matrix::totalBytesOnDevice()) << std::endl;

  cc.trainNetwork();

  nat topPop;
  nat topInd;
  real topFit;

  if (cc.hiddenLayers() > 0) {
    uga->topPerf(topPop, topInd, topFit);

    std::cout << std::endl;
    std::cout << "Unit genetics:" << std::endl;

    uga->dispGenetics(topPop);
  }

  oga->topPerf(topPop, topInd, topFit);

  std::cout << std::endl;
  std::cout << "Output genetics:" << std::endl;

  oga->dispGenetics(topPop);

  if (yes_no("Test network?"))
    cc.testNetwork();

  if (yes_no("Write network to disk?"))
    cc.writeNetwork("network.txt");
}

void record_performance(PerfRecord& pr)
{
  nat popCount = pr.popCount;
  nat popSize = pr.popSize;
  nat threadsPerBlock = 128;

  CascadeCorrelation cc;

  cc.setNetworkParameters(506, 0.2, 13, 1, 1.0, 0.0, 114, 50, true);
  cc.setUnitEAParameters(popCount, popSize, CascadeCorrelation::GeneticAlgorithmEAID);
  cc.setOutputEAParameters(popCount, popSize, CascadeCorrelation::GeneticAlgorithmEAID);
  cc.buildNetwork();

  GeneticAlgorithm* uga = dynamic_cast<GeneticAlgorithm*>(cc.getUnitEA());
  GeneticAlgorithm* oga = dynamic_cast<GeneticAlgorithm*>(cc.getOutputEA());

  uga->setSpeciesGenetics(-1.0, 0.0, 0.0, 1.0, -3.0, 0.0, 0.0, 3.0, 0.04, 0.06, 0.44, 0.46, 0.44, 0.46, 0.4, 0.6, 0.001, 0.1, GeneticAlgorithm::Random);
  uga->setEvolutionEpochs(50);
  uga->setThreadsPerBlock(threadsPerBlock);

  oga->setSpeciesGenetics(-1.0, 0.0, 0.0, 1.0, -2.0, 0.0, 0.0, 2.0, 0.04, 0.06, 0.44, 0.46, 0.44, 0.46, 0.4, 0.6, 0.001, 0.1, GeneticAlgorithm::Random);
  oga->setEvolutionEpochs(50);
  oga->setThreadsPerBlock(threadsPerBlock);

  cc.loadDataSet("housing.dat");

  cc.enableSilentMode();

  Timer t;

  t.start();

  cc.trainNetwork();

  t.stop();

  pr.trainingTime = t.sec();

  cc.topPerf(pr.hiddenLayers, pr.trainingError, pr.testError);
}

void run_performance_tests()
{
  const nat RUNS = 100;

  PerfRecord prTable[RUNS];

  File records("records.csv", std::ios_base::out);

  records << "Population Count" << "," << "Population Size" << ","
    << "Average Hidden Layers" << "," << "Average Training Error" << "," << "Average Test Error" << "," << "Average Training Time" << ","
    << "Best Hidden Layers" << "," << "Best Training Error" << "," << "Best Test Error" << "," << "Best Training Time" << ","
    << "Worst Hidden Layers" << "," << "Worst Training Error" << "," << "Worst Test Error" << "," << "Worst Training Time" << File::endl;

  for (nat size = 1024; size <= 1024; size <<= 1) {

    std::cout << std::endl;

    for (nat popSize = size; popSize >= 32; popSize >>= 1) {

      nat popCount = size / popSize;

      for (nat i = 0; i < RUNS; i++) {
        std::cout << "\rtesting " << popCount << "x" << popSize << "... " << i << "/" << RUNS;

        prTable[i].popCount = popCount;
        prTable[i].popSize = popSize;

        record_performance(prTable[i]);
      }

      std::cout << "\rtesting " << popCount << "x" << popSize << "... " << RUNS << "/" << RUNS << std::endl;

      nat best = 0;
      nat worst = 0;

      real avgHiddenLayers = prTable[0].hiddenLayers;
      real avgTrainingError = prTable[0].trainingError;
      real avgTestError = prTable[0].testError;
      real avgTrainingTime = prTable[0].trainingTime;

      for (nat i = 1; i < RUNS; i++) {
        avgHiddenLayers += prTable[i].hiddenLayers;
        avgTrainingError += prTable[i].trainingError;
        avgTestError += prTable[i].testError;
        avgTrainingTime += prTable[i].trainingTime;

        if (prTable[i].testError < prTable[best].testError)
          best = i;

        if (prTable[i].testError > prTable[worst].testError)
          worst = i;
      }

      avgHiddenLayers /= RUNS;
      avgTrainingError /= RUNS;
      avgTestError /= RUNS;
      avgTrainingTime /= RUNS;

      records << popCount << "," << popSize << ","
        << avgHiddenLayers << "," << avgTrainingError << "," << avgTestError << "," << avgTrainingTime << ","
        << prTable[best].hiddenLayers << "," << prTable[best].trainingError << "," << prTable[best].testError << "," << prTable[best].trainingTime << ","
        << prTable[worst].hiddenLayers << "," << prTable[worst].trainingError << "," << prTable[worst].testError << "," << prTable[worst].trainingTime << File::endl;
    }
  }

  records.close();
}

void record_performance_2(PerfRecord& pr, Matrix& dataset, bool* inTestSet)
{
  nat popCount = pr.popCount;
  nat popSize = pr.popSize;
  nat threadsPerBlock = 128;

  CascadeCorrelation cc;

  cc.setNetworkParameters(506, 0.2, 13, 1, 1.0, 0.0, 114, 50, false);
  cc.setUnitEAParameters(popCount, popSize, CascadeCorrelation::GeneticAlgorithmEAID);
  cc.setOutputEAParameters(popCount, popSize, CascadeCorrelation::GeneticAlgorithmEAID);
  cc.buildNetwork();

  GeneticAlgorithm* uga = dynamic_cast<GeneticAlgorithm*>(cc.getUnitEA());
  GeneticAlgorithm* oga = dynamic_cast<GeneticAlgorithm*>(cc.getOutputEA());

  uga->setSpeciesGenetics(-1.0, 0.0, 0.0, 1.0, -3.0, 0.0, 0.0, 3.0, 0.04, 0.06, 0.44, 0.46, 0.44, 0.46, 0.4, 0.6, 0.001, 0.1, GeneticAlgorithm::Random);
  uga->setEvolutionEpochs(50);
  uga->setThreadsPerBlock(threadsPerBlock);

  oga->setSpeciesGenetics(-1.0, 0.0, 0.0, 1.0, -2.0, 0.0, 0.0, 2.0, 0.04, 0.06, 0.44, 0.46, 0.44, 0.46, 0.4, 0.6, 0.001, 0.1, GeneticAlgorithm::Random);
  oga->setEvolutionEpochs(50);
  oga->setThreadsPerBlock(threadsPerBlock);

  cc.loadDataSet(dataset, inTestSet);

  cc.enableSilentMode();

  Timer t;

  t.start();

  cc.trainNetwork();

  t.stop();

  pr.trainingTime = t.sec();

  cc.topPerf(pr.hiddenLayers, pr.trainingError, pr.testError);
}

void run_performance_tests_2()
{
  const nat TESTS = 6;

  const nat RUNS = 100;

  bool inTestSet[506];

  PerfRecord prTable[RUNS * 2];

  Matrix dataset(512, 16);

  dataset.read("housing.dat", 506, 14);

  File records("records.csv", std::ios_base::app);

  records << "Population Count" << "," << "Population Size" << "," << "Training Time" << "," << "Hidden Layers" << "," << "Training Error" << "," << "Test Error" << ","
    << "Population Count" << "," << "Population Size" << "," << "Training Time" << "," << "Hidden Layers" << "," << "Training Error" << "," << "Test Error" << File::endl;

  for (nat t = 0; t < TESTS; t++) {
    std::fill(inTestSet, inTestSet + 405, false);
    std::fill(inTestSet + 405, inTestSet + 506, true);
    std::random_shuffle(inTestSet, inTestSet + 506);

    nat popCount = 32;
    nat popSize = 128;

    std::cout << std::endl;
    std::cout << "TEST " << t + 1 << "/" << TESTS << std::endl;

    for (nat i = 0; i < RUNS; i++) {
      std::cout << "\rtesting " << popCount << "x" << popSize << "... " << i << "/" << RUNS;

      prTable[i].popCount = popCount;
      prTable[i].popSize = popSize;

      record_performance_2(prTable[i], dataset, inTestSet);
    }

    std::cout << "\rtesting " << popCount << "x" << popSize << "... " << RUNS << "/" << RUNS << std::endl;

    popCount = 1;
    popSize = 4096;

    for (nat i = 0; i < RUNS; i++) {
      std::cout << "\rtesting " << popCount << "x" << popSize << "... " << i << "/" << RUNS;

      prTable[i + RUNS].popCount = popCount;
      prTable[i + RUNS].popSize = popSize;

      record_performance_2(prTable[i + RUNS], dataset, inTestSet);
    }

    std::cout << "\rtesting " << popCount << "x" << popSize << "... " << RUNS << "/" << RUNS << std::endl;

    for (nat i = 0; i < RUNS; i++)
      records << prTable[i].popCount << "," << prTable[i].popSize << "," << prTable[i].trainingTime << "," << prTable[i].hiddenLayers << "," << prTable[i].trainingError << "," << prTable[i].testError << ","
        << prTable[i + RUNS].popCount << "," << prTable[i + RUNS].popSize << "," << prTable[i + RUNS].trainingTime << "," << prTable[i + RUNS].hiddenLayers << "," << prTable[i + RUNS].trainingError << "," << prTable[i + RUNS].testError << File::endl;

    records << File::endl;
  }

  records.close();
}

void run_performance_tests_3()
{
  const nat RUNS = 100;

  PerfRecord prTable[RUNS];

  File records("records.csv", std::ios_base::out);

  records << "Population Count" << "," << "Population Size" << "," << "Training Time" << "," << "Hidden Layers" << "," << "Training Error" << "," << "Test Error" << File::endl;

  for (nat size = 1024; size <= 1024; size <<= 1) {

    std::cout << std::endl;

    for (nat popSize = size; popSize >= 32; popSize >>= 1) {

      nat popCount = size / popSize;

      for (nat i = 0; i < RUNS; i++) {
        std::cout << "\rtesting " << popCount << "x" << popSize << "... " << i << "/" << RUNS;

        prTable[i].popCount = popCount;
        prTable[i].popSize = popSize;

        record_performance(prTable[i]);
      }

      std::cout << "\rtesting " << popCount << "x" << popSize << "... " << RUNS << "/" << RUNS << std::endl;

      for (nat i = 0; i < RUNS; i++)
        records << prTable[i].popCount << "," << prTable[i].popSize << "," << prTable[i].trainingTime << "," << prTable[i].hiddenLayers << "," << prTable[i].trainingError << "," << prTable[i].testError << File::endl;

      records << File::endl;
    }
  }

  records.close();
}

int main()
{
  std::cout << "Initializing..." << std::endl;

  int deviceCount = 0;
  cuda_call(cudaGetDeviceCount(&deviceCount));

  if (deviceCount == 0)
    error("CUDA device not found");

  cudaDeviceProp deviceProp;
  cuda_call(cudaSetDevice(0));
  cuda_call(cudaGetDeviceProperties(&deviceProp, 0));

  if (deviceProp.major < 2)
    error("CUDA device not supported");

  std::cout << std::endl;
  std::cout << "Device: " << deviceProp.name << std::endl;

  std::cout << std::endl;
  std::cout << "CUDA cores: " << cuda_cores(deviceProp) << std::endl;

  srand(time(0));

  //test_matrix_mult();

  //test_bitonic_sort();

  //test_genetic_algorithm();

  //test_cascade_correlation();

  //run_performance_tests();

  //run_performance_tests_2();

  run_performance_tests_3();

  cudaDeviceReset();

  pause();

#ifdef DETECT_MEMORY_LEAKS
  _CrtDumpMemoryLeaks();
#endif

  return 0;
}
