/**
 * @file agemoea.hpp
 * @author Satyam Shukla
 *
 * AGE-MOEA is a multi-objective optimization algorithm, widely used in
 * many real-world applications. AGE-MOEA generates offsprings using
 * crossover and mutation and then selects the next generation according
 * to non-dominated-sorting and survival score comparison.
 *
 * 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_AGEMOEA_AGEMOEA_HPP
#define ENSMALLEN_AGEMOEA_AGEMOEA_HPP

namespace ens {

/**
 * This class implements the AGEMOEA algorithm.
 *
 * The algorithm works by generating a candidate population from a fixed
 * starting point. At each stage of optimization, a new population of children
 * is generated. This new population along with its predecessor is sorted using
 * non-domination as the metric. Following this, the population is further
 * segregated in fronts. A new population is generated from these fronts having
 * size equal to that of the starting population.
 *
 * During evolution, two parents are randomly chosen using binary tournament
 * selection. A pair of children are generated by crossing over these two
 * candidates followed by mutation.
 *
 * The best front (Pareto optimal) is returned by the Optimize() method.
 *
 * For more information, see the following:
 *
 * @code
 * @inproceedings{panichella2019adaptive,
 * title={An adaptive evolutionary algorithm based on non-euclidean geometry for many-objective optimization},
 * author={Panichella, Annibale},
 * booktitle={Proceedings of the genetic and evolutionary computation conference},
 * pages={595--603},
 * year={2019}
 * }
 * @endcode
 *
 */
class AGEMOEA
{
 public:
  /**
   * Constructor for the AGE-MOEA optimizer.
   *
   * The default values provided over here are not necessarily suitable for a
   * given function. Therefore it is highly recommended to adjust the
   * parameters according to the problem.
   *
   * @param populationSize The number of candidates in the population.
   *     This should be atleast 4 in size and a multiple of 4.
   * @param maxGenerations The maximum number of generations allowed for NSGA-II.
   * @param crossoverProb The probability that a crossover will occur.
   * @param distributionIndex The crowding degree of the mutation.
   * @param epsilon The minimum difference required to distinguish between
   *     candidate solutions.
   * @param eta The distance parameters of the crossover distribution.
   * @param lowerBound Lower bound of the coordinates of the initial population.
   * @param upperBound Upper bound of the coordinates of the initial population.
   */
  AGEMOEA(const size_t populationSize = 100,
          const size_t maxGenerations = 2000,
          const double crossoverProb = 0.6,
          const double distributionIndex = 20,
          const double epsilon = 1e-6,
          const double eta = 20,
          const arma::vec& lowerBound = arma::zeros(1, 1),
          const arma::vec& upperBound = arma::ones(1, 1));

  /**
   * Constructor for the AGE-MOEA optimizer. This constructor provides an overload
   * to use `lowerBound` and `upperBound` of type double.
   *
   * The default values provided over here are not necessarily suitable for a
   * given function. Therefore it is highly recommended to adjust the
   * parameters according to the problem.
   *
   * @param populationSize The number of candidates in the population.
   *     This should be atleast 4 in size and a multiple of 4.
   * @param maxGenerations The maximum number of generations allowed for NSGA-II.
   * @param crossoverProb The probability that a crossover will occur.
   * @param distributionIndex The crowding degree of the mutation.
   * @param epsilon The minimum difference required to distinguish between
   *     candidate solutions.
   * @param eta The distance parameters of the crossover distribution
   * @param lowerBound Lower bound of the coordinates of the initial population.
   * @param upperBound Upper bound of the coordinates of the initial population.
   */
  AGEMOEA(const size_t populationSize = 100,
          const size_t maxGenerations = 2000,
          const double crossoverProb = 0.6,
          const double distributionIndex = 20,
          const double epsilon = 1e-6,
          const double eta = 20,
          const double lowerBound = 0,
          const double upperBound = 1);

  /**
   * Optimize a set of objectives. The initial population is generated using the
   * starting point. The output is the best generated front.
   *
   * @tparam ArbitraryFunctionType std::tuple of multiple objectives.
   * @tparam MatType Type of matrix to optimize.
   * @tparam CallbackTypes Types of callback functions.
   * @param objectives Vector of objective functions to optimize for.
   * @param iterate Starting point.
   * @param callbacks Callback functions.
   * @return MatType::elem_type The minimum of the accumulated sum over the
   *     objective values in the best front.
   */
  template<typename MatType,
           typename... ArbitraryFunctionType,
           typename... CallbackTypes>
  typename MatType::elem_type Optimize(
      std::tuple<ArbitraryFunctionType...>& objectives,
      MatType& iterate,
      CallbackTypes&&... callbacks);

  /**
   * Optimize a set of objectives. The initial population is generated using the
   * starting point. The output is the best generated front.
   *
   * @tparam ArbitraryFunctionType std::tuple of multiple objectives.
   * @tparam MatType Type of matrix to optimize.
   * @tparam CubeType The type of cube used to store the front and Pareto set.
   * @tparam CallbackTypes Types of callback functions.
   * @param objectives Vector of objective functions to optimize for.
   * @param iterate Starting point.
   * @param front The generated front.
   * @param paretoSet The generated Pareto set.
   * @param callbacks Callback functions.
   * @return MatType::elem_type The minimum of the accumulated sum over the
   *     objective values in the best front.
   */
  template<typename MatType,
           typename CubeType,
           typename... ArbitraryFunctionType,
           typename... CallbackTypes>
  typename MatType::elem_type Optimize(
      std::tuple<ArbitraryFunctionType...>& objectives,
      MatType& iterate,
      CubeType& front,
      CubeType& paretoSet,
      CallbackTypes&&... callbacks);

  //! Get the population size.
  size_t PopulationSize() const { return populationSize; }
  //! Modify the population size.
  size_t& PopulationSize() { return populationSize; }

  //! Get the maximum number of generations.
  size_t MaxGenerations() const { return maxGenerations; }
  //! Modify the maximum number of generations.
  size_t& MaxGenerations() { return maxGenerations; }

  //! Get the crossover rate.
  double CrossoverRate() const { return crossoverProb; }
  //! Modify the crossover rate.
  double& CrossoverRate() { return crossoverProb; }

  //! Retrieve value of the distribution index.
  double DistributionIndex() const { return distributionIndex; }
  //! Modify the value of the distribution index.
  double& DistributionIndex() { return distributionIndex; }

  //! Retrieve value of eta.
  double Eta() const { return eta; }
  //! Modify the value of eta.
  double& Eta() { return eta; }

  //! Get the tolerance.
  double Epsilon() const { return epsilon; }
  //! Modify the tolerance.
  double& Epsilon() { return epsilon; }

  //! Retrieve value of lowerBound.
  const arma::vec& LowerBound() const { return lowerBound; }
  //! Modify value of lowerBound.
  arma::vec& LowerBound() { return lowerBound; }

  //! Retrieve value of upperBound.
  const arma::vec& UpperBound() const { return upperBound; }
  //! Modify value of upperBound.
  arma::vec& UpperBound() { return upperBound; }

 private:
  /**
   * Evaluate objectives for the elite population.
   *
   * @tparam ArbitraryFunctionType std::tuple of multiple function types.
   * @tparam MatType Type of matrix to optimize.
   * @param population The elite population.
   * @param objectives The set of objectives.
   * @param calculatedObjectives Vector to store calculated objectives.
   */
  template<std::size_t I = 0,
           typename InputMatType,
           typename ObjectiveMatType,
           typename ...ArbitraryFunctionType>
  typename std::enable_if<I == sizeof...(ArbitraryFunctionType), void>::type
  EvaluateObjectives(std::vector<InputMatType>&,
                     std::tuple<ArbitraryFunctionType...>&,
                     std::vector<ObjectiveMatType>&);

  template<std::size_t I = 0,
           typename InputMatType,
           typename ObjectiveMatType,
           typename ...ArbitraryFunctionType>
  typename std::enable_if<I < sizeof...(ArbitraryFunctionType), void>::type
  EvaluateObjectives(std::vector<InputMatType>& population,
                     std::tuple<ArbitraryFunctionType...>& objectives,
                     std::vector<ObjectiveMatType>& calculatedObjectives);

  /**
   * Reproduce candidates from the elite population to generate a new
   * population.
   *
   * @tparam MatType Type of matrix to optimize.
   * @param objectives The set of objectives.
   * @param lowerBound Lower bound of the coordinates of the initial population.
   * @param upperBound Upper bound of the coordinates of the initial population.
   */
  template<typename MatType>
  void BinaryTournamentSelection(std::vector<MatType>& population,
                                 const MatType& lowerBound,
                                 const MatType& upperBound);

  /**
   * Crossover two parents to create a pair of new children.
   *
   * @tparam MatType Type of matrix to optimize.
   * @param childA A newly generated candidate.
   * @param childB Another newly generated candidate.
   * @param parentA First parent from elite population.
   * @param parentB Second parent from elite population.
   * @param lowerBound The lower bound of the objectives.
   * @param upperBound The upper bound of the objectives.
   */
  template<typename MatType>
  void Crossover(MatType& childA,
                 MatType& childB,
                 const MatType& parentA,
                 const MatType& parentB,
                 const MatType& lowerBound,
                 const MatType& upperBound);

  /**
   * Mutate the coordinates for a candidate.
   *
   * @tparam MatType Type of matrix to optimize.
   * @param candidate The candidate whose coordinates are being modified.
   * @param mutationRate The probablity of a mutation to occur.
   * @param lowerBound Lower bound of the coordinates of the initial population.
   * @param upperBound Upper bound of the coordinates of the initial population.
   */
  template<typename MatType>
  void Mutate(MatType& candidate,
              double mutationRate,
              const MatType& lowerBound,
              const MatType& upperBound);

  /**
   * Sort the candidate population using their domination count and the set of
   * dominated nodes.
   *
   * @tparam MatType Type of matrix to optimize.
   * @param fronts The population is sorted into these Pareto fronts. The first
   *     front is the best, the second worse and so on.
   * @param ranks The assigned ranks, used for crowding distance based sorting.
   * @param calculatedObjectives The previously calculated objectives.
   */
  template<typename MatType>
  void FastNonDominatedSort(
      std::vector<std::vector<size_t> >& fronts,
      std::vector<size_t>& ranks,
      std::vector<arma::Col<typename MatType::elem_type> >&
          calculatedObjectives);

  /**
   * Operator to check if one candidate Pareto-dominates the other.
   *
   * A candidate is said to dominate the other if it is at least as good as the
   * other candidate for all the objectives and there exists at least one
   * objective for which it is strictly better than the other candidate.
   *
   * @tparam MatType Type of matrix to optimize.
   * @param calculatedObjectives The previously calculated objectives.
   * @param candidateP The candidate being compared from the elite population.
   * @param candidateQ The candidate being compared against.
   * @return true if candidateP Pareto dominates candidateQ, otherwise, false.
   */
  template<typename MatType>
  bool Dominates(
      std::vector<arma::Col<typename MatType::elem_type>>& calculatedObjectives,
      size_t candidateP,
      size_t candidateQ);

  /**
   * Assigns Survival Score metric for sorting.
   *
   * @param front The previously generated Pareto fronts.
   * @param idealPoint The ideal point of teh first front.
   * @param calculatedObjectives The previously calculated objectives.
   * @param survivalScore The Survival Score vector to be updated for each
   *     individual in the population.
   * @param normalize The normlization vector of the fronts.
   * @param dimension The dimension of the first front.
   * @param fNum teh current front index.
   */
  template <typename MatType>
  void SurvivalScoreAssignment(
      const std::vector<size_t>& front,
      const arma::Col<typename MatType::elem_type>& idealPoint,
      std::vector<arma::Col<typename MatType::elem_type>>& calculatedObjectives,
      std::vector<typename MatType::elem_type>& survivalScore,
      arma::Col<typename MatType::elem_type>& normalize,
      typename MatType::elem_type& dimension,
      size_t fNum);

  /**
   * The operator used in the AGE-MOEA survival score based sorting.
   *
   * If a candidate has a lower rank then it is preferred.
   * Otherwise, if the ranks are equal then the candidate with the larger
   * Survival Score is preferred.
   *
   * @param idxP The index of the first cadidate from the elite population being
   *     sorted.
   * @param idxQ The index of the second cadidate from the elite population
   *     being sorted.
   * @param ranks The previously calculated ranks.
   * @param survivalScore The Survival score for each individual in
   *     the population.
   * @return true if the first candidate is preferred, otherwise, false.
   */
  template<typename MatType>
  bool SurvivalScoreOperator(
      size_t idxP,
      size_t idxQ,
      const std::vector<size_t>& ranks,
      const std::vector<typename MatType::elem_type>& survivalScore);

  /**
   * Normalizes the front given the extreme points in the current front.
   *
   * @tparam The type of population datapoints.
   * @param calculatedObjectives The current population evaluated objectives.
   * @param normalization The normalizing vector.
   * @param front The previously generated Pareto front.
   * @param extreme The indexes of the extreme points in the front.
   */
  template <typename MatType>
  void NormalizeFront(
      std::vector<arma::Col<typename MatType::elem_type>>& calculatedObjectives,
      arma::Col<typename MatType::elem_type>& normalization,
      const std::vector<size_t>& front,
      const arma::Row<size_t>& extreme);

  /**
   * Get the geometry information p of Lp norm (p > 0).
   *
   * @param calculatedObjectives The current population evaluated objectives.
   * @param front The previously generated Pareto fronts.
   * @param extreme The indexes of the extreme points in the front.
   * @return The variable p in the Lp norm that best fits the geometry of the
   *     current front.
   */
  template <typename MatType>
  typename MatType::elem_type GetGeometry(
      std::vector<arma::Col<typename MatType::elem_type> >&
          calculatedObjectives,
      const std::vector<size_t>& front,
      const arma::Row<size_t>& extreme);

  /**
   * Finds the pairwise Lp distance between all the points in the front.
   *
   * @param final The current population evaluated objectives.
   * @param calculatedObjectives The current population evaluated objectives.
   * @param front The front of the current generation.
   * @param dimension The calculated dimension of the front.
   */
  template <typename MatType>
  void PairwiseDistance(
      MatType& final,
      std::vector<arma::Col<typename MatType::elem_type> >&
          calculatedObjectives,
      const std::vector<size_t>& front,
      const typename MatType::elem_type dimension);

  /**
   * Finding the indexes of the extreme points in the front.
   *
   * @param indexes vector containing the slected indexes.
   * @param calculatedObjectives The current population objectives.
   * @param front The front of the current generation.
   */
  template <typename MatType>
  void FindExtremePoints(
      arma::Row<size_t>& indexes,
      std::vector<arma::Col<typename MatType::elem_type> >& calculatedObjectives,
      const std::vector<size_t>& front);

  /**
   * Finding the distance of each point in the front from the line formed
   * by pointA and pointB.
   *
   * @param distance The vector containing the distances of the points in the
   *    from from the line.
   * @param calculatedObjectives Reference to the current population evaluated
   *    objectives.
   * @param front The front of the current generation(indices of population).
   * @param pointA The first point on the line.
   * @param pointB The second point on the line.
   */
  template <typename MatType>
  void PointToLineDistance(
      arma::Row<typename MatType::elem_type>& distances,
      std::vector<arma::Col<typename MatType::elem_type> >&
          calculatedObjectives,
      const std::vector<size_t>& front,
      const arma::Col<typename MatType::elem_type>& pointA,
      const arma::Col<typename MatType::elem_type>& pointB);

  /**
   * Find the Diversity score corresponding the solution S using the selected
   * set.
   *
   * @param selected The current selected set.
   * @param pairwiseDistance The current pairwise distance for the whole front.
   * @param S The relative index of S being considered within the front.
   * @return The diversity score for S which the sum of the two smallest
   *     elements.
  */
 template <typename MatType>
 typename MatType::elem_type DiversityScore(std::set<size_t>& selected,
                                            const MatType& pairwiseDistance,
                                            size_t S);

  //! The number of objectives being optimised for.
  size_t numObjectives;

  //! The numbeer of variables used per objectives.
  size_t numVariables;

  //! The number of candidates in the population.
  size_t populationSize;

  //! Maximum number of generations before termination criteria is met.
  size_t maxGenerations;

  //! Probability that crossover will occur.
  double crossoverProb;

  //! The crowding degree of the mutation. Higher value produces a mutant
  //! resembling its parent.
  double distributionIndex;

  //! The tolerance for termination.
  double epsilon;

  //! The distance parameters of the crossover distribution.
  double eta;

  //! Lower bound of the initial swarm.
  arma::vec lowerBound;

  //! Upper bound of the initial swarm.
  arma::vec upperBound;
};

} // namespace ens

// Include implementation.
#include "agemoea_impl.hpp"

#endif
