#include <RAT/OptimisedComponent.hh>
#include <RAT/Powell.hh>
#include <RAT/DS/FitResult.hh>
#include <RAT/Log.hh>
using namespace RAT::Optimisers;

#include <iostream>
#include <cmath>

using namespace std;

void Powell::Initialise( const std::string& )
{
  fIterMax = 100;    // Number of allowed iterations of optimiser
  fUp = 0.5; // FIXME: Assumes we're using a log likelihood minimisation
}

double Powell::Minimise()
{
  // Minimise (call PowellOptimise)
  fMinFactor = 1.0;
  return fMinFactor*PowellOptimise();
}

double Powell::Maximise()
{
  // Maximise
  fMinFactor = -1.0;
  return fMinFactor*PowellOptimise();
}

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

  // Make a copy of the initial point
  vector<double> initParams = params;

  const int size= params.size();
  const int size2 = size*size;

  // Define intial direction matrix
  // Take one unit vector for each parameter we need to optimise
  // and set the rest of the matrix to zero

  vector<double> xi(size2, 0);

  for(int i=0; i<size; i++)
    xi[i*size + i] = 1;

  double ftol=1e-7;  // Fractional tolerance of fit
  int iter=0;          // Num of iterations
  double fret=0;       // Function value at p
  int errorsReported=0; // Are the errors good (0 if OK, !=0 if invalid)
  double initY; // initial param value

  powell(params, xi, ftol, iter, fret, initY);

  ReportErrors(params, initParams, initY, iter, fret, posErrors, negErrors, errorsReported);

  fParams = params;
  if(errorsReported!=0)
    {
      debug << "Powell::Unable to report errors, fail mode " << errorsReported << "; setting to seeded values" << newline;
      debug << "|Init - minimized| params:";
      for(unsigned int i=0;i<params.size();i++)
        debug << "\t" << fabs(initParams[i] - params[i]);
      debug << newline << "Yinit " << initY << "\tYfin " << fret << newline;
      debug << "NITER: " << iter << newline << newline;
    }
  fPositiveErrors = posErrors;
  fNegativeErrors = negErrors;

  // Did fit finish within the max number of iterations?
  if(iter < fIterMax) fValid = true;
  else fValid = false;

  return fret;
}


void Powell::ReportErrors(std::vector<double> &p, const std::vector<double> &initP, const double initY, int &iter, double &fret, std::vector<double> &posErrors, std::vector<double> &negErrors, int& errorsValid)
{

  // iter==0: minimisation hasn't happened, don't bother looking for an error.
  // iter>0: xi[ibig] will have the direction of steepest descent, use xi in each direction to assist with finding errors
  // Assume we're able to use secant root finding (smooth in region of minimum)
  // Simple error reporting, first attempt to bracket the root

  errorsValid = 0;

  if(iter==0)
    {
      errorsValid = 1;
      return; // Cannot do anything; ignore (quite possible fit should be invalid)
    }

  double yRoot = fret + fUp; // fret was the LL value from the minimisation, just look for 2LL for an error

  vector<double> uncertainties(p.size(), 0);

  // flag for if we find a better minimum whilst trying to find the 1 sigma point
  bool new_min=false;
  for(size_t i=0;i<p.size();i++)
    {
      // for each point, run secant method
      vector<double> pa = p;
      vector<double> pb = p;

      //      double shift = fabs( xi[i*p.size() + (p.size()-1)] );
      // Set the initial shift to be the difference between start and end point / (ystart - ymin)/0.5
      double shift = fabs(p[i] - initP[i]) / ( (initY - fret) / p.size() );
      if(shift<1e-2)
        shift = 1e-2; // don't use silly values!

      pb[i] = pb[i] + shift;
      double yl = fMinFactor * (*fComponent)( pa );
      double y = fMinFactor * (*fComponent)( pb );
      // if y is not > yRoot, we need to adjust the value
      int itry = 0;

      while( y < yRoot )
        {
          itry++;
          if(itry > 5)
            break;
          pb[i] += shift * std::pow(10, (double)itry); // just move attempt further out, use order of magnitude steps (for speed)
          y = fMinFactor * (*fComponent)( pb );
        }

      if(itry>5)
        {
          errorsValid = 2;
          return; // give up
        }

      double x1 = pa[i];
      double x2 = pb[i];
      double xl, rts;
      if(fabs(yl - yRoot) < fabs(y - yRoot))
        {
          rts = x1;
          double ytemp = y;
          xl = x2;
          y = yl;
          yl = ytemp;
        }
      else
        {
          xl = x1;
          rts = x2;
        }

      itry = 0;
      // have bracketed the root; now find it
      for(int j=0;j<50;j++)
        {
          itry ++;

          double dx = 0.0;
          //Avoid divide by 0 problems which happen (rarely) when likelihood is very flat
          if(y!=yl){
            dx = (xl - rts) * (y-yRoot) / (y - yl);
          } else {
            dx = (xl - rts) * (y-yRoot) / 0.01;
          }
          xl = rts;
          yl = y;
          rts += dx;
          pb[i] = rts;
          y = fMinFactor * (*fComponent)( pb );
          if(y<fret){
            //If a better minimum has been found, we should take this as the new minimum and start again
            p[i] = rts;
            fret = y;
            //Make use of the iter counter to ensure there's a termination condition
            iter = iter + itry;
            new_min = true;
            break;
          }
          if( fabs(y - yRoot) < 0.01)
            break; // Convergence
        }
      if(new_min) break;
      uncertainties[i] = pb[i] - pa[i];
      if(std::isnan(uncertainties[i]))
        {
          debug << "Attempted to assign nan uncertainty: " << i << "\t" << pb[i] << "\t" << pa[i] << newline;
          errorsValid = 3;
          return;
        }
    } // i (p.size())
    // If better min was found, repeat process with new minimum fret and point p
    if(new_min && iter<fIterMax) {
        ReportErrors(p,initP,initY,iter,fret,posErrors,negErrors,errorsValid);
        return;
    }

  // Success
  for(unsigned int i=0; i<p.size(); i++)
    posErrors[i] = negErrors[i] = uncertainties[i];

}

// The functions below are used to carry out the actual minimisation.
// They are copied heavily from the the Powell Optimisation section
// in the Numerical Recipes book.

void Powell::powell(std::vector<double> &p, std::vector<double>& xi, double ftol, int &iter, double &fret, double &initY)
{
  // Minimization of a function func of n variables.  Input consists of
  // an initial starting point p[1..n]; an initial matrix xi[1..n][1..n],
  // whose columns contain the initial set of directions (usually the n
  // unit vectors); and ftol, the fractional tolerance in the function
  // value such that failure to decrease by more than this amount on
  // one iteration signals doneness.  On output, p is set to the best
  // point found, xi is the then-current direction set, fret is the
  // returned function value at p, and iter is the number of iterations
  // taken.  The routine linmin is used.

  int ibig=0;
  const int n = p.size();

  double fp=0, del=0.0, fptt=0, t=0;
  vector<double> xit(n, 0), pt(n, 0), ptt(n, 0);

  // Save the initial point
  fret= fMinFactor * (*fComponent)( p );
  initY = fret;
  for(int j=0; j<n; j++) pt[j]=p[j];


  for(iter=0; ; ++iter)
    {
      fp=fret;
      ibig=0;
      del=0.0;
      for(int i=0; i<n; i++)
        {                    // Loop over all directions
          for(int j=0; j<n; j++)
            xit[j] = xi[j*n + i];
          fptt=fret;
          linmin(p,xit,fret);                      // Minimise along it
          if (fabs(fptt-(fret)) > del)
            {           // Record the largest decrease so far
              del=fabs(fptt-(fret));
              ibig=i+1;
            }
        }
      if (2.0*fabs(fp-fret) <= ftol*(fabs(fp)+fabs(fret)))
        {   // Termination criterion
          return;
        }

      if (iter == fIterMax) return;

      for(int j=0; j<n; j++)
        {                    // Construct the extrapolated point and
          ptt[j]=2.0*p[j]-pt[j];                   // the average direction moved. Save the
          xit[j]=p[j]-pt[j];                       // old starting point.
          pt[j]=p[j];
        }

      fptt = fMinFactor * (*fComponent)( ptt );                        // Function value at extrapolated point
      if(fptt < fp)
        {
          t=2.0*(fp-2.0*fret+fptt)*pow((fp-fret-del),2)-del*pow((fp-fptt),2);
          if (t < 0.0)
            {
              linmin(p,xit,fret);
              for(int j=0; j<n; j++)
                {
                  xi[j*n + ibig-1] = xi[j*n + n-1];
                  xi[j*n + n-1] = xit[j];
                }
            }
        }
    }
}

double Powell::f1dim(double x)
{
  //Must accompany linmin
  vector<double> xt(fNcom,0);

  for(int j=0; j<fNcom; j++) xt[j]=fPcom[j]+x*fXicom[j];

  double f=fMinFactor*(*fComponent)( xt );

  return f;
}

void Powell::linmin(std::vector<double> &p, std::vector<double> &xi, double &fret)
{
  // Given an n-dimensional point p[1..n] and an n-dimensional direction
  // xi[1..n], moves and resets p to where the function func(p) takes on
  // a minimum along the direction xi from p, and replaces xi by the actual
  // vector displacement that p was moved.  Also returns as fret the value
  // of func at the returned location p.  This is actually all accomplished
  // by calling the routines mnbrak and brent.

  double tol = 2e-4;              // Tolerance to be passed to brent
  const int n=p.size();
  double xx=1, xmin=0, fx=0, fb=0, fa=0, bx=0, ax=0;

  fNcom=n;                        // Define global variables
  fPcom=p;
  fXicom=xi;

  ax=0;                           // Initial guess for brackets
  xx=1;

  // Attempt to bracket the minimum
  mnbrak_linmin(ax,xx,bx,fa,fx,fb);

  // Carry out a one-dimensional minimisation
  fret=brent_linmin(ax,xx,bx,tol,xmin);

  // Construct vectors to return
  for(int j=0;j<n;j++)
    {
      xi[j] *= xmin;
      p[j] += xi[j];
    }
}

#define SIGN(a,b) ((b) >= 0.0 ? fabs(a) : -fabs(a))

void Powell::mnbrak_linmin(double &ax, double &bx, double &cx, double &fa, double &fb, double &fc)
{
  // Given a function f1dim and distinct intial points ax and bx, this
  // routine searches in the downhill direction (defined by the function
  // as evaluated at the initial points) and returns the new points
  // ax, bx and cx that bracket a minimum of the function. Also returned
  // are the function values at the three points, fa, fb and fc.

  // Here gold is the default ratio by which successive intervals are
  // magnified, glimit is the maximum magnification allowed for a
  // parabolic-fit step.
  double gold = 1.618034, tiny = 1e-20, glimit = 100;

  double ulim=0,u=0,r=0,q=0,fu=0,dum=0;

  fa=f1dim(ax);
  fb=f1dim(bx);

  if (fb > fa) {
    shift(dum,ax,bx,dum);                    // Switch a and b so that we can go downhill
    shift(dum,fb,fa,dum);                    // in the direction from a to b
  }

  cx=bx+gold*(bx-ax);                        // First guess for c
  fc=f1dim(cx);
  while (fb > fc) {
    r=(bx-ax)*(fb-fc);                       // Compute u by parabolic extrapolation from
    q=(bx-cx)*(fb-fa);                       // a, b, c. Use tiny to prevent division by zero.
    u=(bx)-((bx-cx)*q-(bx-ax)*r)/
      (2.0*SIGN(max(fabs(q-r),tiny),q-r));
    ulim=bx+glimit*(cx-bx);

    if ((bx-u)*(u-cx) > 0.0){                // Parabolic u is between b and c,
      fu=f1dim(u);
      if (fu < fc){                          // Minimum between b and c
	ax=bx;
	bx=u;
	fa=fb;
	fb=fu;
	return;
      } else if (fu > fb){                   // Minimum between a and u
	cx=u;
	fc=fu;
	return;
      }
      u=cx+gold*(cx-bx);                     // Parabolic fit was no use. Use default magnification
      fu=f1dim(u);
    } else if ((cx-u)*(u-ulim) > 0.0){       // Parabolic fit is between c and its allowed limit
      fu=f1dim(u);
      if (fu < fc){
	shift(bx,cx,u,u+gold*(u-cx));
	shift(fb,fc,fu,f1dim(u));
      }
    } else if ((u-ulim)*(ulim-cx) >= 0.0) {  // Limit parabolic u to max allowed value
      u=ulim;
      fu=f1dim(u);
    } else {
      u=(cx)+gold*(cx-bx);                   // Reject parabolic u, use default magnification
      fu=f1dim(u);
    }
    shift(ax,bx,cx,u);                       // Eliminate oldest point and continue
    shift(fa,fb,fc,fu);
  }
}

float Powell::brent_linmin(double ax, double bx, double cx, double tol, double &xmin)
{
  // Given a function f1dim and given a bracketing triplet of abscissas ax, bx, cx
  // (such that bx is between ax and cx and f1dim(bx) is less than both f1dim(ax) and
  // f1dim(cx)), this routine isolates the minimum to a fractional precision of tol
  // using Brent's method. The abscissa of the minimum is returned as xmin, and the
  // minimum function is returned as brent, the returned function value.

  // Here itmax is the maximum allowed number of iterations; gold is the golden ratio;
  // zeps is a small number that protects against trying to achieve fractional accuracy
  // for a minimum that happens to be exactly zero
  int itmax = 100;
  double cgold = 0.3819660, zeps = 1e-10;

  double a=0, b=0, d=0, etemp=0, fu=0, fv=0, fw=0, fx=0, p=0, q=0, r=0, tol1=0, tol2=0, u=0, v=0, w=0, x=0, xm=0, e=0.0;

  a=(ax < cx ? ax : cx);      // a and b must be in ascending order but the input
  b=(ax > cx ? ax : cx);      // abscissas need not be
  x=w=v=bx;
  fw=fv=fx=f1dim(x);

  for(int iter=0; iter<itmax; iter++) {           // Main program loop
    xm=0.5*(a+b);
    tol2=2.0*(tol1=tol*fabs(x)+zeps);
    if (fabs(x-xm) <= (tol2-0.5*(b-a))){     // Test for done here
      xmin=x;
      return fx;
    }
    if (fabs(e) > tol1) {                    // Construct a trial parabolic fit
      r=(x-w)*(fx-fv);
      q=(x-v)*(fx-fw);
      p=(x-v)*q-(x-w)*r;
      q=2.0*(q-r);
      if (q > 0.0) p = -p;
      q=fabs(q);
      etemp=e;
      e=d;
      if (fabs(p) >= fabs(0.5*q*etemp) || p <= q*(a-x) || p >= q*(b-x))
	d=cgold*(e=(x >= xm ? a-x : b-x));
      // The above conditions determine the acceptability of the parabolic fit.
      // Here we take the golden section step into the larger of the two segments.
      else {
	d=p/q;                               // Take the parabolic step
	u=x+d;
	if (u-a < tol2 || b-u < tol2)
	  d=SIGN(tol1,xm-x);
      }
    } else {
      d=cgold*(e=(x >= xm ? a-x : b-x));
    }
    u=(fabs(d) >= tol1 ? x+d : x+SIGN(tol1,d));
    fu=f1dim(u);                             // This is the one function evaluation per iteration
    if (fu <= fx) {                          // Now decide what to do with our function evaluation
      if (u >= x) a=x; else b=x;
      shift(v,w,x,u);                        // Housekeeping follows:
      shift(fv,fw,fx,fu);
    } else {
      if (u < x) a=u; else b=u;
      if (fu <= fw || w == x) {
	v=w;
	w=u;
	fv=fw;
	fw=fu;
      } else if (fu <= fv || v == x || v == w) {
	v=u;
	fv=fu;
      }
    }                                        // Done with housekeeping. Back for another iteration.
  }
  RAT::Log::Die("Powell::brent_linmin: Too many iterations in brent");
  xmin=x;
  return fx;
}

void Powell::shift(double &a, double &b, double &c, double d)
{
  a=b;
  b=c;
  c=d;
}