/**
 * @file nsga2_impl.hpp
 * @author Sayan Goswami
 * @author Nanubala Gnana Sai
 *
 * Implementation of the NSGA-II algorithm. Used for multi-objective
 * optimization problems on arbitrary functions.
 *
 * ensmallen is free software; you may redistribute it and/or modify it under
 * the terms of the 3-clause BSD license.  You should have received a copy of
 * the 3-clause BSD license along with ensmallen.  If not, see
 * http://www.opensource.org/licenses/BSD-3-Clause for more Information.
 */

#ifndef ENSMALLEN_NSGA2_NSGA2_IMPL_HPP
#define ENSMALLEN_NSGA2_NSGA2_IMPL_HPP

#include "nsga2.hpp"
#include <assert.h>

namespace ens {

inline NSGA2::NSGA2(const size_t populationSize,
                    const size_t maxGenerations,
                    const double crossoverProb,
                    const double mutationProb,
                    const double mutationStrength,
                    const double epsilon,
                    const arma::vec& lowerBound,
                    const arma::vec& upperBound) :
    numObjectives(0),
    numVariables(0),
    populationSize(populationSize),
    maxGenerations(maxGenerations),
    crossoverProb(crossoverProb),
    mutationProb(mutationProb),
    mutationStrength(mutationStrength),
    epsilon(epsilon),
    lowerBound(lowerBound),
    upperBound(upperBound)
{ /* Nothing to do here. */ }

inline NSGA2::NSGA2(const size_t populationSize,
                    const size_t maxGenerations,
                    const double crossoverProb,
                    const double mutationProb,
                    const double mutationStrength,
                    const double epsilon,
                    const double lowerBound,
                    const double upperBound) :
    numObjectives(0),
    numVariables(0),
    populationSize(populationSize),
    maxGenerations(maxGenerations),
    crossoverProb(crossoverProb),
    mutationProb(mutationProb),
    mutationStrength(mutationStrength),
    epsilon(epsilon),
    lowerBound(lowerBound * arma::ones(1, 1)),
    upperBound(upperBound * arma::ones(1, 1))
{ /* Nothing to do here. */ }

//! Optimize the function.
template<typename MatType,
         typename... ArbitraryFunctionType,
         typename... CallbackTypes>
typename MatType::elem_type NSGA2::Optimize(
    std::tuple<ArbitraryFunctionType...>& objectives,
    MatType& iterateIn,
    CallbackTypes&&... callbacks)
{
  typedef typename ForwardType<MatType>::bcube CubeType;
  CubeType paretoFront, paretoSet;
  return Optimize(objectives, iterateIn, paretoFront, paretoSet,
      std::forward<CallbackTypes>(callbacks)...);
}

//! Optimize the function.
template<typename MatType,
         typename CubeType,
         typename... ArbitraryFunctionType,
         typename... CallbackTypes>
typename MatType::elem_type NSGA2::Optimize(
    std::tuple<ArbitraryFunctionType...>& objectives,
    MatType& iterateIn,
    CubeType& paretoFrontIn,
    CubeType& paretoSetIn,
    CallbackTypes&&... callbacks)
{
  // Make sure for evolution to work at least four candidates are present.
  if (populationSize < 4 && populationSize % 4 != 0)
  {
    throw std::logic_error("NSGA2::Optimize(): population size should be at"
        " least 4, and, a multiple of 4!");
  }

  // Convenience typedefs.
  typedef typename MatType::elem_type ElemType;
  typedef typename MatTypeTraits<MatType>::BaseMatType BaseMatType;
  typedef typename ForwardType<CubeType>::bmat CubeBaseMatType;

  BaseMatType& iterate = (BaseMatType&) iterateIn;

  // Make sure that we have the methods that we need.  Long name...
  traits::CheckArbitraryFunctionTypeAPI<ArbitraryFunctionType...,
      BaseMatType>();
  RequireDenseFloatingPointType<BaseMatType>();

  // Check if lower bound is a vector of a single dimension.
  if (lowerBound.n_rows == 1)
    lowerBound = lowerBound(0, 0) * arma::ones(iterate.n_rows, iterate.n_cols);

  // Check if upper bound is a vector of a single dimension.
  if (upperBound.n_rows == 1)
    upperBound = upperBound(0, 0) * arma::ones(iterate.n_rows, iterate.n_cols);

  // Check the dimensions of lowerBound and upperBound.
  assert(lowerBound.n_rows == iterate.n_rows && "The dimensions of "
      "lowerBound are not the same as the dimensions of iterate.");
  assert(upperBound.n_rows == iterate.n_rows && "The dimensions of "
      "upperBound are not the same as the dimensions of iterate.");

  numObjectives = sizeof...(ArbitraryFunctionType);
  numVariables = iterate.n_rows;

  // Cache calculated objectives as a matrix: (numObjectives x populationSize).
  arma::Mat<ElemType> calculatedObjectives(numObjectives, populationSize,
      arma::fill::zeros);

  // Population size reserved to 2 * populationSize + 1 to accommodate
  // for the size of intermediate candidate population.
  std::vector<BaseMatType> population;
  population.reserve(2 * populationSize + 1);

  // Pareto fronts, initialized during non-dominated sorting.
  // Stores indices of population belonging to a certain front.
  std::vector<std::vector<size_t> > fronts;
  // Initialised in CrowdingDistanceAssignment.
  std::vector<ElemType> crowdingDistance;
  // Initialised during non-dominated sorting.
  std::vector<size_t> ranks;

  //! Useful temporaries for float-like comparisons.
  const BaseMatType castedLowerBound = conv_to<BaseMatType>::from(
      lowerBound);
  const BaseMatType castedUpperBound = conv_to<BaseMatType>::from(
      upperBound);

  // Controls early termination of the optimization process.
  bool terminate = false;

  // Generate the population based on a uniform distribution around the given
  // starting point.
  for (size_t i = 0; i < populationSize; i++)
  {
    population.push_back(randu<BaseMatType>(iterate.n_rows,
        iterate.n_cols) - ElemType(0.5) + iterate);

    // Constrain all genes to be within bounds.
    population[i] = min(max(population[i], castedLowerBound), castedUpperBound);
  }

  Info << "NSGA2 initialized successfully. Optimization started." << std::endl;

  // Iterate until maximum number of generations is obtained.
  Callback::BeginOptimization(*this, objectives, iterate, callbacks...);

  for (size_t generation = 1; generation <= maxGenerations && !terminate;
      generation++)
  {
    Info << "NSGA2: iteration " << generation << "." << std::endl;

    // Create new population of candidate from the present elite population.
    // Have P_t, generate G_t using P_t.
    BinaryTournamentSelection(population, castedLowerBound, castedUpperBound);

    // Evaluate the objectives for the new population.
    calculatedObjectives.zeros(numObjectives, population.size());
    EvaluateObjectives(population, objectives, calculatedObjectives);

    // Perform fast non dominated sort on P_t ∪ G_t.
    ranks.resize(population.size());
    FastNonDominatedSort(fronts, ranks, calculatedObjectives);

    // Perform crowding distance assignment.
    crowdingDistance.resize(population.size());
    std::fill(crowdingDistance.begin(), crowdingDistance.end(), 0.);
    for (size_t fNum = 0; fNum < fronts.size(); fNum++)
    {
      CrowdingDistanceAssignment(fronts[fNum], calculatedObjectives,
          crowdingDistance);
    }

    // Sort based on crowding distance.
    std::sort(population.begin(), population.end(),
      [this, ranks, crowdingDistance, population]
        (BaseMatType candidateP, BaseMatType candidateQ)
          {
            size_t idxP{}, idxQ{};
            for (size_t i = 0; i < population.size(); i++)
            {
              if (approx_equal(population[i], candidateP, "absdiff",
                  ElemType(epsilon)))
                idxP = i;

              if (approx_equal(population[i], candidateQ, "absdiff",
                  ElemType(epsilon)))
                idxQ = i;
            }

            return CrowdingOperator<BaseMatType>(idxP, idxQ, ranks,
                crowdingDistance);
          }
    );

    // Yield a new population P_{t+1} of size populationSize.
    // Discards unfit population from the R_{t} to yield P_{t+1}.
    population.resize(populationSize);

    terminate |= Callback::GenerationalStepTaken(*this, objectives, iterate,
        calculatedObjectives, fronts, callbacks...);
  }

  // Set the candidates from the Pareto Set as the output.
  paretoSetIn.set_size(population[0].n_rows, population[0].n_cols,
      fronts[0].size());
  // The Pareto Set is stored, can be obtained via ParetoSet() getter.
  for (size_t solutionIdx = 0; solutionIdx < fronts[0].size(); ++solutionIdx)
  {
    paretoSetIn.slice(solutionIdx) = conv_to<CubeBaseMatType>::from(
        population[fronts[0][solutionIdx]]);
  }

  // Set the candidates from the Pareto Front as the output.
  paretoFrontIn.set_size(calculatedObjectives.n_rows, 1, fronts[0].size());
  for (size_t solutionIdx = 0; solutionIdx < fronts[0].size(); ++solutionIdx)
  {
    paretoFrontIn.slice(solutionIdx) = conv_to<CubeBaseMatType>::from(
        calculatedObjectives.col(fronts[0][solutionIdx]));
  }

  // Assign iterate to first element of the Pareto Set.
  iterate = population[fronts[0][0]];

  Callback::EndOptimization(*this, objectives, iterate, callbacks...);

  ElemType performance = std::numeric_limits<ElemType>::max();

  for (size_t i = 0; i < calculatedObjectives.n_cols; ++i)
    performance = std::min(performance, arma::accu(calculatedObjectives.col(i)));

  return performance;
}

//! No objectives to evaluate.
template<std::size_t I,
         typename MatType,
         typename ObjectiveMatType,
         typename ...ArbitraryFunctionType>
typename std::enable_if<I == sizeof...(ArbitraryFunctionType), void>::type
NSGA2::EvaluateObjectives(
    std::vector<MatType>&,
    std::tuple<ArbitraryFunctionType...>&,
    ObjectiveMatType&)
{
  // Nothing to do here.
}

//! Evaluate the objectives for the entire population.
template<std::size_t I,
         typename MatType,
         typename ObjectiveMatType,
         typename ...ArbitraryFunctionType>
typename std::enable_if<I < sizeof...(ArbitraryFunctionType), void>::type
NSGA2::EvaluateObjectives(
    std::vector<MatType>& population,
    std::tuple<ArbitraryFunctionType...>& objectives,
    ObjectiveMatType& calculatedObjectives)
{
  for (size_t i = 0; i < populationSize; i++)
  {
    calculatedObjectives(I, i) =
        std::get<I>(objectives).Evaluate(population[i]);
    EvaluateObjectives<I + 1, MatType, ObjectiveMatType,
        ArbitraryFunctionType...>(population, objectives, calculatedObjectives);
  }
}

//! Reproduce and generate new candidates.
template<typename InputMatType>
void NSGA2::BinaryTournamentSelection(
    std::vector<InputMatType>& population,
    const InputMatType& lowerBound,
    const InputMatType& upperBound)
{
  std::vector<InputMatType> children;

  while (children.size() < population.size())
  {
    // Choose two random parents for reproduction from the elite population.
    size_t indexA = arma::randi<size_t>(
        arma::distr_param(0, populationSize - 1));
    size_t indexB = arma::randi<size_t>(
        arma::distr_param(0, populationSize - 1));

    // Make sure that the parents differ.
    if (indexA == indexB)
    {
      if (indexB < populationSize - 1)
        indexB++;
      else
        indexB--;
    }

    // Initialize the children to the respective parents.
    InputMatType childA = population[indexA], childB = population[indexB];

    Crossover(childA, childB, population[indexA], population[indexB]);

    Mutate(childA, lowerBound, upperBound);
    Mutate(childB, lowerBound, upperBound);

    // Add the children to the candidate population.
    children.push_back(childA);
    children.push_back(childB);
  }

  // Add the candidates to the elite population.
  population.insert(
      std::end(population), std::begin(children), std::end(children));
}

//! Perform crossover of genes for the children.
template<typename InputMatType>
void NSGA2::Crossover(
    InputMatType& childA,
    InputMatType& childB,
    const InputMatType& parentA,
    const InputMatType& parentB)
{
  typedef typename InputMatType::elem_type ElemType;

  // Indices at which crossover is to occur.
  const InputMatType idx = conv_to<InputMatType>::from(randu<InputMatType>(
      childA.n_rows, childA.n_cols) < ElemType(crossoverProb));

  // Use traits from parentA for indices where idx is 1 and parentB otherwise.
  childA = parentA % idx + parentB % (1 - idx);
  // Use traits from parentB for indices where idx is 1 and parentA otherwise.
  childB = parentA % (1 - idx) + parentA % idx;
}

//! Perform mutation of the candidates weights with some noise.
template<typename InputMatType>
void NSGA2::Mutate(
    InputMatType& child,
    const InputMatType& lowerBound,
    const InputMatType& upperBound)
{
  typedef typename InputMatType::elem_type ElemType;

  child += conv_to<InputMatType>::from(
      InputMatType(child.n_rows, child.n_cols,
          GetFillType<InputMatType>::randu) < ElemType(mutationProb)) %
      (ElemType(mutationStrength) * InputMatType(child.n_rows, child.n_cols,
          GetFillType<InputMatType>::randn));

  // Constrain all genes to be between bounds.
  child = min(max(child, lowerBound), upperBound);
}

//! Sort population into Pareto fronts.
template<typename MatType>
void NSGA2::FastNonDominatedSort(
    std::vector<std::vector<size_t> >& fronts,
    std::vector<size_t>& ranks,
    MatType& calculatedObjectives)
{
  std::map<size_t, size_t> dominationCount;
  std::map<size_t, std::set<size_t> > dominated;

  // Reset and initialize fronts.
  fronts.clear();
  fronts.push_back(std::vector<size_t>());

  for (size_t p = 0; p < populationSize; p++)
  {
    dominated[p] = std::set<size_t>();
    dominationCount[p] = 0;

    for (size_t q = 0; q < populationSize; q++)
    {
      if (Dominates<MatType>(calculatedObjectives, p, q))
        dominated[p].insert(q);
      else if (Dominates<MatType>(calculatedObjectives, q, p))
        dominationCount[p] += 1;
    }

    if (dominationCount[p] == 0)
    {
      ranks[p] = 0;
      fronts[0].push_back(p);
    }
  }

  size_t i = 0;

  while (!fronts[i].empty())
  {
    std::vector<size_t> nextFront;

    for (size_t p: fronts[i])
    {
      for (size_t q: dominated[p])
      {
        dominationCount[q]--;

        if (dominationCount[q] == 0)
        {
          ranks[q] = i + 1;
          nextFront.push_back(q);
        }
      }
    }

    i++;
    fronts.push_back(nextFront);
  }
  // Remove the empty final set.
  fronts.pop_back();
}

//! Check if a candidate Pareto dominates another candidate.
template<typename MatType>
inline bool NSGA2::Dominates(
    MatType& calculatedObjectives,
    size_t candidateP,
    size_t candidateQ)
{
  bool allBetterOrEqual = true;
  bool atleastOneBetter = false;
  const size_t n_objectives = calculatedObjectives.n_rows;

  for (size_t i = 0; i < n_objectives; i++)
  {
    // P is worse than Q for the i-th objective function.
    if (calculatedObjectives(i, candidateP) >
        calculatedObjectives(i, candidateQ))
      allBetterOrEqual = false;

    // P is better than Q for the i-th objective function.
    else if (calculatedObjectives(i, candidateP) <
             calculatedObjectives(i, candidateQ))
      atleastOneBetter = true;
  }

  return allBetterOrEqual && atleastOneBetter;
}

//! Assign crowding distance to the population.
template <typename MatType>
void NSGA2::CrowdingDistanceAssignment(
    const std::vector<size_t>& front,
    MatType& calculatedObjectives,
    std::vector<typename MatType::elem_type>& crowdingDistance)
{
  // Convenience typedefs.
  typedef typename MatType::elem_type ElemType;
  typedef typename ForwardType<MatType>::uvec UVecType;
  typedef typename ForwardType<MatType>::bcol BaseColType;

  size_t fSize = front.size();
  // Stores the sorted indices of the fronts.
  UVecType sortedIdx = regspace<UVecType>(0, 1, fSize - 1);

  for (size_t m = 0; m < numObjectives; m++)
  {
    // Cache fValues of individuals for current objective.
    BaseColType fValues(fSize);
    for (size_t k = 0; k < fSize; ++k)
      fValues(k) = calculatedObjectives(m, size_t(front[k]));

    // Sort front indices by ascending fValues for current objective.
    sortedIdx = sort_index(fValues, "ascend");

    crowdingDistance[front[sortedIdx(0)]] =
        std::numeric_limits<ElemType>::max();
    crowdingDistance[front[sortedIdx(fSize - 1)]] =
        std::numeric_limits<ElemType>::max();
    ElemType minFval = fValues(sortedIdx(0));
    ElemType maxFval = fValues(sortedIdx(fSize - 1));
    ElemType scale =
        std::abs(maxFval - minFval) == 0 ? 1 : std::abs(maxFval - minFval);

    for (size_t i = 1; i < fSize - 1; i++)
    {
      crowdingDistance[front[sortedIdx(i)]] +=
          (fValues(sortedIdx(i + 1)) - fValues(sortedIdx(i - 1))) / scale;
    }
  }
}

//! Comparator for crowding distance based sorting.
template<typename InputMatType>
bool NSGA2::CrowdingOperator(
    size_t idxP,
    size_t idxQ,
    const std::vector<size_t>& ranks,
    const std::vector<typename InputMatType::elem_type>& crowdingDistance)
{
  if (ranks[idxP] < ranks[idxQ])
  {
    return true;
  }
  else if (ranks[idxP] == ranks[idxQ] &&
           crowdingDistance[idxP] > crowdingDistance[idxQ])
  {
    return true;
  }

  return false;
}

} // namespace ens

#endif
