/* ========================================================================== *
 *  Base class for all evolutionary algorithms.
 *
 *  Author: Firas Safadi (fsafadi@ulg.ac.be)
 * ========================================================================== */

#include "EvolutionaryAlgorithm.h"
#include "sort.h"

#define DEFAULT_THREADS_PER_BLOCK 256

#ifndef DOUBLE_PRECISION
#define FABS fabsf
#define TANH tanhf
#else
#define FABS fabs
#define TANH tanh
#endif

class BenchSimulator : public Simulator
{
public:
  BenchSimulator() {}
  ~BenchSimulator() {}
  void sim(EvolutionaryAlgorithm* ea);
};

__global__ void bench_sim_K(real* pop, nat popLdt, real* fit, nat* board, nat totalInd, nat popSize, nat genSize)
{
  nat x = threadIdx.x + blockIdx.x*blockDim.x;

  if (x < totalInd) {
    board[x] = x & (popSize - 1);

    pop += x;

    real* end = pop + genSize*popLdt;

    real f = genSize;

    while (pop < end) {
      f -= FABS(TANH(*pop - 1.0));
      pop += popLdt;
    }

    fit[x] = f;
  }
}

void BenchSimulator::sim(EvolutionaryAlgorithm* ea)
{
  if (!ea)
    error("Cannot run simulation for null evolutionary algorithm");

  AsyncKernelLaunch(bench_sim_K, RDIV(ea->totalInd(), DEFAULT_THREADS_PER_BLOCK), DEFAULT_THREADS_PER_BLOCK, ea->getCurrPop().dp(), ea->getCurrPop().ldt(), ea->getCurrFit().dp(), ea->getCurrBoard().dp(), ea->totalInd(), ea->popSize(), ea->genSize());
}

EvolutionaryAlgorithm::EvolutionaryAlgorithm(nat popCount, nat popSize, nat genSize, nat maxGenSize)
{
  m_iPopulationCount = popCount;
  m_iPopulationSize = popSize;
  m_iTotalIndividuals = popCount*popSize;
  m_iGenomeSize = genSize;
  m_iMaxGenomeSize = maxGenSize;

  validateSpeciesParameters();

  m_bAlive = false;
  m_bPreGenerate = true;

  m_iEvolutionEpochs = 1;
  m_iCurrentGeneration = 0;

  m_iThreadsPerBlock = DEFAULT_THREADS_PER_BLOCK;
  m_iBlocks = RDIV(m_iTotalIndividuals, m_iThreadsPerBlock);

  m_pDefaultSimulator = new BenchSimulator();
  m_pSimulator = m_pDefaultSimulator;

  m_pPopulation.resize(m_iMaxGenomeSize, m_iTotalIndividuals);
  m_pPopulationAlt.resize(m_iMaxGenomeSize, m_iTotalIndividuals);
  m_pFitness.resize(1, m_iTotalIndividuals);
  m_pFitnessAlt.resize(1, m_iTotalIndividuals);
  m_pLeaderboard.resize(1, m_iTotalIndividuals);
  m_pLeaderboardAlt.resize(1, m_iTotalIndividuals);

  m_pPopulation.spawn();
  m_pPopulationAlt.spawn();
  m_pFitness.spawn();
  m_pFitnessAlt.spawn();
  m_pLeaderboard.spawn();
  m_pLeaderboardAlt.spawn();

  m_pPopulation.initDevice(0.0);
  m_pPopulationAlt.initDevice(0.0);
  m_pFitness.initDevice(0.0);
  m_pFitnessAlt.initDevice(0.0);
  m_pLeaderboard.initDevice(0);
  m_pLeaderboardAlt.initDevice(0);

  initRandomGenerator();

  cublas_call(cublasCreate(&m_pHandle));
}

EvolutionaryAlgorithm::~EvolutionaryAlgorithm()
{
  delete m_pDefaultSimulator;

  cuda_call(cudaFree(m_pRandSpring));

  cublas_call(cublasDestroy(m_pHandle));
}

void EvolutionaryAlgorithm::incGenSize()
{
  if (!m_bAlive)
    error("Cannot increase individual length: not alive");

  if (!m_bPreGenerate)
    error("Cannot increase individual length with pre-generation disabled");

  if (m_iGenomeSize == m_iMaxGenomeSize)
    error("Maximum individual length reached");

  m_iGenomeSize++;
}

void EvolutionaryAlgorithm::setSimulator(Simulator* sim)
{
  if (!m_pSimulator)
    error("Simulator cannot be null");

  m_pSimulator = sim;
}

void EvolutionaryAlgorithm::topPerf(nat& pop, nat& ind, real& fit)
{
  if (!m_bAlive)
    error("Cannot get top individual performance: not alive");

  getCurrFit().syncHostWithDevice();
  getCurrBoard().syncHostWithDevice();

  real* fp = getCurrFit().hp();
  nat* bp = getCurrBoard().hp();

  fit = fp[0];
  ind = bp[0];
  pop = 0;

  for (nat i = 1; i < m_iPopulationCount; i++) {
    if (fp[i*m_iPopulationSize] > fit) {
      fit = fp[i*m_iPopulationSize];
      ind = bp[i*m_iPopulationSize];
      pop = i;
    }
  }
}

void EvolutionaryAlgorithm::setThreadsPerBlock(nat n)
{
  if (m_bAlive)
    error("Cannot modify threads per block while population is alive");

  if (n < 32)
    error("Number of threads per block cannot be lower than 32");

  if (n & (n - 1))
    error("Number of threads per block must be a power of 2");

  m_iThreadsPerBlock = n;
  m_iBlocks = RDIV(m_iTotalIndividuals, n);
}


void EvolutionaryAlgorithm::setEvolutionEpochs(nat n)
{
  if (n == 0)
    error("Number of generations in an evolution cycle must be greater than 0");

  m_iEvolutionEpochs = n;
}


void EvolutionaryAlgorithm::evalPop()
{
  if (!m_bAlive)
    error("Cannot evaluate population: not alive");

  if (!m_pSimulator)
    error("Missing fitness function");

  m_pSimulator->sim(this);
}

void EvolutionaryAlgorithm::sortPop()
{
  if (!m_bAlive)
    error("Cannot sort population: not alive");

  bitonic_sort(getCurrFit().dp(), getCurrBoard().dp(), m_iPopulationCount, m_iPopulationSize);
}

Matrix& EvolutionaryAlgorithm::getPrevPop()
{
  return ODD(m_iCurrentGeneration) ? m_pPopulationAlt : m_pPopulation;
}

Matrix& EvolutionaryAlgorithm::getCurrPop()
{
  return ODD(m_iCurrentGeneration) ? m_pPopulation : m_pPopulationAlt;
}

Matrix& EvolutionaryAlgorithm::getPrevFit()
{
  return ODD(m_iCurrentGeneration) ? m_pFitnessAlt : m_pFitness;
}

Matrix& EvolutionaryAlgorithm::getCurrFit()
{
  return ODD(m_iCurrentGeneration) ? m_pFitness : m_pFitnessAlt;
}

Natrix& EvolutionaryAlgorithm::getPrevBoard()
{
  return ODD(m_iCurrentGeneration) ? m_pLeaderboardAlt : m_pLeaderboard;
}

Natrix& EvolutionaryAlgorithm::getCurrBoard()
{
  return ODD(m_iCurrentGeneration) ? m_pLeaderboard : m_pLeaderboardAlt;
}

__global__ void init_spring_K(curandState* spring, bytes seed, nat totalInd)
{
  nat x = threadIdx.x + blockIdx.x*blockDim.x;

  if (x < totalInd)
    curand_init(seed + x, 0, 0, spring + x); //consider replacing seed with clock64()
}

void EvolutionaryAlgorithm::initRandomGenerator()
{
  cuda_call(cudaMalloc(&m_pRandSpring, m_iTotalIndividuals*sizeof(curandState)));

  AsyncKernelLaunch(init_spring_K, m_iBlocks, m_iThreadsPerBlock, m_pRandSpring, RDTSC(), m_iTotalIndividuals);
}

void EvolutionaryAlgorithm::validateSpeciesParameters() const
{
  if (m_iPopulationCount == 0 || m_iPopulationSize < 2 || (m_iPopulationSize & (m_iPopulationSize - 1)) || m_iGenomeSize == 0 || m_iMaxGenomeSize == 0 || m_iGenomeSize > m_iMaxGenomeSize)
    error("Invalid population parameters");
}
