#include <globals.hh>
#include <Randomize.hh>

#include <RAT/OptimisedComponent.hh>
#include <RAT/SimulatedAnnealing.hh>
#include <RAT/DS/FitResult.hh>
#include <RAT/Processor.hh>
#include <RAT/Log.hh>
#include <CLHEP/Random/RandGauss.h>
using namespace RAT::DS;
using namespace RAT::Optimisers;

#include <iostream>
#include <cmath>
using namespace std;


void
SimulatedAnnealing::Initialise(const std::string& )
{
  // Defaults from SNOMAN path fitter
  fIterations = 150;
  fTolerance = 1.0e-4;
  fScheduleSteps = 10;
  fSchedulePower = 4;
  fStartTemperature = -9999; // -ve number means use default
}

void
SimulatedAnnealing::SetI( const std::string& param, int value )
{
  if( param == string( "iterations" ) )
    fIterations = value;
  else if( param == string( "scheduleSteps" ) )
    fScheduleSteps = value;
  else if( param == string( "schedulePower" ) )
    fSchedulePower = value;
  else
    throw Processor::ParamUnknown( param );

}

void
SimulatedAnnealing::SetD( const std::string& param, double value )
{
  if( param == string( "tolerance" ) )
    fTolerance = value;
  else if( param == string( "startingTemperature" ) )
    fStartTemperature = value;
  else
    throw Processor::ParamUnknown( param );
}

double
SimulatedAnnealing::Minimise()
{
  fMinFactor = 1.0;
  return AnnealOptimise();
}

double
SimulatedAnnealing::Maximise()
{
  fMinFactor = -1.0;
  return AnnealOptimise();
}

double
SimulatedAnnealing::AnnealOptimise()
{

  // Get initial parameters
  vector<double> params = fComponent->GetParams();
  vector<double> posErrors = fComponent->GetPositiveErrors();
  vector<double> negErrors = fComponent->GetNegativeErrors();

  // Need to get a set of seeds from e.g. the MetaNSeed
  const int ndim = params.size();
  const int ndim1 = ndim+1;

  vector<double> p(ndim1*ndim, 0); // p[1 .. ndim+1][1 .. ndim]
  vector<double> ptry(ndim, 0);
  vector<double> y(ndim1, 0);
  vector<double> pb(ndim, 0);
  vector<double> notnan;

  for(int i=0; i<ndim1; i++)
    {
      int itry = 0;
      do
        {
          for(int j=0; j<ndim; j++)
            {
              // Using the seed errors, generate points on the simplex
              ptry[j] = CLHEP::RandGauss::shoot(params[j], posErrors[j]);
              p[i*ndim + j] = ptry[j];
            }
          // Evaluate for y values
          y[i] = fMinFactor * (*fComponent)(ptry);

          itry++;
          if( itry>20 )
            {
              // Give up!
              debug << "SimulatedAnnealing: Keep creating seed with nan PDF value" << newline;
              break;
            }
        }while( std::isnan(y[i]) );

      if(!std::isnan(y[i]))
        notnan.push_back(y[i]);

    }

  if(notnan.size()==0)
    {
      debug << "SimulatedAnnealing: All values are nan!" << newline;
      throw OptimiserFailError( "SimulatedAnnealing: All values are nan!" );
    }

  double yworst = *max_element(notnan.begin(), notnan.end());
  for(int i=0;i<ndim1;i++)
    if(std::isnan(y[i]))
      // just set to the worst value plus a fluctuation
      y[i] = yworst + G4UniformRand();

  int ybIndex = distance( y.begin(), min_element(y.begin(), y.end()));
  double yb = y[ybIndex];

  for(int i=0;i<ndim;i++)
    pb[i] = p[ybIndex*ndim + i];

  int iterations = fIterations;
  double temperature = fStartTemperature;
  if(fStartTemperature<0)
    // Start with a temperature that is the range between ybest and yworst
    // Note that if temperature is 0 then all points on the simplex have the
    // same y value - impossible to optimise.
    temperature = yworst - yb;

  // Now call the annealing according to annealing schedule
  for(int i=0; i<fScheduleSteps; i++){
    double tempTemp = temperature * pow((1 - float(i) / fScheduleSteps), fSchedulePower);
    iterations = fIterations;
    Amebsa(p, y, ndim, pb, yb, fTolerance, iterations, tempTemp);
  }

  fParams = pb;
  fPositiveErrors = posErrors;
  fNegativeErrors = negErrors;

  // Not clear how to check validity of anneal method
  fValid = true;

  return yb;
}


void
SimulatedAnnealing::GetPSum(std::vector<double>& psum, const int& ndim, const std::vector<double>& p)
{
  int mpts = ndim+1;
  for (int i=0; i<ndim; i++)
    {
      double sum = 0.0;
      for(int j=0; j<mpts; j++)
        sum += p[j*ndim + i];
      psum[i] = sum;
    }
}


void
SimulatedAnnealing::Amebsa(std::vector<double>& p, std::vector<double>& y, int ndim, std::vector<double>& pb,
                           double& yb, double ftol, int& iter, double temptr)
{
  // Multidimensional minimization of the function funk(x) where x[1..ndim] is a vector in
  // ndim dimensions, by simulated annealing combined with the downhill simplex method of
  // Nelder and Mead. The input matrix p[1..ndim+1][1..ndim] has ndim+1 rows, each an
  // ndim- dimensional vector which is a vertex of the starting simplex. Also input are the
  // following: the vector y[1..ndim+1], whose components must be pre-initialized to the values
  // of funk evaluated at the ndim+1 vertices (rows) of p; ftol, the fractional convergence
  // tolerance to be achieved in the function value for an early return; iter, and temptr The
  // routine makes iter function evaluations at an annealing temperature temptr, then returns.
  // You should then decrease temptr according to your annealing schedule, reset iter, and call
  // the routine again (leaving other arguments unaltered between calls). If iter is returned
  // with a positive value, then early convergence and return occurred. If you initialize yb to
  // a very large value on the first call, then yb and pb[1..ndim] will subsequently return the
  // best function value and point ever encountered (even if it is no longer a point in the simplex).

  int mpts=ndim+1;
  vector<double> psum(ndim, 0);
  double tt = -temptr;

  GetPSum(psum, ndim, p);

  for(;;)
    {
      // Which point is the highest (worst), next-highest and lowest (best)
      int ilo = 0;
      int ihi = 1;
      // Whenever we "look at" a vertex, it gets a random thermal fluctuation
      double ylo = y[0] + tt * log(G4UniformRand());
      double ynhi = ylo;
      double yhi = y[1] + tt * log(G4UniformRand());
      if(ylo > yhi)
        {
          ihi = 0;
          ilo = 1;
          ynhi = yhi;
          yhi = ylo;
          ylo = ynhi;
        }
      for(int i=2; i<mpts; i++)
        { // Loop over the points in the simplex
          double yt = y[i] + tt * log(G4UniformRand());
          if(yt <= ylo)
            {
              ilo = i;
              ylo = yt;
            }
          if(yt > yhi)
            {
              ynhi = yhi;
              ihi = i;
              yhi = yt;
            }
          else if(yt > ynhi)
            {
              ynhi = yt;
            }
        }
      double rtol = 2.0 * fabs(yhi-ylo) / (fabs(yhi) + fabs(ylo));
      // Compute the fractional range from highest to lowest and
      // return if satisfactory
      if(rtol < ftol || iter < 0)
        { // If returning, put best point and value in slot 1
          double swap = y[0];
          y[0] = y[ilo];
          y[ilo] = swap;
          for(int n=0; n<ndim; n++)
            {
              swap = p[0 + n]; // keep 0 to remind of vector layout
              p[0 + n] = p[ilo*ndim + n];
              p[ilo*ndim + n] = swap;
            }
          break;
        }
      iter -= 2;
      // Begin a new iteration. First extrapolate by a factor -1 through the face of
      // the simplex across from the high point, i.e. reflect the simplex from the high point.
      double ytry = Amotsa(p, y, psum, ndim, pb, yb, ihi, yhi, -1.0, tt);
      if(ytry <= ylo)
        {
          // Gives a result better than the best point, so try an additional extrapolation
          // by a factor of 2
          ytry = Amotsa(p, y, psum, ndim, pb, yb, ihi, yhi, 2.0, tt);
        }
      else if(ytry >= ynhi)
        {
          // The reflected point is worse than the second-highest, so look for an intermediate
          // lower point, i.e. do a one-dimensional contraction.
          double ysave = yhi;
          ytry = Amotsa(p, y, psum, ndim, pb, yb, ihi, yhi, 0.5, tt);
          if(ytry >= ysave)
            { // Can't seem to get rid of that high point.  Better crontract around the lowest
              // (best) point.
              for(int i=0; i<mpts; i++)
                {
                  if(i!=ilo)
                    {
                      for(int j=0; j<ndim; j++)
                        {
                          psum[j] = 0.5 * (p[i*ndim + j] + p[ilo*ndim + j]);
                          p[i*ndim + j] = psum[j];
                        }
                      y[i] = fMinFactor * (*fComponent)(psum);
                    }
                }
              iter -= ndim;
              // Recompute PSUM
              GetPSum(psum, ndim, p);
            }
        }
      else
        {
          ++(iter); // Correct the evaluation count.
        }
    }
}


double
SimulatedAnnealing::Amotsa(std::vector<double>& p, std::vector<double>& y, std::vector<double>& psum, int ndim,
                           std::vector<double>& pb, double& yb, int ihi, double& yhi, double fac, const double& tt)
{
  // Extrapolates by a factor fac through the face of the simplex across from the high point,
  // tries it, and replaces the high point if the new point is better.
  double fac1 = (1.0-fac)/ndim;
  double fac2 = fac1-fac;

  vector<double> ptry(ndim, 0);

  for (int j=0; j<ndim; j++)
    ptry[j] = psum[j] * fac1 - p[ihi*ndim + j] * fac2;

  double ytry = fMinFactor * (*fComponent)(ptry);

  if (ytry <= yb)
    { // Save the best-ever.
      for (int j=0; j<ndim; j++)
        pb[j] = ptry[j];
      yb = ytry;
    }

  // We added a thermal fluctuation to all the current vertices, but we subtract it here
  // so as to give the simplex a thermal Brownian motion: it likes to accept any
  // suggested change
  double yflu = ytry - tt * log(G4UniformRand());

  if (yflu < yhi)
    {
      y[ihi] = ytry;
      yhi = yflu;
      for (int j=0; j<ndim; j++)
        {
          psum[j] += ptry[j] - p[ihi*ndim + j];
          p[ihi*ndim + j] = ptry[j];
        }
    }

  return yflu;

}