#include <stdlib.h>
#include <math.h>
#include <stdio.h>
#include <assert.h>
#include "gpk_std.h"
#include "gpklib.h"
#include "gpklib/numerics.h"
#include "gpklib/ieee754.h"

int amoeba_anneal_noisy = FALSE;

#define D	if(amoeba_anneal_noisy)


int amoeba_anneal_qq = 2;

static double ran1aqq()
{
 double tmp;
 OPENMP(omp critical(amoeba_anneal_rand_qq))
 { tmp = ran1a(&amoeba_anneal_qq); }
 return tmp;
}

static double randexp()
{
 return -log(ran1aqq());
}


static double randH()
{
 return exp(1.0*(ran1aqq()-0.5));
}



static double f(double *p, int n, void *c, int nc,
		double (*funk)(C double *p, int n, void *c, int nc),
		void (*project)(double *p, int n, void *c, int nc)
		)
{
 if(project)
 	OPENMP(omp critical(amoeba_anneal_project))
 	(*project)(p, n, c, nc);
 return (*funk)(p, n, c, nc);
}


LOCAL void copy_into(double *to, C double *from, int n)
{
 int i;
 OPENMP(omp parallel for private(i) schedule(guided, 1024))
 for(i=0; i<n; i++)
 	to[i] = from[i];
}


LOCAL double ni(double x)
{ return isnan(x) ? infinity() : x;}


LOCAL void writeout(C double *p, int n, FILE *afp)
{
 int i;
 for(i=0;i<n;i++)
	fprintf(afp, "%g%c", p[i], i==n-1 ? '\n' : ' ');
 fflush(afp);
}


#define assert_finite(x) x_assert(finite(x), __LINE__, __FILE__)

static void x_assert(int tf, int line, const char *file)
{
 if(!tf) {
 	fprintf(stderr, "ERR: non-finite value returned from funk, line %d in %s\n",
			line, file);
	abort();
	}
}


/* Returns 1 if the iteration count (itmax) was reached, 0 if not.
 * The initial parameter vectors are placed in "p". Amoeba_anneal_allfree
 * modifies that 2-dimensional array, and the final converged values are 
 * left there when it is finished.
 *
 * */
int amoeba_anneal_allfree(double *C*p,
				int n,
				double ftol,
				double (*funk)(C double *p, int n, void *c, int nc),
				int itmax,
				void *c,
				int nc, double T0, double Tfac)
{
 return amoeba_anneal_allfree2(p, n, n+1, ftol, funk, itmax, c, nc, T0, Tfac, 0);
}


int amoeba_anneal_allfreex(double *const*p,
				int n,
				double ftol,
				double (*funk)(const double *p, int n, void *c, int nc),
				int itmax,
				void *c,
				int nc, double T0, double Tfac,
				void (*project)(double *p, int n, void *c, int nc)
				)
{
 return amoeba_anneal_allfree2(p, n, n+1, ftol, funk, itmax, c, nc, T0, Tfac, project);
}



int amoeba_anneal_allfree2(double *const*p,
				int n, int nvertex,
				double ftol,
				double (*funk)(const double *p, int n, void *c, int nc),
				int itmax,
				void *c,
				int nc, double T0, double Tfac,
				void (*project)(double *p, int n, void *c, int nc)
				)
{
 /* MAXTF is a way of causing the algorithm to change from a
  * minimization-oriented behaviour when you're far from the minimim
  * to a more Markov-Chain Monte-Carlo behaviour when you're near
  * equilibrium.    Basically, if the y-values of the vertices
  * are nearly equal and MAXTF>>log(nvertex), the vertices are picked with
  * more-or-less equal probabilities.
  */
 C double MAXTF = 2.0;
 int i, ihi, inhi, ilo, iter;
 double *y, *pbar, *pr, ypr, ypr_real;
 double T;
 FILE *afp = 0;
 fprintf(stderr, "T0=%g Tfac=%g nc=%d n=%d\n", T0, Tfac, nc, n);

 D afp = fopen("amoeba.tmp", "w");

 assert(n > 0);
 assert(nvertex > n);
 iter = 0;
 y = (double *)malloc(sizeof(*y) * nvertex);
 assert(y);
 pbar = (double *)malloc(sizeof(*pbar) * n);
 assert(pbar);
 pr = (double *)malloc(sizeof(*pr) * n);
 assert(pr);
 OPENMP(omp parallel for schedule(guided, 4))
 for(i=0; i<nvertex; i++) {
 	y[i] = ni(f(p[i], n, c, nc, funk, project));
	assert_finite(y[i]);
 }

 for(iter=0,T=T0; iter<itmax; iter++)	{
	 {
	 double ylo = infinity(), yhi = -infinity(), ynhi = -infinity();
	 for(ilo=ihi=inhi=-1,i=0; i<nvertex;i++)	{
		C double ytmp = y[i]-MAXTF*T*randexp();
		assert_finite(ytmp);
	 	if(ytmp < ylo)	{ ilo = i; ylo = ytmp; }
		if(ytmp >= yhi)	{
			if(finite(yhi))	{ inhi = ihi; ynhi = yhi; }
/*
D			fprintf(stderr, "accepting ytmp=%g to replace yhi=%g for y[%d]=%g\n",
					ytmp, yhi, i, y[i]);
					*/
			ihi = i; yhi = ytmp;
			}
		else if(ytmp>ynhi && finite(ytmp))	{ inhi = i; ynhi = ytmp; }
		}
	/* Inhi indexes the next-to-highest value, or the largest finite
		value ( if ihi indexes a non-finite value). */
	if(inhi < 0)	{inhi = ilo; ynhi = ylo;}
D	fprintf(stderr, "T=%g ylo=%d=%g ynhi=%d=%g yhi=%d=%g\n",
					T, ilo, ylo, inhi, ynhi, ihi, yhi);
D	fprintf(stderr, "T=%g y[lo=%d]=%g y[inhi=%d]=%g y[ihi=%d]=%g\n",
					T, ilo, y[ilo], inhi, y[inhi], ihi, y[ihi]);
	 assert_finite(ylo);
	 assert_finite(ynhi);
	 assert(ynhi >= ylo);
	 assert(!(yhi<ynhi));
	}

	 assert(ihi != ilo);
	 if(T+fabs(y[ihi]-y[ilo]) < (fabs(y[ihi]) + fabs(y[ilo]))*ftol)
		break;
	/* Evaluate the test vector. This is a vector through the center of */
	/* the face which is opposite the highest point. */
        /* OPENMP(omp parallel for collapse(2) schedule(guided, 20)) */
        OPENMP(omp parallel for private(i) schedule(dynamic, 20))
	for(i=0; i<n; i++) {
		int j;
		pbar[i] = 0.0;
		for(j=0; j<nvertex; j++)
			if(j != ihi)
				pbar[i] += p[j][i];
		pbar[i] /= (nvertex-1);
		}
	C double Alpha = randH();
	OPENMP(omp parallel for private(i) schedule(dynamic, 512))
	for(i=0;i<n;i++)	{
 		pr[i] = (1+Alpha)*pbar[i] - Alpha * p[ihi][i];
		}
	/* evaluate the function at the reflected point */
	ypr = (ypr_real = ni(f(pr, n, c, nc, funk, project))) - randexp()*T;
D	fprintf(stderr, "Tested reflection = %g (%g real)\n", ypr, ypr_real);
	if(ypr < y[ihi])	{
D		fprintf(stderr, "reflection accepted - excellent\n");
D		writeout(pr, n, afp);
		copy_into(p[ihi], pr, n);
		y[ihi] = ypr_real;
		T *= Tfac;
		continue;
		}


	C double Beta = randH();
	OPENMP(omp parallel for private(i) schedule(dynamic, 256))
	for(i=0; i<n; i++)
		pr[i] = (1-Beta)*pbar[i] + Beta*p[ihi][i];
	ypr = (ypr_real = ni(f(pr, n, c, nc, funk, project))) - randexp()*T;
	if(ypr < y[ihi])	{
D		fprintf(stderr, "local accepted - OK\n");
D		writeout(pr, n, afp);
		copy_into(p[ihi], pr, n);
		y[ihi] = ypr_real;
		T *= sqrt(Tfac);
		continue;
		}
	}
 D fclose(afp);
 free((char *)y);
 free((char *)pbar);
 free((char *)pr);

 return iter>=itmax;
}


#ifdef TEST

LOCAL double p2s = 0.0, p2c = 0.0;

double func(C double *p, int n, void *c, int nc)
{
 C double f = p[0]*p[0] + p[1]*p[1];
 p2c += 1;
 p2s += f;
 C double tmp =  p[0]>0 ? (p[1]>0 ? f : quiet_nan(0)) : infinity();
 fprintf(stderr, "%g %g -> %g %g avg= %g\n", p[0], p[1], f, tmp, p2s/p2c);

 return tmp;
}

main()
{
 double x0[2], x1[2], x2[2];
 double *xpp[3];

 xpp[0] = x0; xpp[1] = x1; xpp[2] = x2;
 x0[0] = 0.0003; x0[1] = 0.001;
 x1[0] = 0.001; x1[1] = .004;
 x2[0] = 0.002; x2[1] = .0004;

 amoeba_anneal_allfree(xpp, 2, 1e-6, func, 1000, 0, 0, 1.0, 1.0);
}

#endif
