/*
   pa1.c

   First of two implementation files for the computation of piece-wise
   analytic fits based on sample data.  Implements those definitions
   in the header file pa.h that do not require a random number
   generator.

   Ramses van Zon, December 4, 2009

   Reference:
   [] Ramses van Zon and Jeremy Schofield, "Constructing smooth
      potentials of mean force, radial distribution functions and
      probability densities from sampled data", 
      arxiv.0912.0465 [cond-mat.stat-mech].
*/

#include <stdlib.h>
#include <stdio.h>
#include <math.h>
#include "pa.h"

/* some local helper functions */
static double pa_sqr( double x )
{ 
  return x*x;
}

static double pa_min( double x, double y )
{ 
  return x<y?x:y; 
}

/* computation of the q values from the Kolmogorov-Smirnov test */
static double pa_q( int n, const double *f, int *j )
{
  double a, b, d, e, c, g, h, q;
  int i;

  b = sqrt(n);
  e = g = h = q = 0.0;
  c = 2;
  *j = 0;
  for (i = 0; i < n; ++i) {
    d = f[i]-e;
    if (d > h) {
      *j = i;
      h = d;
    }
    e = (double)(i+1)/n;
    d = e-f[i];
    if (d > h) {
      *j = i;
      h = d;
    }
  }
  d = b*h+0.12*h+0.11*h/b;
  a = -2*d*d;
  for(i = 1;i<= 200;++i){
    d = c*exp(a*i*i);
    q += d;
    if (fabs(d) <= g)
      break;
    c = -c;
    g = fabs(d)/1000;
  }
  if ( i > 200 )
    printf("#pa_q: no convergence.\n");
  return q;
}

/* perform a simple block-jackknife. */
void pajack( int j, int m, int n, const double *r, double *x )
{
  int i, k, nb, jnb;

  nb = n/m;
  jnb = j*nb;
  for (i = 0; i < n-nb; ++i) {
    k = i;
    if (k > jnb)
      k += nb;
    x[i] = r[k];
  }
}

/* allocate memory for the arrays used by pafit. */ 
struct Pafit* paalloc( int km, int mm )
{
  int i;
  struct Pafit *pa;

  pa = malloc(sizeof(struct Pafit));
  pa->k = 0;
  pa->a = (double*)malloc(sizeof(double)*(km+1));
  pa->b = (double*)malloc(sizeof(double)*km);
  pa->c = (double*)malloc(sizeof(double)*km);
  pa->d = (double**)malloc(sizeof(double*)*(km+1));
  pa->f = (double*)malloc(sizeof(double)*km);
  pa->g = (double*)malloc(sizeof(double)*km);
  for (i = 0; i < km; ++i)
    (pa->d)[i] = (double*)malloc(sizeof(double)*(mm+1));
  (pa->d)[km] = 0;
  
  return pa;
}

/* release the memory allocated by paalloc. */
void pafree( struct Pafit *pa )
{
  int i;

  for (i = 0; pa->d[i]; ++i)
    free(pa->d[i]);
  free(pa->a);
  free(pa->b);
  free(pa->c);
  free(pa->d);
  free(pa->f);
  free(pa->g);
  free(pa);
}

/* computes the parameters of the piecewise analytic fit to a
   probability distrubution from data. */
double pafit( int n, const double *r, double qm, int km, int mm,
              int nodiscont, struct Pafit *pa )
{
  int     i, j , m , mu, nu, imax, nToDo; 
  double  q, mp, xs, x1, x2, s1, s2, c1, c2, pr, pl, dp, qAim;
  double  mpx, smpx, cmpx, mpx2, dx, sdx, cdx;
  int     *index, *ibegin, *iend, *np, *toDo, *done;
  double  *qcut, *F, *z, *dummy;

  index = (int*)malloc(km*sizeof(int));
  ibegin = (int*)malloc(km*sizeof(int));
  iend = (int*)malloc(km*sizeof(int));
  np = (int*)malloc(km*sizeof(int));
  toDo = (int*)malloc(km*sizeof(int));
  qcut = (double*)malloc(km*sizeof(double));          
  F = (double*)malloc(n*sizeof(double));
  z = (double*)malloc(n*sizeof(double));
  dummy = (double*)malloc((mm+1)*sizeof(double));
  done = (int*)calloc(km,sizeof(int));
  pa->k = 1; /*start analysis with one interval*/
  ibegin[0] = 0;
  iend[0] = n-1;
  qcut[0] = qm;
  toDo[0] = 0; 
  nToDo = 1;
  while (nToDo) {  
    nu = toDo[--nToDo]; /*get interval index*/
    np[nu] = iend[nu]-ibegin[nu]+1;
    for(i = ibegin[nu];i<= iend[nu];++i) 
      F[i] = z[i] = (r[i]-r[ibegin[nu]])/(r[iend[nu]]-r[ibegin[nu]]);
    q = pa_q(np[nu],F+ibegin[nu],&imax);  
    m = 0; /*start with zero modes*/
    while ( q<qcut[nu] && m<mm ) {   
      ++m; /*add a fourier mode*/
      mpx = 0.0;
      smpx = 0.0;
      cmpx = 1.0;
      mp = m*M_PI;
      x2 = 0.0;
      s2 = 0.0;
      c2 = 1.0;
      pa->d[nu][m] = 0.0;
      for (i = 1; i < iend[nu]-ibegin[nu]; ++i) {
        x1 = x2;
        s1 = s2;
        c1 = c2;
        xs = i/(double)np[nu];
        x2 = z[i+ibegin[nu]];
        mpx2 = mp*x2;
        dx = mpx2-mpx;
        if (dx < 0.018 ) {
          /*compute faster using a series (accurate up to 10^-6)*/
          sdx = dx;
          cdx = 1.0-0.5*dx*dx;
          s2 = smpx*cdx+cmpx*sdx;
          c2 = cmpx*cdx-smpx*sdx;
        } else {
          s2 = sin(mpx2); 
          c2 = cos(mpx2);
          mpx = mpx2;
          smpx = s2;
          cmpx = c2;            
        }
        pa->d[nu][m]+= 2*((xs-x1)*c1-(xs-x2)*c2+(s1-s2)/mp)/mp;
      }
      for (i = ibegin[nu]+1; i < iend[nu]; ++i)
        F[i] += pa->d[nu][m]*sin(mp*z[i]);
      q = pa_q(np[nu], F+ibegin[nu], &imax);
    }
    pa->d[nu][0] = (double)m; /*store number of fourier modes*/
    if (q < qcut[nu] && pa->k < km) { 
      if (r[ibegin[nu]+imax] != r[ibegin[nu]] 
          && r[ibegin[nu]+imax] != r[iend[nu]] ) { 
        /*add an interval*/
        ibegin[pa->k] = ibegin[nu]+imax;
        iend[pa->k] = iend[nu];
        iend[nu] = ibegin[pa->k];
        qcut[pa->k] = (1.0-(double)imax/np[nu])*qcut[nu];
        qcut[nu]*= (double)imax/np[nu];
        toDo[nToDo++] = nu;
        toDo[nToDo++] = pa->k;
        ++(pa->k);
      } else {
        printf("#pafit: no convergence\n");
        printf("#pafit: convergence trouble at %f\n", r[ibegin[nu]+imax] );
      }
    }
  }
  /*sort intervals*/
  for (mu = 0; mu < pa->k; ++mu) 
    index[mu] = mu;
  i = 1;
  while (i<pa->k && ibegin[index[i-1]]<ibegin[index[i]]) 
    ++i;    
  while (i<pa->k) {
    mu = index[i];
    for (j = i; j>0 && ibegin[index[j-1]]>ibegin[mu]; j--)
      index[j] = index[j-1];
    index[j] = mu;
    while ( i<pa->k && ibegin[index[i-1]]<ibegin[index[i]] ) 
      ++i;
  } 
  for (mu = 0; mu < pa->k; ++mu) /*apply permutation in "index"*/
    if ( !done[mu] && index[mu] != mu ) {
      for (i = 0; i <= mm; ++i) 
        dummy[i] = pa->d[mu][i]; 
      for (nu = mu; index[nu] != mu; nu = index[nu]) {
        for(i = 0; i <=  mm; ++i) 
          pa->d[nu][i] = pa->d[index[nu]][i];
        done[nu] = 1;
      } 
      for (i = 0; i <= mm; ++i) 
        pa->d[nu][i] = dummy[i];
      done[nu] = 1;
    }
  free(done);
  free(dummy);
  for (mu = 0; mu < pa->k; ++mu) {
    pa->a[mu] = r[ibegin[index[mu]]];
    pa->b[mu] = 0; 
    pa->c[mu] = 0; 
    pa->f[mu] = (double)np[index[mu]]/n;
    pa->g[mu] = (mu!= 0)?(pa->g[mu-1]+pa->f[mu-1]):0.0;
    F[ibegin[index[mu]]] = 0.0;
    for (i = ibegin[index[mu]]; i < iend[index[mu]]; ++i) 
      F[i] = pa->g[mu]+pa->f[mu]*F[i];   
  }
  pa->a[pa->k] = r[n-1];
  F[n-1] = 1.0;  
  q = pa_q(n,F,&imax); /*final q test*/
  if (nodiscont && pa->k>1 ) { /*smooth out discontinuities*/    
    qAim = pa_min(q,qm);
    for (mu = 1; mu < pa->k; ++mu) {
      pl = 0.0;
      pr = 0.0; 
      m = (int)(pa->d[mu-1][0]);
      for (j = 1; j <= m; ++j) 
        pl+= ((j&1)!= 0?-j:j)*pa->d[mu-1][j];
      m = (int)(pa->d[mu][0]);
      for (j = 1; j <= m; ++j)
        pr+= j*pa->d[mu][j];
      dp = pa->f[mu]*(1+M_PI*pr)/(pa->a[mu+1]-pa->a[mu]) 
        - pa->f[mu-1]*(1+M_PI*pl)/(pa->a[mu]-pa->a[mu-1]);
      pa->c[mu] = 0.5*pa_min(pa->a[mu+1]-pa->a[mu],pa->a[mu]-pa->a[mu-1]);
      pa->b[mu] = 0.25*dp/pa->c[mu];         
    }
    i = 0;
    for (nu = 1; nu < pa->k; ++nu) {
      for( ; r[i] < pa->a[nu]-pa->c[nu]; ++i)
        z[i] = F[i];
      for( ; r[i] < pa->a[nu]; ++i)
        z[i] = F[i]+pa->b[nu]*pa_sqr(r[i]-(pa->a[nu]-pa->c[nu]));
      for( ; r[i] < pa->a[nu]+pa->c[nu]; ++i) 
        z[i] = F[i]+pa->b[nu]*pa_sqr(r[i]-(pa->a[nu]+pa->c[nu]));
    }
    for ( ; i < n; ++i)
      z[i] = F[i];
    q = pa_q(n, z, &imax);        
    while (q < qAim) {
      nu = 0;
      while (pa->a[nu+1] < r[imax])
        ++nu;
      if (nu < (pa->k)-1 && pa->a[nu+1]-r[imax]<pa->c[nu+1])
        ++nu;
      i = 0;
      while (r[i] < pa->a[nu]-pa->c[nu])
        ++i;
      for( ; r[i] < pa->a[nu]+pa->c[nu]; ++i) 
        z[i] = F[i];
      pa->c[nu]/= 2;
      pa->b[nu]*= 2;
      if (pa->c[nu] < (pa->a[nu+1]-pa->a[nu])/np[index[nu]])  
        pa->b[nu] = pa->c[nu] = 0;
      i = 0;
      while (r[i] < pa->a[nu]-pa->c[nu])
        ++i;
      for ( ; r[i] < pa->a[nu]; ++i)
        z[i] = F[i]+pa->b[nu]*pa_sqr(r[i]-(pa->a[nu]-pa->c[nu]));
      for ( ; r[i] < pa->a[nu]+pa->c[nu]; ++i) 
        z[i] = F[i]+pa->b[nu]*pa_sqr(r[i]-(pa->a[nu]+pa->c[nu]));
      q = pa_q(n,z,&imax);
    }
  }
  free(F);
  free(z);
  free(index);
  free(ibegin);
  free(iend);
  free(np);
  free(toDo);
  free(qcut);
  return q;
}

/* evaluate the piecewise analytic probability distribution at a point */
double papd( double r, struct Pafit *pa )
{
  int nu, m ,j;
  double p, patch, z;

  if (r < pa->a[0] || r > pa->a[pa->k] || pa->k == 0) {
    return 0;
  } else {
    nu = 0;
    while (pa->a[nu+1] < r)
      ++nu;
    z = M_PI*(r-pa->a[nu])/(pa->a[nu+1]-pa->a[nu]);
    m = (int)(pa->d[nu][0]);
    p = 0;
    for (j = 1;j<= m;++j)
      p += pa->d[nu][j]*j*M_PI*cos(j*z);
    patch = 0;
    if (nu>0 && r-pa->a[nu]<pa->c[nu])
        patch += 2*pa->b[nu]*(r-(pa->a[nu]+pa->c[nu]));
    if (nu<pa->k-1 && pa->a[nu+1]-r<pa->c[nu+1])
        patch += 2*pa->b[nu+1]*(r-(pa->a[nu+1]-pa->c[nu+1]));
    return pa->f[nu]/(pa->a[nu+1]-pa->a[nu])*(1+p)+patch;
  }
}

/* evaluate the piecewise analytic cumulative distribution. */
double pacd( double r, struct Pafit *pa )
{
  if (r < pa->a[0])
    return 0.0;
  else if (r > pa->a[pa->k])
    return 1.0;
  else {
    int nu, m, j;
    double z, zp, dz, patch;
    nu = 0;
    while(pa->a[nu+1]<r)
      ++nu;
    z = (r-pa->a[nu])/(pa->a[nu+1]-pa->a[nu]);
    zp = z*M_PI;
    m = (int)(pa->d[nu][0]);
    dz = 0.0;
    for(j = 1; j <= m; ++j)
      dz += pa->d[nu][j]*sin(j*zp);
    patch = 0.0;
    if(nu>0 && r-pa->a[nu]<pa->c[nu])
        patch += pa->b[nu]*pa_sqr(r-(pa->a[nu]+pa->c[nu]));
    if(nu<pa->k-1 && pa->a[nu+1]-r<pa->c[nu+1])
        patch += pa->b[nu+1]*pa_sqr(r-(pa->a[nu+1]-pa->c[nu+1]));
    return pa->g[nu]+pa->f[nu]*(z+dz)+patch;
  }
}

/* write the fit parameters found by pafit to a file. */
void pawrite( FILE *file, struct Pafit *pa )
{
  int mu, j;

  fprintf(file,"# k %d\n",pa->k);
  for (mu = 0; mu < pa->k; mu++) {
    fprintf(file,"# a%d %f",mu,pa->a[mu]);
    if (mu && pa->b[mu]!= 0.0)
      fprintf(file," b%d %f c%d %f",mu,pa->b[mu],mu,pa->c[mu]);
    fprintf(file," f%d %f d%d %d ",mu,pa->f[mu],mu,(int)(pa->d[mu][0]));
    for (j = 1; j <= (int)(pa->d[mu][0]); j++)
      fprintf(file,"%f ",pa->d[mu][j]);
    fprintf(file,"\n");
  }
  fprintf(file,"# a%d %f\n",pa->k,pa->a[pa->k]);
  fprintf(file,"# EOP\n");
}