/*

==========================================================================

amcmc.c -- a C program (called from R) for doing adaptive MCMC

Copyright (c) 2007 by Jeffrey S. Rosenthal (probability.ca/jeff)

(Last modified March 2009.)

Instructions for use at:  probability.ca/amcmc

Licensed for general copying, distribution and modification according to
the GNU General Public License (www.gnu.org/copyleft/gpl.html).

==========================================================================

*/


#include <R.h>
#include <Rdefines.h>
#include <math.h>
#include <sys/time.h>

#define TARGACCEPT 0.44
#define MAXK 1000

double drand48();

extern double myfunctional(), mydensity();
extern int mydim();

int cdensityflag;
int cfunctionalflag;
int K;
int vectorlength;

SEXP amcmc(SEXP densfn, SEXP functional, SEXP Karg, SEXP numbatchesarg,
	SEXP batchlengtharg, SEXP logvalarg, SEXP adaptvalarg,
	SEXP burninfracarg, SEXP initvalarg, SEXP verbosearg,
	SEXP cdensityarg, SEXP cfunctionalarg, SEXP vectorlengtharg,
	SEXP writearg, SEXP env)
{
  SEXP Rreturnval;

  double returnval, Revalfn();
  int ifloor(), imax(), divisible();
  double valsum = 0.0;
  int tmpK = ifloor( REAL(Karg)[0] );
  int NUMBATCHES = ifloor( REAL(numbatchesarg)[0] );
  int BATCHLENGTH = ifloor( REAL(batchlengtharg)[0] );
  int ADAPTFLAG = ifloor( REAL(adaptvalarg)[0] );
  int LOGFLAG = ifloor( REAL(logvalarg)[0] );
  int tmpcdensityflag = ifloor( REAL(cdensityarg)[0] );
  int tmpcfunctionalflag = ifloor( REAL(cfunctionalarg)[0] );
  int tmpvectorlength = ifloor( REAL(vectorlengtharg)[0] );
  int writeflag = ifloor( REAL(writearg)[0] );
  double burninfrac = REAL(burninfracarg)[0];
  double initval = REAL(initvalarg)[0];
  int VERBOSELEVEL = ifloor( REAL(verbosearg)[0] );
  int i, j, ii, jj, kk, PRINTLENGTH;

  int batchnum, numit, numcounted, accepted, acount[MAXK];
  double t[MAXK];
  double newt[MAXK];
  double logsigma[MAXK];
  double logalpha, sqnorm(), min(), normal(), targlogdens(), adaptamount();
  double funcest, funcsum, funcval, curlogdens, newlogdens;
  int functionalready;
  double begintime, endtime, currtime, burnintime;
  struct timeval tmptv;
  FILE *fp, *fps;

  cdensityflag = tmpcdensityflag;
  cfunctionalflag = tmpcfunctionalflag;
  vectorlength = tmpvectorlength;
  if (cdensityflag > 0)
    K = mydim();
  else if (vectorlength > 0)
    K = vectorlength;
  else
    K = tmpK;

  PRINTLENGTH = imax(100, NUMBATCHES/1000);
  if (VERBOSELEVEL >= 3)
    PRINTLENGTH = PRINTLENGTH / 10;

  gettimeofday (&tmptv, (struct timezone *)NULL);
  begintime = 1.0 * tmptv.tv_sec + 0.000001 * tmptv.tv_usec;
  /* printf("\nBeginning variable-at-time adaption run.\n"); */

if (VERBOSELEVEL >= 1) {
printf("\nStarting C function 'amcmc'.\nK = %d. NUMBATCHES = %d. BATCHLENGTH = %d. LOGFLAG = %d. ADAPTFLAG = %d.\n", K, NUMBATCHES, BATCHLENGTH, LOGFLAG, ADAPTFLAG);
printf("burninfrac = %f.  cdensityflag = %d.  cfunctionalflag = %d.\n",
			burninfrac, cdensityflag, cfunctionalflag);
printf("\n");
}

    /* INITIALISATIONS. */
    numit = numcounted = 0;
    for (kk=0; kk<K; kk++) {
        t[kk] = initval;
        logsigma[kk] = 0.0;
        newt[kk] = t[kk];
    }
    curlogdens = targlogdens(t, densfn, LOGFLAG, env);
    funcval = funcest = funcsum = 0.0;
    functionalready = 0;
    burnintime = NUMBATCHES * BATCHLENGTH * burninfrac;
    if (writeflag) {
        if ( ((fp = fopen("amcmcvals","w")) == NULL) ||
             ((fps = fopen("amcmcsigmas","w")) == NULL) ) {
          fprintf(stderr, "Unable to open file to write values.\n");
	  writeflag = 0;  /* continue, but don't write anything */
        }
	fprintf(fp, "\namcmcvals = matrix( c( ");
	fprintf(fps, "\namcmcsigmas = matrix( c( ");
	for (kk=0; kk<K; kk++) {
	    fprintf(fp, "%d, ", t[kk]);
	    fprintf(fps, "%d, ", logsigma[kk]);
	}
        fprintf(fp, "\n");
        fprintf(fps, "\n");
    }

    /* MAIN ITERATIVE LOOP. */
    for (batchnum=1; batchnum<=NUMBATCHES; batchnum++) {
    
        for (kk=0; kk<K; kk++) {
          acount[kk] = 0;
          if (abs(t[kk]-newt[kk]) > 0.01) {
            fprintf(stderr, "\nERROR: t and newt not synchronised.\n\n");
            exit(1);
          }
        }
    
        for (ii=1; ii<=BATCHLENGTH; ii++) {

          numit++; 
          /* funcval = log(1.0+sqnorm(t)); */
    
          for (kk=0; kk<K; kk++) {

            /* PROPOSE MOVE OF kk'th COORD. */
            newt[kk] = t[kk] + exp(logsigma[kk]) * normal();
            newlogdens = targlogdens(newt, densfn, LOGFLAG, env);
            /* logalpha = targlogdens(newt, densfn, LOGFLAG, env)
				- targlogdens(t, densfn, LOGFLAG, env); */
            logalpha = newlogdens - curlogdens;
if (VERBOSELEVEL >= 4)
printf("Coord %d from %f to %f (ls=%f, la=%f) is ... ",
kk, t[kk], newt[kk], logsigma[kk], logalpha);
            accepted = ( log(drand48()) < logalpha );
            if (accepted) {
              t[kk] = newt[kk];
              acount[kk] = acount[kk] + 1;
	      curlogdens = newlogdens;
	      functionalready = 0;
if (VERBOSELEVEL >= 4)
printf("accepted!\n");
            } else {
              newt[kk] = t[kk];
if (VERBOSELEVEL >= 4)
printf("rejected!\n");
            }

          } /* End of kk for loop. */

          /* funcval = log(1.0+sqnorm(t)); */

/* printf("t[0] = %f ... in C: %f ... in R: %f\n", t[0], myfunctional(t),
Revalfn(functional, t, env) ); */

	  if (numit > burnintime) {
	    if (functionalready==0) {
	      /* Need to re-evaluate functional. */
              if (cfunctionalflag==1)
                funcval = myfunctional(t);
              else
                funcval = Revalfn(functional, t, env);
	      functionalready = 1;
	    }
if (VERBOSELEVEL >= 5)
printf("EVALUATING: t[0] = %f, funcval = %f\n", t[0], funcval);
            funcsum = funcsum + funcval;
	    numcounted++;
	  }
  
        } /* End of ii for loop. */

        /* Update various estimates etc. */
        funcest = funcsum / numcounted;

        /* DO THE ADAPTING. */

if(ADAPTFLAG > 0) {
        for (jj=0; jj<K; jj++) {
            if (acount[jj] > BATCHLENGTH * TARGACCEPT)
              logsigma[jj] = logsigma[jj] + adaptamount(batchnum);
            else if (acount[jj] < BATCHLENGTH * TARGACCEPT)
              logsigma[jj] = logsigma[jj] - adaptamount(batchnum);
        }
}

      /* OUTPUT VALUES TO SCREEN. */
if (VERBOSELEVEL >= 2) {
      if (divisible(batchnum, PRINTLENGTH)) {
        for(jj=0; jj<K; jj++) {
          printf("t[%d]=%.2f, ls[%d]=%.2f; ", jj, t[jj], jj, logsigma[jj]);
        }
        printf("sq=%f, funcest=%f ", sqnorm(t), funcest );
        gettimeofday (&tmptv, (struct timezone *)NULL);
        currtime = 1.0 * tmptv.tv_sec + 0.000001 * tmptv.tv_usec;
        printf("(%d/%d, %d, %.2f secs)\n", batchnum, NUMBATCHES, numcounted,
                                                currtime-begintime);
     }
}

     /* WRITE VALUES TO FILE. */
     if (writeflag) {
	for (jj=0; jj<K; jj++) {
	    fprintf(fp, "%f", t[jj]);
	    fprintf(fps, "%f", logsigma[jj]);
	    if ( (jj < K-1) || (batchnum < NUMBATCHES) ) {
	      fprintf(fp, ", ");
	      fprintf(fps, ", ");
	    }
	}
        fprintf(fp, "\n");
        fprintf(fps, "\n");
    }

} /* End of batchnum for loop. */

gettimeofday (&tmptv, (struct timezone *)NULL);
endtime = 1.0 * tmptv.tv_sec + 0.000001 * tmptv.tv_usec;

if (VERBOSELEVEL >= 1) {
  printf("\nFinishing C function 'amcmc'.  Number of variables = %d.\n", K);
  /* printf("Batch size = %d full scans.\n", BATCHLENGTH); */
  /* printf("Printed to screen every %d batches.\n", PRINTLENGTH); */
  printf("Did %d batches, each consisting of %d updates of each variable.\n", NUMBATCHES, BATCHLENGTH);
  printf("Ellapsed time = %.2f seconds.\n\n", endtime-begintime);
}
  if (writeflag) {
    fprintf(fp, "),\nnrow=%d)\n\n", K);
    fclose(fp);
    fprintf(fps, "),\nnrow=%d)\n\n", K);
    fclose(fps);
  }

  returnval = funcest;

  /* return(returnval) */
  PROTECT(Rreturnval = NEW_NUMERIC(1));
  NUMERIC_POINTER(Rreturnval)[0] = returnval;
  UNPROTECT(1);
  return( Rreturnval );
}


/* ADAPTAMOUNT */
double adaptamount( int iteration )
{
    double min();

    /* return( 1.0 / imax(100, sqrt(iteration)) ); */
    return( min( 0.01, 1.0 / sqrt((double)iteration) ) );
    /* return(0.01); */
}


double Revalfn(SEXP fn, double *arg, SEXP env)
{
  SEXP R_fcall, argg, tmpval[MAXK], tmpsex;
  double m, *p_tmp[MAXK];
  int iii;

  if (vectorlength > 0) {

    PROTECT(argg = allocVector(REALSXP, K));
    for (iii=0; iii<K; iii++)
      REAL(argg)[iii] = arg[iii];
    PROTECT(R_fcall = lang2(fn, argg));
    PROTECT(tmpsex = coerceVector( eval(R_fcall, env), REALSXP ) );
    m = REAL(tmpsex)[0];
    UNPROTECT(3);

  } else {

    PROTECT(argg = allocVector(VECSXP, K));
    for (iii=0; iii<K; iii++) {
      PROTECT(tmpval[iii] = NEW_NUMERIC(1));
      p_tmp[iii] = NUMERIC_POINTER(tmpval[iii]);
      p_tmp[iii][0] = arg[iii];
      SET_VECTOR_ELT(argg, iii, tmpval[iii]);
    }
    PROTECT(R_fcall = lang2(fn, argg));
    PROTECT(tmpsex = coerceVector( eval(R_fcall, env), REALSXP ) );
    m = REAL(tmpsex)[0];
    UNPROTECT(K+3);

  }

  return(m);
}


/* IFLOOR */
int ifloor(double x)  /* returns floor(x) as an integer */
{
    return((int)floor(x));
}


/* MIN */
double min( double xx, double yy)
{
  if (xx<yy)
    return(xx);
  else
    return(yy);
}


/* NORMAL:  return a standard normal random number. */
double normal()
{
    double R, theta, drand48();

    R = - log(drand48());
    theta = 2 * PI * drand48();

    return( sqrt(2*R) * cos(theta));
}


/* SEEDRAND: SEED RANDOM NUMBER GENERATOR. */
void seedrand()
{
    int i, seed;
    struct timeval tmptv;
    gettimeofday (&tmptv, (struct timezone *)NULL);
    /* seed = (int) (tmptv.tv_usec - 1000000 *
                (int) ( ((double)tmptv.tv_usec) / 1000000.0 ) ); */
    seed = (int) tmptv.tv_usec;
    srand48(seed);
    (void)drand48();  /* Spin it once. */
    /* return(0); */
}


/* TARGLOGDENS */
double targlogdens( double t[], SEXP densfn, int thelogflag, SEXP env )
{
    double Revalfn(), tmpval;

    if (cdensityflag==1)
	tmpval = mydensity(t);
    else
	tmpval = Revalfn(densfn, t, env);

    if (thelogflag > 0)
      return( tmpval );
    else
      return( log( tmpval ) );
}


/* SQNORM */
double sqnorm( double t[] )
{

    int ii;
    double ss;
    ss = 0.0;
    for (ii=0; ii<K; ii++)
        ss = ss + t[ii]*t[ii];
    return(ss);

    /* return( A ); */
    /* return( (t[0]-mu0)*(t[0]-mu0) ); */
}


int divisible (int n1, int n2)
{
  return (n1 == n2 * (n1/n2));
}


/* IMAX */
int imax( int n1, int n2)
{
  if (n1>n2)
    return(n1);
  else
    return(n2);
}