/* Nelder-Mead simplex algorithm for unconstrained non-linear function minimization

minimize function f(x,n,p) starting at initial vector x0 of size n: simplex<Real> S(f,x0,n);

Based on https://fog.misty.com/perry/ss/src/simplex.c
 from SS: https://fog.misty.com/perry/ss/ref.html

R. Perry, 11 June 2022
*/

#ifndef _SIMPLEX_H_
#define _SIMPLEX_H_

/* matalloc: allocates space for m-by-n matrix, returns pointer.
The matrix elements are contiguous, since space for all of the rows
is obtained with one malloc() call.  a[-1] (as well as a[0]) is set
to point to the storage so that rows may be interchanged by swapping
row pointers, and matfree() can still reference the storage through
a[-1] even if a[0] was changed.
*/
template<typename Real> Real **matalloc( Real xtol, int m, int n) {
// xtol not used, just there so template knows what Real is
  Real **a;
  a = (Real **) emalloc( (m+1) * sizeof(Real *)); // get space for m+1 pointers to Real
  a[0] = a[1] = (Real *) emalloc( m*n*sizeof(Real)); // get space for all of the rows, m*n Real
 // set the pointers a[1]..a[m-1] to point to their row
  ++a; for( int i = 1; i < m; i++) a[i] = a[i-1] + n;
  return a;
}

// free the matrix data and row pointers
template<typename Real> void matfree( Real **a) {
  if( a) { free(a[-1]); free(a-1); }
}

template <typename Real = double> struct simplex {

  int N;	// number of independent variables
  int L, NH, H;	// indices of Lowest, Next-to-Highest, Highest,
  int C, D, R;	//  Centroid, Direction, Reflection,
  int S, T;	//  Stepsize and Temporary points
  int count;	// function evaluation count
  Real size;	// simplex size
  Real *d;	// distances from x[L]
  Real *f;	// function values
  Real **x;	// x[i], i=0...N, the N+1 simplex points
		//   i=N+1...N+5, Centroid, Direction, Reflection, Stepsize, Temporary
  int debug;    // enable debug output
  void *params; // extra parameters argument for F()
  Real step;	// optional step size, default is 0.1*x0 if step is zero

  Real (*F)(const Real *x, int n, void *params); // function to minimize

 // compute distance between two points (inf norm)
  Real dist( Real *x1, Real *x2) { Real d = 0, e;
    for( int i = 0; i < N; ++i) { e = fabs( x2[i] - x1[i]); if( e > d) d = e; }
    return d;
  }

 // set simplex distances
  void setdist( void) {
    Real *xl = x[L]; for( int i = 0; i <= N; ++i)
    d[i] = (i == L) ? 0 : dist( xl, x[i]);
  }

 // set simplex size, assumes distances are already set
  void setsize( void) {
    size = 0; for( int i = 0; i <= N; ++i) if( d[i] > size) size = d[i];
  }

 // determine the ordering indices L, NH, and H.
 // calls setdist() if L != oldL; always calls setsize() at the end
  void order( int oldL) {
   // f[0] and f[1] determine the initial values for L, NH, and H
    if( f[0] < f[1]) { L = 0; H = 1; } else { L = 1; H = 0; }
    NH = L;
   // examine f[2]...f[N] to complete the ordering
    for( int i = 2; i <= N; ++i)
      if( f[i] < f[L]) { L = i; } // new Lowest point
      else if( f[i] >= f[H]) { NH = H; H = i; } // new Highest point
      else if( f[i] >= f[NH]) { NH = i; } // new Next-to-Highest point
    if( L != oldL) setdist();
    setsize();
  }

 // replace highest point with reflection, extension, or contraction
  void update( int p, const char *msg) {
    if( debug) std::cout << msg << '\n';
    Real *t = x[H]; x[H] = x[p]; x[p] = t; // swap pointers instead of copying
    f[H] = f[p]; d[H] = dist( x[L], x[H]);
    order( L); // order() will call setsize()
  }

 // compute centroid, direction, and reflection point
  Real reflect( void) {
    Real *xc = x[C], *xd = x[D], *xr = x[R];
    for( int j = 0; j < N; ++j) xc[j] = 0;
    for( int i = 0; i <= N; ++i) if( i != H)
      for( int j = 0; j < N; ++j) xc[j] += x[i][j];
    for( int j = 0; j < N; ++j) xc[j] /= N;
    for( int j = 0; j < N; ++j) xd[j] = xc[j] - x[H][j];
    for( int j = 0; j < N; ++j) xr[j] = xc[j] + xd[j];
    ++count; return f[R] = F( xr, N, params);
  }

 // compute expansion point
  Real expand( void) {
    Real *xc = x[C], *xd = x[D], *xe = x[T];
    for( int j = 0; j < N; ++j) xe[j] = xc[j] + 2*xd[j];
    ++count; return f[T] = F( xe, N, params);
  }

 // compute outside contraction point
  Real contract_outside( void) {
    Real *xc = x[C], *xd = x[D], *t = x[T];
    for( int j = 0; j < N; ++j) t[j] = xc[j] + 0.5*xd[j];
    ++count; return f[T] = F( t, N, params);
  }

 // compute inside contraction point
  Real contract_inside( void) {
    Real *xc = x[C], *xd = x[D], *t = x[T];
    for( int j = 0; j < N; ++j) t[j] = xc[j] - 0.5*xd[j];
    ++count; return f[T] = F( t, N, params);
  }

 // shrink around lowest point, return non-zero if size does not decrease
  int shrink( void) {
    if( debug) std::cout << "shrink\n";
    Real *xl = x[L];
    for( int i = 0; i <= N; ++i) if( i != L) {
      for( int j = 0; j < N; ++j) x[i][j] = 0.5*(x[i][j] + xl[j]);
      f[i] = F( x[i], N, params);
    }
    count += N;
   // check that new size is smaller
    Real oldsize = size;
    order( -1); // order() will call setdist() and setsize()
    if( size < oldsize) return 0;
    if( debug) std::cout << "shrink failed to decrease size\n";
    return 1;
  }

 // constructor, performs the minimization
  simplex( Real (*_F)(const Real *x, int n, void *params), const Real *_x0, int _N, void *_params,
    Real _step, Real xtol, Real ftol, int _debug = 0, int Flimit = 0) {

    F = _F; N = _N; params = _params; step = _step; debug = _debug;
    C = N+1; D = N+2; R = N+3; S = N+4; T = N+5;
    if( Flimit == 0) Flimit = 200*N;
    x = matalloc( xtol, N+1+5, N);
    f = (Real *) emalloc( (N+1+5)*sizeof(Real));
    d = (Real *) emalloc( (N+1+5)*sizeof(Real));

   // initialize starting point and stepsize
    Real *x0 = x[0], *dx = x[S];
    for( int j = 0; j < N; ++j, ++x0, ++_x0, ++dx) {
      *x0 = *_x0; *dx = (step) ? step : ((*x0) ? (0.1 * *x0) : 0.1);
    }

   // initialize x[1]...x[N] from x[0] and step sizes
    for( int i = 1, j = 0; i <= N; ++i, ++j) {
      memcpy( x[i], x[0], N*sizeof(Real));
      Real delta = x[S][j]; x[i][j] += delta; // make sure the point is altered
      while( x[i][j] == x[0][j]) { delta = 10*delta + xtol; x[i][j] += delta; }
    }

   // initialize count and function values
    int iter = 0; count = N+1; for( int i = 0; i <= N; ++i) f[i] = F( x[i], N, params);

   // order() will call setdist() and setsize()
    order( -1);

   // main loop, performed until function and size are both converged,
   // or function count limit is reached
    while( 1)
    {
      int done = 0; Real fR, fE, fC, fL = f[L], fNH = f[NH], fH = f[H];

      if( fH - fL <= ftol) { ++done;
        if( debug) std::cout << "function converged within tolerance\n";
      }
      if( size <= xtol) { ++done;
        if( debug) std::cout << "size converged within tolerance\n";
      }
      if( done == 2) break; // function and size both converged
      if( count >= Flimit) { std::cout << "simplex: function count limit reached\n"; break; }
      ++iter; fR = reflect();

      if( fR < fL) { fE = expand();
        if( fE < fR) update( T, "expand");
        else update( R, "reflect");
      } else if( fR < fNH) {
	update( R, "reflect");
      } else if( fR < fH) {
        fC = contract_outside();
        if( fC <= fR) update( T, "contract outside");
        else if( shrink()) break;
      } else { fC = contract_inside(); // fR >= fH
        if( fC < fH) update( T, "contract inside");
        else if( shrink()) break;
      }
      if( debug && iter % 10 == 0) std::cout << "iter = " << iter << ", size = " << size
	<< ", fL = " << f[L] << ", fH = " << f[H] << '\n';
    }
    // if( debug) display();
  }

  ~simplex() { matfree(x); free( f); free( d); }
};

#endif // _SIMPLEX_H_