/******************************************************************************** 
 *
 * Bayesian Regression and Adaptive Sampling with Gaussian Process Trees
 * Copyright (C) 2005, University of California
 * 
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 * 
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 * 
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
 *
 * Questions? Contact Robert B. Gramacy (rbgramacy@ams.ucsc.edu)
 *
 ********************************************************************************/


extern "C"
{
#include "matrix.h"
#include "rand_draws.h"  
#include "rhelp.h"
}
#include "tgp.h"
#include "model.h"
#include "params.h"
#include <stdlib.h>
#include <string.h>
#include <assert.h>
#include <fstream>
#include <time.h>

extern "C"
{

Tgp* tgpm = NULL;
void *tgp_state = NULL;

void tgp(int* state_in, 
	 double *X_in, int *n_in, int *d_in, double *Z_in, double *XX_in, int *nn_in,
	 int *BTE, int* R_in, int* linburn_in, double *params_in, int *printlev_in,
	 double *Zp_mean_out, double *ZZ_mean_out, double *Zp_q_out, double *ZZ_q_out,
	 double *Zp_q1_out, double *Zp_median_out, double *Zp_q2_out,
	 double *ZZ_q1_out, double *ZZ_median_out, double *ZZ_q2_out,
	 double *Ds2x_out, double *ego_out)
{
  BFILE = BETAFILE = NULL;

  /* create the RNG state */
  void *tgp_state = 
    newRNGstate((unsigned long) (state_in[0] * 100000 + state_in[1] * 100 + state_in[2]));

  /* copy the input parameters to the tgp class object where all the MCMC 
     work gets done */
  tgpm = new Tgp(tgp_state, *n_in, *d_in, *nn_in, BTE[0], BTE[1], BTE[2], *R_in, 
		 *linburn_in, (bool) (Zp_mean_out!=NULL), (bool) (Ds2x_out!=NULL), 
		 (bool) (ego_out != NULL), X_in, Z_in, XX_in, params_in, *printlev_in);

  /* maybe print the booleans and betas out to a file */
  if(bprint) { 
    BFILE = fopen("b.out", "w");
    BETAFILE = fopen("beta.out", "w");
  }

  /* tgp MCMC rounds are done here */
  tgpm->Rounds();

  /* gather the posterior predictive statistics from the MCMC rounds */
  tgpm->GetStats(Zp_mean_out, ZZ_mean_out, Zp_q_out, ZZ_q_out, Zp_q1_out, Zp_median_out, 
		 Zp_q2_out, ZZ_q1_out, ZZ_median_out, ZZ_q2_out, Ds2x_out, ego_out);

  /* delete the tgp model */
  delete tgpm; tgpm = NULL;
  
  /* close the beta summary files */

  if(bprint) {
    fclose(BFILE); BFILE = NULL;
    fclose(BETAFILE);  BETAFILE = NULL;
  }

  /* destroy the RNG */
  deleteRNGstate(tgp_state);
  tgp_state = NULL;

  /* free blank line before returning to R prompt */
  if(*printlev_in >= 1) myprintf(stdout, "\n");
}


/*
 * Tgp: (constructor) 
 *
 * copies the input passed to the tgp function from R via
 * .C("tgp", ..., PACKAGE="tgp").  Then, it calls the init
 * function in order to get everything ready for MCMC rounds.
 */

Tgp::Tgp(void *state, int n, int d, int nn, int B, int T, int E, int R, int linburn, 
	 bool pred_n, bool delta_s2, bool ego, double *X, double *Z, double *XX, 
	 double *dparams, int verb)
{
  itime = time(NULL);
	
  this->state = NULL;
  this->X = this->XX = NULL;
  this->rect = NULL;
  this->Z = NULL;
  params = NULL;
  model = NULL;
  cumpreds = preds = NULL;

  /* RNG state */
  this->state = state;

  /* integral dimension parameters */
  this->n = (unsigned int) n;
  this->d = (unsigned int) d;
  this->nn = (unsigned int) nn;
  this->B = B;
  this->T = T;
  this->E = E;
  this->R = R;
  this->linburn = linburn;
  this->pred_n = pred_n;
  this->delta_s2 = delta_s2;
  this->ego = ego;
  this->verb = verb;

  /* copy X from input */
  this->X = new_matrix(n, d);
  dupv(this->X[0], X, n*d);
  
  /* copy Z from input */
  this->Z = new_dup_vector(Z, n);
  
  /* copy X from input */
  this->XX = new_matrix(nn, d);
  if(this->XX) dupv(this->XX[0], XX, nn*d);
  
  /* use default parameters */
  params = new Params(d);
  if((int) dparams[0] != -1) params->read_double(dparams);
  else myprintf(stdout, "Using default params.\n");

  Init();
}


/*
 * ~Tgp: (destructor)
 *
 * typical destructur function.  Checks to see if the class objects
 * are NULL first becuase this might be called from within 
 * tgp_cleanup if tgp was interrupted during computation
 */

Tgp::~Tgp(void)
{
  /* clean up */
  if(model) { delete model; model = NULL; }
  if(params) { delete params; params = NULL; }
  if(XX) { delete_matrix(XX);  XX = NULL; }
  if(Z) { free(Z); Z = NULL; }
  
  /* clean up */
  if(rect) { delete_matrix(rect); rect = NULL; }
  if(X) { delete_matrix(X); X = NULL; }
  if(cumpreds) { delete_preds(cumpreds); }
}


/*
 * Init:
 *
 * get everything ready for MCMC rounds -- called from within the
 * the Tgp constructor function, in order to separate the copying
 * of the input parameters from the initialization of the model
 * and predictive data.
 */

void Tgp::Init(void)
{
  /* get  the rectangle */
  rect = getXdataRect(X, n, d, XX, nn);
  
  /* construct the new model */
  model = new Model(params, d, rect, 0, state);
  model->Init(X, n, d, Z);
  model->Outfile(stdout, verb);
  
  /* structure for accumulating predictive information */
  cumpreds = new_preds(XX, nn, pred_n*n, d, rect, R*(T-B), delta_s2, ego, E);

  /* print the parameters of this module */
  if(verb >= 2) Print(stdout);
}  


/*
 * Rounds:
 *
 * Actually do the MCMC for sampling from the posterior of the tgp model
 * based on the parameterization given to the Tgp constructor.  
 */

void Tgp::Rounds(void)
{

  for(unsigned int i=0; i<R; i++) {

    itime = my_r_process_events(itime);
    
     /* Linear Model Initialization rounds -B thru 1 */
    if(linburn) model->Linburn(B, state);

    /* do model rounds 1 thru B (burn in) */
    model->Burnin(B, state);
	
    /* do the MCMC rounds B,...,T */
    preds = new_preds(XX, nn, pred_n*n, d, rect, T-B, delta_s2, ego, E);
    model->Sample(preds, T-B, state);

    /* print tree statistics */
    if(verb >= 1) model->PrintTreeStats(stdout);

    /* accumulate predictive information */
    import_preds(cumpreds, preds->R * i, preds);		
    delete_preds(preds); preds = NULL;

    /* done with this repetition; prune the tree all the way back */
    if(R > 1) {
      myprintf(stdout, "finished repetition %d of %d\n", i+1, R);
      model->cut_root();
    }
  }

  /* cap of the printing */
  if(verb >= 1) myflush(stdout);

  /* these might not do anything, if they're turned off */
  model->print_linarea();
  model->PrintPosteriors();
}


/*
 * GetStats:
 *
 * Coalate the statistics from the samples of the posterior predictive
 * distribution gathered during the MCMC Tgp::Rounds() function.
 */

void Tgp::GetStats(double *Zp_mean, double *ZZ_mean, double *Zp_q, double *ZZ_q,
	      double *Zp_q1, double *Zp_median, double *Zp_q2, double *ZZ_q1, 
	      double *ZZ_median, double *ZZ_q2, double *Ds2x, double *ego)
{
  itime = my_r_process_events(itime);

  if(pred_n) { /* calculate means and quantiles */
    mean_of_columns(Zp_mean, cumpreds->Zp, cumpreds->R, n);
    qsummary(Zp_q, Zp_q1, Zp_median, Zp_q2, cumpreds->Zp, cumpreds->R, n);
  }
  if(nn > 0) { /* predictive data locations (XX) */
    mean_of_columns(ZZ_mean, cumpreds->ZZ, cumpreds->R, nn);
    qsummary(ZZ_q, ZZ_q1, ZZ_median, ZZ_q2, cumpreds->ZZ, cumpreds->R, cumpreds->nn);
    /* warning: this makes a permanent change to cumpreds->Ds2xy, should change this! */
    if(cumpreds->Ds2xy) norm_Ds2xy(cumpreds->Ds2xy, cumpreds->R, cumpreds->nn);
    if(delta_s2) {
      assert(cumpreds->Ds2xy);
      mean_of_rows(Ds2x, cumpreds->Ds2xy, nn, nn);
    }
    if(ego) {
      scalev(cumpreds->ego, cumpreds->nn, 1.0/cumpreds->R);
      dupv(ego, cumpreds->ego, nn);
    }
  }
}


/*
 * tgp_cleanup
 *
 * function for freeing memory when tgp is interrupted
 * by R, so that there won't be a (big) memory leak.  It frees
 * the major chunks of memory, but does not guarentee to 
 * free up everything
 */

void tgp_cleanup(void)
{
  if(tgpm) { 
    delete tgpm; 
    tgpm = NULL; 
    myprintf(stderr, "INTERRUPT: tgp model leaked, is now destroyed\n");
  }

  if(tgp_state) {
    deleteRNGstate(tgp_state);
    tgp_state = NULL;
    myprintf(stderr, "INTERRUPT: tgp RNG leaked, is now destroyed\n");
  }

  if(BFILE) {
    fclose(BFILE);
    BFILE = NULL;
    myprintf(stderr, "INTERRUPT: tgp BFILE leaked, is now destroyed\n");
  }
  if(BETAFILE) {
    fclose(BETAFILE);
    BETAFILE = NULL;
    myprintf(stderr, "INTERRUPT: tgp BETAFILE leaked, is now destroyed\n");
  }
}

} /* extern "C" */


/* 
 * getXdataRect:
 * 
 * given the data Xall (Nxd), infer the rectancle
 * from IFace class
 */

double ** getXdataRect(double **X, unsigned int n, unsigned int d, double **XX, unsigned int nn)
{
  unsigned int N = nn+n;
  double **Xall = new_matrix(N, d);
  dupv(Xall[0], X[0], n*d);
  if(nn > 0) dupv(Xall[n], XX[0], nn*d);
  
  double **rect = get_data_rect(Xall, N, d);
  delete_matrix(Xall);
  
  return rect;
}


/* 
 * Print:
 *
 * print the settings of the parameters used by this module:
 * which basically summarize the data and MCMC-related inputs
 * followed by a call to the model Print function
 */

void Tgp::Print(FILE *outfile)
{
  myprintf(stdout, "\n");

  /* DEBUG: print the input parameters */
  myprintf(stdout, "n=%d, d=%d, nn=%d\nBTE=(%d,%d,%d), R=%d, linburn=%d\n", 
	   n, d, nn, B, T, E, R, linburn);

  /* print the random number generator state */
  printRNGstate(state, stdout);

  /* print predictive statistic types */
  if(pred_n || delta_s2 || ego) myprintf(stdout, "preds:");
  if(pred_n) myprintf(stdout, " data");
  if(delta_s2) myprintf(stdout, " ALC");
  if(ego) myprintf(stdout, " EGO");
  if(pred_n || delta_s2 || ego) myprintf(stdout, "\n");
  myflush(stdout);

  /* print the model, uses the internal model 
     printing variable OUTFILE */
  model->Print();
}