/* ========================================================================== *
 *  Multi-population genetic algorithm.
 *
 *  Author: Firas Safadi (fsafadi@ulg.ac.be)
 * ========================================================================== */

#include "GeneticAlgorithm.h"

#define MAX_STOCHASTIC_ACCEPTANCE_TRIES 20

#ifndef DOUBLE_PRECISION
#define EXP10 exp10f
#define CURAND_UNIFORM curand_uniform
#else
#define EXP10 exp10
#define CURAND_UNIFORM curand_uniform_double
#endif

GeneticAlgorithm::GeneticAlgorithm(nat popCount, nat popSize, nat genSize, nat maxGenSize) :
EvolutionaryAlgorithm(popCount, popSize, genSize, maxGenSize)
{
  m_pInfBoundVec.resize(1, popCount);
  m_pSupBoundVec.resize(1, popCount);
  m_pExpInfBoundVec.resize(1, popCount);
  m_pExpSupBoundVec.resize(1, popCount);
  m_pEliteProportionVec.resize(1, popCount);
  m_pBreedProportionVec.resize(1, popCount);
  m_pMutantProportionVec.resize(1, popCount);
  m_pEliteVec.resize(1, popCount);
  m_pBreedVec.resize(1, popCount);
  m_pMutantVec.resize(1, popCount);
  m_pXoTypeVec.resize(1, popCount);
  m_pMutationRateVec.resize(1, popCount);
  m_pMutationGammaVec.resize(1, popCount);

  m_pInfBoundVec.spawn();
  m_pSupBoundVec.spawn();
  m_pExpInfBoundVec.spawn();
  m_pExpSupBoundVec.spawn();
  m_pEliteProportionVec.spawn();
  m_pBreedProportionVec.spawn();
  m_pMutantProportionVec.spawn();
  m_pEliteVec.spawn();
  m_pBreedVec.spawn();
  m_pMutantVec.spawn();
  m_pXoTypeVec.spawn();
  m_pMutationRateVec.spawn();
  m_pMutationGammaVec.spawn();

  m_eXoType = Random;
}

GeneticAlgorithm::~GeneticAlgorithm()
{

}

__device__ inline void d_gen_ind(curandState* spring, real* pop, nat popLdt, nat genSize, real infBound, real supBound, real expInfBound, real expSupBound)
{
  curandState localState = *spring;

  real* end = pop + genSize*popLdt;

  while (pop < end) {
    *pop = ((supBound - infBound)*CURAND_UNIFORM(&localState) + infBound) * EXP10((expSupBound - expInfBound)*CURAND_UNIFORM(&localState) + expInfBound);
    pop += popLdt;
  }

  *spring = localState;
}

__global__ void gen_pop_K(curandState* spring, real* pop, nat popLdt, nat totalInd, nat popSize, nat genSize, real* infBound, real* supBound, real* expInfBound, real* expSupBound)
{
  nat x = threadIdx.x + blockIdx.x*blockDim.x;

  if (x < totalInd) {
    nat i = x / popSize;

    real iBound = infBound[i];
    real sBound = supBound[i];
    real eIBound = expInfBound[i];
    real eSBound = expSupBound[i];

    d_gen_ind(spring + x, pop + x, popLdt, genSize, iBound, sBound, eIBound, eSBound);
  }
}

void GeneticAlgorithm::hatchSpecies()
{
  if (m_bAlive)
    error("Cannot hatch population: already alive");

  validateGenetics();

  AsyncKernelLaunch(gen_pop_K, m_iBlocks, m_iThreadsPerBlock,
    m_pRandSpring, m_pPopulation.dp(), m_pPopulation.ldt(),
    m_iTotalIndividuals, m_iPopulationSize, m_bPreGenerate ? m_iMaxGenomeSize : m_iGenomeSize,
    m_pInfBoundVec.dp(), m_pSupBoundVec.dp(), m_pExpInfBoundVec.dp(), m_pExpSupBoundVec.dp());

  m_iCurrentGeneration = 1;

  m_bAlive = true;

  evalPop();

  sortPop();
}

__device__ inline void d_elite_evo(real* prevPop, nat prevPopLdt, real* currPop, nat currPopLdt, nat genSize)
{
  real* end = prevPop + genSize*prevPopLdt;

  while (prevPop < end) {
    *prevPop = *currPop;
    prevPop += prevPopLdt;
    currPop += currPopLdt;
  }
}

__device__ inline nat d_select_ind(curandState* spring, real* currFit, nat* currBoard, nat popSize)
{
  popSize--;

  real maxFit = currFit[0];
  nat triesLeft = MAX_STOCHASTIC_ACCEPTANCE_TRIES;

  for (;;) {
    nat i = curand(spring) & popSize;

    if (--triesLeft == 0 || CURAND_UNIFORM(spring) <= currFit[i] / maxFit)
      return currBoard[i];
  }
}

__device__ inline void d_1pxo_evo(curandState* spring, real* prevPop, nat prevPopLdt, real* currPop, nat currPopLdt, real* currFit, nat* currBoard, nat popSize, nat genSize)
{
  curandState localState = *spring;

  nat a = d_select_ind(&localState, currFit, currBoard, popSize);
  nat b = d_select_ind(&localState, currFit, currBoard, popSize);

  nat xop = curand(&localState) % (genSize - 1) + 1;

  real* curr = currPop + a;
  real* end = prevPop + xop*prevPopLdt;

  while (prevPop < end) {
    *prevPop = *curr;
    prevPop += prevPopLdt;
    curr += currPopLdt;
  }

  curr = currPop + xop*currPopLdt + b;
  end = prevPop + (genSize - xop)*prevPopLdt;

  while (prevPop < end) {
    *prevPop = *curr;
    prevPop += prevPopLdt;
    curr += currPopLdt;
  }

  *spring = localState;
}

__device__ inline void d_uxo_evo(curandState* spring, real* prevPop, nat prevPopLdt, real* currPop, nat currPopLdt, real* currFit, nat* currBoard, nat popSize, nat genSize)
{
  curandState localState = *spring;

  nat a = d_select_ind(&localState, currFit, currBoard, popSize);
  nat b = d_select_ind(&localState, currFit, currBoard, popSize);

  real* curr_a = currPop + a;
  real* curr_b = currPop + b;
  real* end = prevPop + genSize*prevPopLdt;

  while (prevPop < end) {
    real draw = CURAND_UNIFORM(&localState);

    *prevPop = draw > 0.5 ? *curr_b : *curr_a;
    prevPop += prevPopLdt;
    curr_a += currPopLdt;
    curr_b += currPopLdt;
  }

  *spring = localState;
}

__device__ inline void d_mutant_evo(curandState* spring, real* prevPop, nat prevPopLdt, real* currPop, nat currPopLdt, real* currFit, nat* currBoard, nat popSize, nat genSize, real mutationRate, real mutationGamma)
{
  curandState localState = *spring;

  currPop += d_select_ind(&localState, currFit, currBoard, popSize);

  real* end = prevPop + genSize*prevPopLdt;

  while (prevPop < end) {
    real draw = CURAND_UNIFORM(&localState);

    *prevPop = draw > mutationRate ? *currPop : *currPop*(1.0 + mutationGamma*(2.0*CURAND_UNIFORM(&localState) - 1.0));
    prevPop += prevPopLdt;
    currPop += currPopLdt;
  }

  *spring = localState;
}

__global__ void evo_pop_K(curandState* spring, real* prevPop, nat prevPopLdt, real* currPop, nat currPopLdt, real* currFit, nat* currBoard, nat totalInd, nat popSize, nat genSize, nat maxGenSize, nat* eliteCount, nat* breedCount, nat* mutantCount, nat* xoType, real* mutationRate, real* mutationGamma, real* infBound, real* supBound, real* expInfBound, real* expSupBound)
{
  nat x = threadIdx.x + blockIdx.x*blockDim.x;

  if (x < totalInd) {
    nat i = x / popSize;
    
    nat eliteBp = eliteCount[i];
    nat breedBp = eliteBp + breedCount[i];
    nat mutantBp = breedBp + mutantCount[i];

    nat xType = xoType[i];
    real mRate = mutationRate[i];
    real mGamma = mutationGamma[i];

    real iBound = infBound[i];
    real sBound = supBound[i];
    real eIBound = expInfBound[i];
    real eSBound = expSupBound[i];

    nat offset = x & ~(popSize - 1);

    spring += offset;
    prevPop += offset;
    currPop += offset;
    currFit += offset;
    currBoard += offset;

    x &= popSize - 1;

    if (x < eliteBp) {
      d_elite_evo(prevPop + x, prevPopLdt, currPop + currBoard[x], currPopLdt, genSize);
    }

    else if (x < breedBp) {
      if (xType == GeneticAlgorithm::OnePoint)
        d_1pxo_evo(spring + x, prevPop + x, prevPopLdt, currPop, currPopLdt, currFit, currBoard, popSize, genSize);
      else if (xType == GeneticAlgorithm::Uniform)
        d_uxo_evo(spring + x, prevPop + x, prevPopLdt, currPop, currPopLdt, currFit, currBoard, popSize, genSize);
      else if (xType == GeneticAlgorithm::Mixed) {
        if (ODD(x))
          d_uxo_evo(spring + x, prevPop + x, prevPopLdt, currPop, currPopLdt, currFit, currBoard, popSize, genSize);
        else
          d_1pxo_evo(spring + x, prevPop + x, prevPopLdt, currPop, currPopLdt, currFit, currBoard, popSize, genSize);
      }
    }

    else if (x < mutantBp) {
      d_mutant_evo(spring + x, prevPop + x, prevPopLdt, currPop, currPopLdt, currFit, currBoard, popSize, genSize, mRate, mGamma);
    }

    else if (x < popSize) {
      d_gen_ind(spring + x, prevPop + x, prevPopLdt, genSize, iBound, sBound, eIBound, eSBound);
    }
  }
}

void GeneticAlgorithm::evolveSpecies()
{
  if (!m_bAlive)
    error("Cannot evolve population: not alive");

  for (nat i = 0; i < m_iEvolutionEpochs; i++) {
    AsyncKernelLaunch(evo_pop_K, m_iBlocks, m_iThreadsPerBlock,
      m_pRandSpring,
      getPrevPop().dp(), getPrevPop().ldt(),
      getCurrPop().dp(), getCurrPop().ldt(),
      getCurrFit().dp(), getCurrBoard().dp(),
      m_iTotalIndividuals, m_iPopulationSize, m_iGenomeSize, m_iMaxGenomeSize,
      m_pEliteVec.dp(), m_pBreedVec.dp(), m_pMutantVec.dp(),
      m_pXoTypeVec.dp(), m_pMutationRateVec.dp(), m_pMutationGammaVec.dp(),
      m_pInfBoundVec.dp(), m_pSupBoundVec.dp(), m_pExpInfBoundVec.dp(), m_pExpSupBoundVec.dp());

    m_iCurrentGeneration++;

    evalPop();

    sortPop();
  }
}

void GeneticAlgorithm::hatchSpecies(nat genSize)
{
  if (genSize > m_iMaxGenomeSize)
    error("Genome size cannot be reset: value out of range");

  m_iGenomeSize = genSize;

  if (m_bAlive) {
    AsyncKernelLaunch(gen_pop_K, m_iBlocks, m_iThreadsPerBlock,
      m_pRandSpring, m_pPopulation.dp(), m_pPopulation.ldt(),
      m_iTotalIndividuals, m_iPopulationSize, m_bPreGenerate ? m_iMaxGenomeSize : m_iGenomeSize,
      m_pInfBoundVec.dp(), m_pSupBoundVec.dp(), m_pExpInfBoundVec.dp(), m_pExpSupBoundVec.dp());

    m_iCurrentGeneration = 1;

    evalPop();

    sortPop();
  }
  
  else {
    hatchSpecies();
  }
}

void GeneticAlgorithm::wipeSpecies()
{
  m_iCurrentGeneration = 0;
  m_bAlive = false;
}

void GeneticAlgorithm::setPopGenetics(nat i, real infBound, real supBound, real expInfBound, real expSupBound, real eliteProp, real breedProp, real mutantProp, real mutationRate, real mutationGamma, CrossoverType xoType)
{
  if (m_bAlive)
    error("Cannot change living population parameters");

  m_pInfBoundVec(0, i) = infBound;
  m_pSupBoundVec(0, i) = supBound;

  m_pExpInfBoundVec(0, i) = expInfBound;
  m_pExpSupBoundVec(0, i) = expSupBound;

  m_pEliteProportionVec(0, i) = eliteProp;
  m_pBreedProportionVec(0, i) = breedProp;
  m_pMutantProportionVec(0, i) = mutantProp;

  m_pXoTypeVec(0, i) = xoType == Random ? rand() % 3 : xoType;

  m_pMutationRateVec(0, i) = mutationRate;
  m_pMutationGammaVec(0, i) = mutationGamma;

  validatePopGenetics(i);

  m_pEliteVec(0, i) = static_cast<nat>(m_iPopulationSize*eliteProp);
  m_pBreedVec(0, i) = static_cast<nat>(m_iPopulationSize*breedProp);
  m_pMutantVec(0, i) = static_cast<nat>(m_iPopulationSize*mutantProp);

  pushSpeciesToDevice();
}

void GeneticAlgorithm::setSpeciesGenetics(
  real infBoundLow, real infBoundHigh, real supBoundLow, real supBoundHigh,
  real expInfBoundLow, real expInfBoundHigh, real expSupBoundLow, real expSupBoundHigh,
  real elitePropLow, real elitePropHigh, real breedPropLow, real breedPropHigh, real mutantPropLow, real mutantPropHigh,
  real mutationRateLow, real mutationRateHigh,
  real mutationGammaLow, real mutationGammaHigh,
  CrossoverType xoType)
{
  if (m_bAlive)
    error("Cannot change living population parameters");

  m_pInfBoundRange.set(infBoundLow, infBoundHigh);
  m_pSupBoundRange.set(supBoundLow, supBoundHigh);

  m_pExpInfBoundRange.set(expInfBoundLow, expInfBoundHigh);
  m_pExpSupBoundRange.set(expSupBoundLow, expSupBoundHigh);

  m_pEliteProportionRange.set(elitePropLow, elitePropHigh);
  m_pBreedProportionRange.set(breedPropLow, breedPropHigh);
  m_pMutantProportionRange.set(mutantPropLow, mutantPropHigh);

  m_pMutationRateRange.set(mutationRateLow, mutationRateHigh);
  m_pMutationGammaRange.set(mutationGammaLow, mutationGammaHigh);

  m_eXoType = xoType;

  validateGenetics();

  generateSpecies();

  pushSpeciesToDevice();
}

void GeneticAlgorithm::dispGenetics() const
{
  for (nat i = 0; i < m_iPopulationCount; i++)
    dispGenetics(i);
}

void GeneticAlgorithm::dispGenetics(nat i) const
{
  std::cout << std::endl;
  std::cout << "Base              :    [" << m_pInfBoundVec(0, i) << ", " << m_pSupBoundVec(0, i) << "]" << std::endl;

  std::cout << "Exponent          :    [" << m_pExpInfBoundVec(0, i) << ", " << m_pExpSupBoundVec(0, i) << "]" << std::endl;

  std::cout << "Elite %           :    " << m_pEliteProportionVec(0, i) << " (" << m_pEliteVec(0, i) << " p)" << std::endl;
  std::cout << "Breed %           :    " << m_pBreedProportionVec(0, i) << " (" << m_pBreedVec(0, i) << " p)" << std::endl;
  std::cout << "Mutant %          :    " << m_pMutantProportionVec(0, i) << " (" << m_pMutantVec(0, i) << " p)" << std::endl;
  std::cout << "Random %          :    " << 1.0 - (m_pEliteProportionVec(0, i) + m_pBreedProportionVec(0, i) + m_pMutantProportionVec(0, i)) << " (" << m_iPopulationSize - (m_pEliteVec(0, i) + m_pBreedVec(0, i) + m_pMutantVec(0, i)) << " p)" << std::endl;

  std::cout << "Crossover type    :    " << (m_pXoTypeVec(0, i) != OnePoint ? m_pXoTypeVec(0, i) != Uniform ? "Mixed" : "Uniform" : "OnePoint") << std::endl;

  std::cout << "Mutation rate     :    " << m_pMutationRateVec(0, i) << std::endl;
  std::cout << "Mutation gamma    :    " << m_pMutationGammaVec(0, i) << std::endl;
}

void GeneticAlgorithm::validateGenetics()
{
  if (m_pInfBoundRange.high > m_pSupBoundRange.low)
    error("Lower base bound cannot be greater than upper bound");

  if (m_pExpInfBoundRange.high > m_pExpSupBoundRange.low)
    error("Lower exponent bound cannot be greater than upper bound");

  if (m_pEliteProportionRange.low < 0.0 || m_pBreedProportionRange.low < 0.0 || m_pMutantProportionRange.low < 0.0)
    error("Genetic operator proportion cannot be lower than 0.0");

  if (m_pEliteProportionRange.high + m_pBreedProportionRange.high + m_pMutantProportionRange.high > 1.0)
    error("Sum of genetic operator proportions cannot be greater than 1.0");

  if (m_pMutationRateRange.low < 0.0 || m_pMutationRateRange.high > 1.0)
    error("Mutation rate cannot be lower than 0.0 or greater than 1.0");

  if (m_pMutationGammaRange.low < 0.0 || m_pMutationGammaRange.high > 1.0)
    error("Mutation gamma cannot be lower than 0.0 or greater than 1.0");

  if (m_iGenomeSize < 2 && m_pBreedProportionRange.high > 0.0)
    error("Crossover proportion cannot be greater than 0.0 with only 1 gene");
}

void GeneticAlgorithm::validatePopGenetics(nat i)
{
  if (m_pInfBoundVec(0, i) > m_pSupBoundVec(0, i))
    error("Lower base bound cannot be greater than upper bound");

  if (m_pExpInfBoundVec(0, i) > m_pExpSupBoundVec(0, i))
    error("Lower exponent bound cannot be greater than upper bound");

  if (m_pEliteProportionVec(0, i) < 0.0 || m_pBreedProportionVec(0, i) < 0.0 || m_pMutantProportionVec(0, i) < 0.0)
    error("Genetic operator proportion cannot be lower than 0.0");

  if (m_pEliteProportionVec(0, i) + m_pBreedProportionVec(0, i) + m_pMutantProportionVec(0, i) > 1.0)
    error("Sum of genetic operator proportions cannot be greater than 1.0");

  if (m_pMutationRateVec(0, i) < 0.0 || m_pMutationRateVec(0, i) > 1.0)
    error("Mutation rate cannot be lower than 0.0 or greater than 1.0");

  if (m_pMutationGammaVec(0, i) < 0.0 || m_pMutationGammaVec(0, i) > 1.0)
    error("Mutation gamma cannot be lower than 0.0 or greater than 1.0");

  if (m_iGenomeSize < 2 && m_pBreedProportionVec(0, i) > 0.0)
    error("Crossover proportion cannot be greater than 0.0 with only 1 gene");

  if (m_pXoTypeVec(0, i) >= Random)
    error("Invalid crossover type");
}

void GeneticAlgorithm::generateSpecies()
{
  for (nat i = 0; i < m_iPopulationCount; i++) {
    m_pInfBoundVec(0, i) = rand(m_pInfBoundRange.low, m_pInfBoundRange.high);
    m_pSupBoundVec(0, i) = rand(m_pSupBoundRange.low, m_pSupBoundRange.high);

    m_pExpInfBoundVec(0, i) = rand(m_pExpInfBoundRange.low, m_pExpInfBoundRange.high);
    m_pExpSupBoundVec(0, i) = rand(m_pExpSupBoundRange.low, m_pExpSupBoundRange.high);

    m_pEliteProportionVec(0, i) = rand(m_pEliteProportionRange.low, m_pEliteProportionRange.high);
    m_pBreedProportionVec(0, i) = rand(m_pBreedProportionRange.low, m_pBreedProportionRange.high);
    m_pMutantProportionVec(0, i) = rand(m_pMutantProportionRange.low, m_pMutantProportionRange.high);

    m_pMutationRateVec(0, i) = rand(m_pMutationRateRange.low, m_pMutationRateRange.high);
    m_pMutationGammaVec(0, i) = rand(m_pMutationGammaRange.low, m_pMutationGammaRange.high);

    m_pXoTypeVec(0, i) = m_eXoType == Random ? rand() % 3 : m_eXoType;

    m_pEliteVec(0, i) = static_cast<nat>(m_iPopulationSize*m_pEliteProportionVec(0, i));
    m_pBreedVec(0, i) = static_cast<nat>(m_iPopulationSize*m_pBreedProportionVec(0, i));
    m_pMutantVec(0, i) = static_cast<nat>(m_iPopulationSize*m_pMutantProportionVec(0, i));
  }
}

void GeneticAlgorithm::pushSpeciesToDevice()
{
  m_pInfBoundVec.syncDeviceWithHost();
  m_pSupBoundVec.syncDeviceWithHost();
  m_pExpInfBoundVec.syncDeviceWithHost();
  m_pExpSupBoundVec.syncDeviceWithHost();
  m_pEliteProportionVec.syncDeviceWithHost();
  m_pBreedProportionVec.syncDeviceWithHost();
  m_pMutantProportionVec.syncDeviceWithHost();
  m_pEliteVec.syncDeviceWithHost();
  m_pBreedVec.syncDeviceWithHost();
  m_pMutantVec.syncDeviceWithHost();
  m_pXoTypeVec.syncDeviceWithHost();
  m_pMutationRateVec.syncDeviceWithHost();
  m_pMutationGammaVec.syncDeviceWithHost();
}
