/*=========================================================================
 *
 *  Copyright NumFOCUS
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *         http://www.apache.org/licenses/LICENSE-2.0.txt
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 *
 *=========================================================================*/

#ifndef itkParticleSwarmOptimizerBase_h
#define itkParticleSwarmOptimizerBase_h

#include "itkSingleValuedNonLinearOptimizer.h"
#include "itkMersenneTwisterRandomVariateGenerator.h"
#include "ITKOptimizersExport.h"

namespace itk
{
/** \class ParticleSwarmOptimizerBase
 * \brief Abstract implementation of a Particle Swarm Optimization (PSO) algorithm.
 *
 * The PSO algorithm was originally presented in:<br>
 * J. Kennedy, R. Eberhart, "Particle Swarm Optimization",
 * Proc. IEEE Int. Neural Networks, 1995.<br>
 *
 * The algorithm is a stochastic global search optimization approach.
 * Optimization is performed by maintaining a swarm of particles that traverse
 * the parameter space, searching for the optimal function value. Associated
 * with each particle are its location and speed, in parameter space.
 *
 * Swarm initialization is performed within the user supplied parameter bounds
 * using either a uniform distribution or a normal distribution centered on
 * the initial parameter values supplied by the user. The search terminates when
 * the maximal number of iterations has been reached or when the change in the
 * best value in the past \f$g\f$ generations is below a threshold and the swarm
 * has collapsed (i.e. best personal particle locations are close to the
 * swarm's best location in parameter space).
 *
 * The actual optimization procedure, updating the swarm, is performed in the
 * subclasses, required to implement the UpdateSwarm() method.
 *
 * NOTE: This implementation only performs minimization.
 *
 * \ingroup Numerics Optimizers
 * \ingroup ITKOptimizers
 */
class ITKOptimizers_EXPORT ParticleSwarmOptimizerBase : public SingleValuedNonLinearOptimizer
{
public:
  ITK_DISALLOW_COPY_AND_ASSIGN(ParticleSwarmOptimizerBase);

  /** Standard "Self" type alias. */
  using Self = ParticleSwarmOptimizerBase;
  using Superclass = SingleValuedNonLinearOptimizer;
  using Pointer = SmartPointer<Self>;
  using ConstPointer = SmartPointer<const Self>;

  /** Run-time type information (and related methods). */
  itkTypeMacro(ParticleSwarmOptimizerBase, SingleValuedNonLinearOptimizer)

    using ParameterBoundsType = std::vector<std::pair<ParametersType::ValueType, ParametersType::ValueType>>;

  struct ParticleData
  {
    ParametersType                m_CurrentParameters;
    ParametersType                m_CurrentVelocity;
    CostFunctionType::MeasureType m_CurrentValue;
    ParametersType                m_BestParameters;
    CostFunctionType::MeasureType m_BestValue;
  };

  using SwarmType = std::vector<ParticleData>;
  using NumberOfIterationsType = unsigned int;
  using NumberOfParticlesType = unsigned int;
  using NumberOfGenerationsType = unsigned int;
  using MeasureType = CostFunctionType::MeasureType;
  using ValueType = ParametersType::ValueType;
  using RandomVariateGeneratorType = Statistics::MersenneTwisterRandomVariateGenerator;

  /** Specify whether to initialize the particles using a normal distribution
   * centered on the user supplied initial value or a uniform distribution.
   * If the optimum is expected to be near the initial value it is likely
   * that initializing with a normal distribution will result in faster
   * convergence.*/
  itkSetMacro(InitializeNormalDistribution, bool) itkGetMacro(InitializeNormalDistribution, bool)
    itkBooleanMacro(InitializeNormalDistribution)

    /**
     * Specify the initial swarm. Useful for evaluating PSO variants. If the
     * initial swarm is set it will be used. To revert to random initialization
     * (uniform or normal particle distributions) set using an empty swarm.
     */
    void SetInitialSwarm(const SwarmType & initialSwarm);
  void
  ClearSwarm();

  /**
   * Indicate whether or not to output the swarm information when printing an
   * object. By default this option is turned off as it generates too much
   * information.
   */
  itkSetMacro(PrintSwarm, bool) itkGetMacro(PrintSwarm, bool) itkBooleanMacro(PrintSwarm)

    /** Start optimization. */
    void StartOptimization() override;


  /** Set/Get number of particles in the swarm - the maximal number of function
      evaluations is m_MaximalNumberOfIterations*m_NumberOfParticles */
  void
  SetNumberOfParticles(NumberOfParticlesType n);
  itkGetMacro(NumberOfParticles, NumberOfParticlesType)

    /** Set/Get maximal number of iterations - the maximal number of function
        evaluations is m_MaximalNumberOfIterations*m_NumberOfParticles */
    itkSetMacro(MaximalNumberOfIterations, NumberOfIterationsType)
      itkGetMacro(MaximalNumberOfIterations, NumberOfIterationsType)

    /** Set/Get the number of generations to continue with minimal improvement in
     *  the function value, |f_best(g_i) - f_best(g_k)|<threshold where
     *  k <= i+NumberOfGenerationsWithMinimalImprovement
     *  Minimal value is one.*/
    itkSetMacro(NumberOfGenerationsWithMinimalImprovement, NumberOfGenerationsType)
      itkGetMacro(NumberOfGenerationsWithMinimalImprovement, NumberOfGenerationsType)

    /**Set/Get the parameter bounds. Search for optimal value is inside these
       bounds. NOTE: It is assumed that the first entry is the minimal value,
       second is the maximal value. */
    virtual void SetParameterBounds(ParameterBoundsType & bounds);
  void
  SetParameterBounds(std::pair<ParametersType::ValueType, ParametersType::ValueType> & bounds, unsigned int n);

  ParameterBoundsType
  GetParameterBounds() const;

  /** The optimization algorithm will terminate when the function improvement
   *  in the last m_NumberOfGenerationsWithMinimalImprovement generations
   *  is less than m_FunctionConvergenceTolerance and the maximal distance
   *  between particles and the best particle in each dimension is less than
   *  m_ParametersConvergenceTolerance[i] for the specified percentage of the
   *  particles.
   *  That is, we haven't improved the best function value for a while and in
   *  the parameter space most (m%) of our particles attained their best value
   *  close to the swarm's best value.
   *  NOTE: The use of different tolerances for each dimension is desired when
   *         optimizing over non-commensurate parameters (e.g. rotation and
   *         translation). Alternatively, we could use ITK's parameter scaling
   *         approach. The current approach seems more intuitive.
   */
  itkSetMacro(FunctionConvergenceTolerance, MeasureType) itkGetMacro(FunctionConvergenceTolerance, MeasureType)
    /**Set parameters convergence tolerance using the same value for all, sz,
       parameters*/
    void SetParametersConvergenceTolerance(ValueType convergenceTolerance, unsigned int sz);
  itkSetMacro(ParametersConvergenceTolerance, ParametersType)
    itkGetMacro(ParametersConvergenceTolerance, ParametersType) itkGetMacro(PercentageParticlesConverged, double)
      itkSetMacro(PercentageParticlesConverged, double)

    /**Set the random number seed for the swarm. Use this method to
     * produce reaptible results, typically, for testing.
     */
    itkSetMacro(Seed, RandomVariateGeneratorType::IntegerType)
      itkGetMacro(Seed, RandomVariateGeneratorType::IntegerType)

    /** Use a specific seed to initialize the random number
     * generator. If On, use m_Seed to seed the random number
     * generator. Default is Off. */
    itkSetMacro(UseSeed, bool) itkGetMacro(UseSeed, bool) itkBooleanMacro(UseSeed)

    /** Get the function value for the current position.
     *  NOTE: This value is only valid during and after the execution of the
     *        StartOptimization() method.*/
    MeasureType GetValue() const;

  /** Get the reason for termination */
  const std::string
  GetStopConditionDescription() const override;

  /** Print the swarm information to the given output stream. Each line
   * (particle data) is of the form:
   * current_parameters current_velocity current_value best_parameters best_value
   */
  void
  PrintSwarm(std::ostream & os, Indent indent) const;

protected:
  ParticleSwarmOptimizerBase();
  ~ParticleSwarmOptimizerBase() override;
  void
  PrintSelf(std::ostream & os, Indent indent) const override;
  void
  PrintParamtersType(const ParametersType & x, std::ostream & os) const;

  /**
   * Implement your update rule in this function.*/
  virtual void
  UpdateSwarm() = 0;

  virtual void
  ValidateSettings();

  /**
   * Initialize the particle swarm, and seed the random number generator.
   */
  virtual void
  Initialize();

  void
  RandomInitialization();
  void
  FileInitialization();

  bool                                    m_PrintSwarm;
  std::ostringstream                      m_StopConditionDescription;
  bool                                    m_InitializeNormalDistribution;
  NumberOfParticlesType                   m_NumberOfParticles;
  NumberOfIterationsType                  m_MaximalNumberOfIterations;
  NumberOfGenerationsType                 m_NumberOfGenerationsWithMinimalImprovement;
  ParameterBoundsType                     m_ParameterBounds;
  ParametersType                          m_ParametersConvergenceTolerance;
  double                                  m_PercentageParticlesConverged;
  CostFunctionType::MeasureType           m_FunctionConvergenceTolerance;
  std::vector<ParticleData>               m_Particles;
  CostFunctionType::MeasureType           m_FunctionBestValue{ 0 };
  std::vector<MeasureType>                m_FunctionBestValueMemory;
  ParametersType                          m_ParametersBestValue;
  NumberOfIterationsType                  m_IterationIndex{ 0 };
  RandomVariateGeneratorType::IntegerType m_Seed;
  bool                                    m_UseSeed;
};
} // end namespace itk

#endif
