/*
 * Copyright (C) 1996-2012 Daniel Waggoner and Tao Zha
 *
 * This free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * It 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 General Public License for more details.
 *
 * If you did not received a copy of the GNU General Public License
 * with this software, see <http://www.gnu.org/licenses/>.
 */

#include "dw_switch.h"
#include "dw_switchio.h"
#include "dw_metropolis_theta.h"
#include "dw_rand.h"
#include "dw_error.h"
#include "dw_ascii.h"
#include "dw_parse_cmd.h"
#include "dw_array.h"
#include "dw_metropolis_simulation.h"
#include "dw_metropolis_simulation_prior.h"
#include "dw_std.h"

#include <time.h>
#include <string.h>
#include <ctype.h>
#include <stdlib.h>
#include <limits.h>
#include <math.h>

#include <time.h>

/* Undefine to use wall time instead of process time */
//#define WALL_TIME

//=======================

int SimulatePrior(int burn_in, int draws, int thin, int period, TStateModel *model, FILE *f_out)
{
#ifdef WALL_TIME
  time_t time;
#else
   clock_t begin_time;
#endif
  int i, j, check, ntheta, nq;
  TVector theta, q;

  // Initialization
  ntheta=NumberFreeParametersTheta(model);
  nq=NumberFreeParametersQ(model);
  if (ntheta) theta=CreateVector(ntheta);
  if (nq) q=CreateVector(nq);

  // Burn-in
  ResetCounts_metropolis_theta(model->metropolis_theta);
  ResetMetropolisHastingsCounts_q(model); 
  if (burn_in > 0)
    { 
  printf("(0 / %d)   Beginning burn in.\n",burn_in);
#ifdef WALL_TIME
      begin_time=time((time_t*)NULL);
#else
      begin_time=clock();
#endif
      for (check=period, i=1; i <= burn_in; i++)
	{
	  draw_prior_all_metropolis(model);

	  if (i == check)
	    {
	      check+=period;
#ifdef WALL_TIME
	      printf("(%d / %d)  Total Elapsed Time: %d seconds\n",i,burn_in,(int)(time((time_t*)NULL) - begin_time));
#else
	      printf("(%d / %d)  Total Elapsed Time: %lf seconds\n",i,burn_in,((double)(clock() - begin_time))/CLOCKS_PER_SEC);
#endif
	    }
	}
#ifdef WALL_TIME
      printf("(%d / %d)   Elapsed Time: %d seconds\n",burn_in,burn_in,(int)(time((time_t*)NULL) - begin_time));
#else
      printf("(%d / %d)   Elapsed Time: %lf seconds\n",burn_in,burn_in,((double)(clock() - begin_time))/CLOCKS_PER_SEC);
#endif
    }


  // Simulate
  if (draws > 0)
    { 
      printf("(0 / %d)   Beginning simulation.\n",draws);
#ifdef WALL_TIME
      begin_time=time((time_t*)NULL);
#else
      begin_time=clock();
#endif
      for (check=period, i=1; i <= draws; i++)
	{
	  for (j=thin; j > 0; j--) draw_prior_all_metropolis(model);

	  fprintf(f_out,"%le %le",LogPosterior_StatesIntegratedOut(model),LogLikelihood_StatesIntegratedOut(model));
	  if (ntheta) 
	    {
	      ConvertThetaToFreeParameters(model,pElementV(theta));
	      for (j=0; j < ntheta; j++) fprintf(f_out," %le",ElementV(theta,j));
	    }
	  if (nq)
	    {
	      ConvertQToFreeParameters(model,pElementV(q));
	      for (j=0; j < nq; j++) fprintf(f_out," %le",ElementV(q,j));
	    }
	  fprintf(f_out,"\n");

	  if (i == check)
	    {
	      check+=period;
#ifdef WALL_TIME
	      printf("(%d / %d)  Total Elapsed Time: %d seconds\n",i,draws,(int)(time((time_t*)NULL) - begin_time));
#else
	      printf("(%d / %d)  Total Elapsed Time: %lf seconds\n",i,draws,((double)(clock() - begin_time))/CLOCKS_PER_SEC);
#endif
	    }
	}
#ifdef WALL_TIME
      printf("(%d / %d)   Elapsed Time: %d seconds\n",draws,draws,(int)(time((time_t*)NULL) - begin_time));
#else
      printf("(%d / %d)   Elapsed Time: %lf seconds\n",draws,draws,((double)(clock() - begin_time))/CLOCKS_PER_SEC);
#endif
    }

  // Clean up
  if (ntheta) FreeVector(theta);
  if (nq) FreeVector(q);
  return 0;
}

/*
   Returns 1 if metropolis_theta parameters properly set and 0 otherwise.  See dw_simulate_prior_command_line for 
   command line options.
*/
int SetUpMetropolisSimulationPrior(int n_args, char **args, char *tag, TStateModel *model, int default_period)
{  
  TMetropolis_theta *metropolis=model->metropolis_theta;
  FILE *f_in;
  int i, ndraws, ntheta, nq, period_s, max_period_s, period_a, max_period_a, 
    begin_time, end_time, free_filename_metropolis=0, free_filename_hessian=0;
  PRECISION center_s, center_a;
  char *filename_metropolis, *id, *fmt;
  TMatrix hessian=(TMatrix)NULL, hessian_theta=(TMatrix)NULL, posterior_draws=(TMatrix)NULL,
    variance=(TMatrix)NULL, X=(TMatrix)NULL;
  TVector mean=(TVector)NULL, x=(TVector)NULL, theta=(TVector)NULL;

  // Get Sizes
  ntheta=NumberFreeParametersTheta(model);
  nq=NumberFreeParametersQ(model);

  // Input/output file name for metropolis_theta info
  if (!(filename_metropolis=dw_ParseString_String(n_args,args,"af",(char*)NULL)))
    {
      if (!tag)
	{
	  printf("No adaptive filename (-af) or tag (-t) specified\n");
	  return 0;
	}
      fmt="metropolis_theta_prior_%s.out";
      sprintf(filename_metropolis=(char*)dw_malloc(strlen(fmt) + strlen(tag) - 1),fmt,tag);
      free_filename_metropolis=1;
    }

  // First try: read info from adpative file if -diagonal is not specified
  if (dw_FindArgument_String(n_args,args,"diagonal") == -1)
    if (Read_metropolis_theta((FILE*)NULL,filename_metropolis,(char*)NULL,metropolis))
      {
	// Cleanup and return
	if (free_filename_metropolis) dw_free(filename_metropolis);
	return READ_METROPOLIS_INFO;
      }
    else
      printf("Unable to open %s.  Attempting calibration\n",filename_metropolis);
  else 
    printf("Starting recalibration\n");

  // Setup calibration parameters
  center_s=dw_ParseFloating_String(n_args,args,"cs",0.20);
  //center_a=dw_ParseFloating_String(n_args,args,"ca",0.25);
  period_s=dw_ParseInteger_String(n_args,args,"ps",default_period);
  //period_a=dw_ParseInteger_String(n_args,args,"pa",default_period);
  //max_period_s=dw_ParseInteger_String(n_args,args,"mps",8*period_s);
  max_period_s=dw_ParseInteger_String(n_args,args,"mps",16*period_s);
  //max_period_a=dw_ParseInteger_String(n_args,args,"mpa",16*period_a);

  // Second try: Use diagonal
  // Initializing scale and directions
  Setup_metropolis_theta_diagonal(metropolis);
  Setup_metropolis_theta_blocks_full(metropolis);
 
  // Calibrating
  begin_time=(int)time((time_t*)NULL);
  Calibrate_metropolis_theta_prior(model,center_s,period_s,max_period_s,1);
  //ndraws=Calibrate_metropolis_theta_prior_two_pass(model,center_s,period_s,max_period_s,center_a,period_a,max_period_a,1);
  end_time=(int)time((time_t*)NULL);

  // Print Metropolis info
  WriteMetropolisThetaInfo(filename_metropolis,tag,DIAGONAL_CALIBRATION | PRIOR_CALIBRATION,center_s,period_s,max_period_s,
			   center_a,period_a,max_period_a,ndraws,end_time-begin_time,model);

  // Clean up and return
  if (free_filename_metropolis) dw_free(filename_metropolis);
  return ((dw_FindArgument_String(n_args,args,"special") != -1) ? SPECIAL_CALIBRATION : DIAGONAL_CALIBRATION) | TWO_PASS;
}

//****************************************************************************//
//**************************** Simulation Main  *****************************//
//****************************************************************************//
/*
   Command line parameters

    General parameters
     -t : filename tag
     -ft : filename tag - supersedes -t
     -outtag : filename tag - supersedes -t and -ft

    Adaptive jump ratio parameters
     -cs <real>     : target acceptance rate for single variable calibrations (default = 0.20)
     -ca <real>     : target acceptance rate for overall calibrations (default = 0.25)
     -ps <integer>  : the initial period for the single variable calibrations (default = 20)
     -pa <integer>  : the initial period for the overall calibrations (default = 20)
     -mps <integer> : the maximum period for the single variable calibrations (default = 8*ps)
     -mpa <integer> : the maximum period for the overall calibrations (default = 16*mps)
     -diagonal      : Forces calibration using two-pass diagonal, even if calibration has alread been performed.
     -af <filename> : file for storing the calibrated metropolis parameters. (default = metropolis_theta_<tag>.out)
     -hf <filename> : file containing the hessian matrix. (default = outdataout_<tag>.prn)

    Simulation parameters
     -ndraws : number of draws to save (default = 1000)
     -burnin : number of burn-in draws (default = 0.1 * ndraws)
     -thin : thinning factor.  Total number of draws made is thin*ndraws + burnin (default = 1)
     -of <filename> : output filename (default = simulation_prior_<tag>.out)
   
   Attempts to setup metropolis parameters from command line for simulating theta from the prior. Note if -diagonal
   is not specified, the default behavior is to first look for calibration file and it is not present to use the 
   two-pass diagonal method.
*/
void dw_simulate_prior_command_line(int n_args, char **args, TStateModel *model, TVector start_value, int generator_init)
{
  FILE *f_out;
  char *tag, *default_filename, *filename, *fmt;
  int ndraws, burn_in, thin, default_period=20, n_theta=NumberFreeParametersTheta(model), n_q=NumberFreeParametersQ(model);

  // Initialize generator
  dw_initialize_generator(generator_init);

  // Get tag from command line
  tag=dw_ParseString(n_args,args,'t',(char*)NULL);
  tag=dw_ParseString_String(n_args,args,"ft",tag);
  tag=dw_ParseString_String(n_args,args,"outtag",tag);

  // Push starting value into model 
  if (start_value)
    {
      if (n_theta) ConvertFreeParametersToTheta(model,pElementV(start_value));
      if (n_q) ConvertFreeParametersToQ(model,pElementV(start_value)+n_theta);
    }

  if (SetUpMetropolisSimulationPrior(n_args,args,tag,model,default_period))
    {
      // simulation parameters
      ndraws=dw_ParseInteger_String(n_args,args,"ndraws",1000);
      burn_in=dw_ParseInteger_String(n_args,args,"burnin",ndraws/10);
      thin=dw_ParseInteger_String(n_args,args,"thin",1);

      // Starting Simulation
      fmt="simulation_prior_%s.out";
      sprintf(default_filename=(char*)dw_malloc(strlen(fmt) + strlen(tag) - 1),fmt,tag);
      filename=dw_ParseString_String(n_args,args,"fo",default_filename);
      f_out=fopen(filename,"wt");
      dw_free(default_filename);

      SimulatePrior(burn_in,ndraws,thin,1000,model,f_out);

      if (f_out) fclose(f_out);
    }
  else
    {
      printf("Error setting up Metropolis adaptive simulation\n");
    }
}
