/*
 * Copyright (C) 1996-2011 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_matrix.h"
#include "dw_rand.h"
#include "dw_parse_cmd.h"
#include "dw_ascii.h"
#include "dw_error.h"
#include "VARbase.h"
#include "VARio.h"
#include "switch.h"
#include "switchio.h"
#include "dw_sbvar_command_line.h"
#include "dw_std.h"

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

/*
   Returns the largest of the absolute values of the eigenvalues of the companion
    matrix of B
*/
__attribute__((unused))
static PRECISION LargestRoot(TMatrix B)
{
  TMatrix X;
  TVector re_eig, im_eig;
  int i, k=RowM(B), n=ColM(B);
  PRECISION max, tmp;
  InitializeMatrix(X=CreateMatrix(k,k),0.0);
  for (i=n; i < k; i++) ElementM(X,i-n,i)=1.0;
  InsertSubMatrix(X,B,0,0,0,0,k,n);
  re_eig=CreateVector(k);
  im_eig=CreateVector(k);
  Eigenvalues(re_eig,im_eig,X);
  max=ElementV(re_eig,0)*ElementV(re_eig,0)+ElementV(im_eig,0)*ElementV(im_eig,0);
  for (i=k-1; i > 0; i--)
    if ((tmp=ElementV(re_eig,i)*ElementV(re_eig,i)+ElementV(im_eig,i)*ElementV(im_eig,i)) > max) max=tmp;;
  FreeVector(im_eig);
  FreeVector(re_eig);
  FreeMatrix(X);
  return sqrt(max);
}

/*
   Computes the matrices A0(s)*Xi(s) and B(s) = Ahat(s) * inverse(A0(s))
   where Ahat(s) are the first nvars*nlags rows of Aplus(s), for 0 <= s < h.
   The integer h is the number of regimes.  Both A0_Xi and B must be matrix
   arrays of length h and each element of A0_Xi and B must be of the proper
   size.
*/
void GetIRMatrices(int h, TMatrix *A0_Xi, TMatrix *B, T_VAR_Parameters *p)
{
  TMatrix Aplus;
  int i, j, k;
  PRECISION tmp;
  for (k=h-1; k >= 0; k--)
    {
      MakeA0(A0_Xi[k],k,p);
      Aplus=MakeAplus((TMatrix)NULL,k,p);
      SubMatrix(B[k],Aplus,0,0,p->nvars*p->nlags,p->nvars);
      FreeMatrix(Aplus);
      ProductInverseMM(B[k],B[k],A0_Xi[k]);
      for (j=p->nvars-1; j >= 0; j--)
	for (tmp=sqrt(p->Zeta[j][p->var_states[j][k]]), i=p->nvars-1; i >= 0; i--)
	  ElementM(A0_Xi[k],i,j)*=tmp;
    }
}


/*
   Additional command line parameters

     -ndraws <integer>
        Number of draws to save (default value = 1000)
     -burnin <integer>
        Number of burn-in draws (default = 0.1 * thin * ndraws)
     -thin <integer>
        Thinning factor.  Total number of draws made is thin*ndraws + burnin (default = 1)
     -no_deg <integer : retry> <integer : cutoff>
        Rejects degenerate draws of regimes.  Will redraw regimes up to retry times.  In
        order not to be degenerate, there must be at least cutoff observations for each
        state. (default = 1000 1)
     -append
        Append simulation and info to existing files if they exist and use last parameter
        draw as initial value. (default overwrites existing files)
     -mh <integer>
        Tuning period for Metropolis-Hasting draws (default value = 30000)
     -flat
        Produce flat output file and header file (default value = off)
     -nofree
        Do not produce free parameters file (default value = off)
     -nd1
        Normalize diagonal of A0 to one (flat output only)
*/
//#define DEBUG_1
void dw_sbvar_simulate_command_line(int nargs, char **args)
{
  TStateModel *model;
  T_VAR_Parameters *p;
  FILE *f_flat, *f_free, *f_info;
  char *filename;
  time_t begin_time, end_time;
  int count, check, period=1000, i, tuning, burnin, ndraws, thin, append, writeflat, writefree, nd1;
  TVARCommandLine *cmd=(TVARCommandLine*)NULL;
  TVector parameters;

  // Get tuning peroid, burn-in period, number of iterations, and thinning factor
  ndraws=dw_ParseInteger_String(nargs,args,"ndraws",1000);
  thin=dw_ParseInteger_String(nargs,args,"thin",1);
  burnin=dw_ParseInteger_String(nargs,args,"burnin",(ndraws*thin)/10);
  tuning=dw_ParseInteger_String(nargs,args,"mh",30000);

  // Append file
  append=(dw_FindArgument_String(nargs,args,"append") >= 0) ? 1 : 0;
  
  // Get output options
  writeflat=(dw_FindArgument_String(nargs,args,"flat") >= 0) ? 1 : 0;
  writefree=(dw_FindArgument_String(nargs,args,"nofree") >= 0) ? 0 : 1;
  nd1=(dw_FindArgument_String(nargs,args,"nd1") >= 0) ? 1 : 0;

  // Setup model and initial parameters
  if (!(cmd=CreateTStateModel_VARCommandLine(nargs,args,(TVARCommandLine*)NULL)))
    dw_Error(MEM_ERR);
  else
    if (!cmd->model)
      dw_UserError("Unable to setup SBVAR from command line parameters");
    else
      if (!cmd->parameters_filename_actual)
	dw_UserError("Unable to initialize parameters from command line parameters");
      else
	{
	  model=cmd->model;
	  p=(T_VAR_Parameters*)(model->theta);

	  // Degeneracy behavior
	  if ((i=dw_FindArgument_String(nargs,args,"no_deg")) == -1)
	    model->reject_degenerate_draws=0;
	  else
	    {
	      if ((i+1 >= nargs) || ((model->reject_degenerate_draws=atoi(args[i+1])) <= 0)) model->reject_degenerate_draws=100;
	      if ((i+2 >= nargs) || ((model->degenerate_draws_cutoff=atoi(args[i+2])) <= 0)) model->degenerate_draws_cutoff=1;
	    }


	  // Allocate parameters workspace
	  parameters=CreateVector(NumberFreeParametersTheta(model)+NumberFreeParametersQ(model));

	  // initialize parameters from previous simulation run if append
	  if (append)
            {
              if (writefree)
                {
                  filename=CreateFilenameFromTag("%ssimulation_%s.out",cmd->out_tag,cmd->out_directory);
                  f_free=fopen(filename,"rt");
                  dw_free(filename);
                  if (f_free)
                    {
                      while (GetPosteriorDraw(f_free,model,F_FREE));
                      fclose(f_free);
                    }
                }
              else
                if (writeflat)
                  {
                    filename=CreateFilenameFromTag("%sdraws_%s.out",cmd->out_tag,cmd->out_directory);
                    f_flat=fopen(filename,"rt");
                    dw_free(filename);
                    if (f_flat)
                      {
                        while (GetPosteriorDraw(f_flat,model,F_FLAT));
                        fclose(f_flat);
                      }
                  }
            }

	  // Save initial parameters
	  ConvertThetaToFreeParameters(model,pElementV(parameters));
	  ConvertQToFreeParameters(model,pElementV(parameters)+NumberFreeParametersTheta(model));

	  // Output tag
	  if (!cmd->out_tag)
	    cmd->out_tag=dw_DuplicateString(cmd->in_tag ? cmd->in_tag : "");

	  // Open header and flat file and print headers
	  if (writeflat)
	    {
	      filename=CreateFilenameFromTag("%sdraws_header_%s.out",cmd->out_tag,cmd->out_directory);
	      if ((f_flat=fopen(filename,"wt")))
		{
		  dw_free(filename);
		  WriteBaseTransitionMatricesFlat_Headers(f_flat,model);
		  Write_VAR_ParametersFlat_Headers(f_flat,model);
		  fprintf(f_flat,"\n");
		  fclose(f_flat);
		}

	      filename=CreateFilenameFromTag("%sdraws_%s.out",cmd->out_tag,cmd->out_directory);
	      f_flat=fopen(filename,append ? "at" : "wt");
	      dw_free(filename);
	    }
	  else
	    f_flat=(FILE*)NULL;

	  // Open free file
	  if (writefree)
	    {
	      filename=CreateFilenameFromTag("%ssimulation_%s.out",cmd->out_tag,cmd->out_directory);
	      f_free=fopen(filename,append ? "at" : "wt");
	      dw_free(filename);
	    }
	  else
	    f_free=(FILE*)NULL;

	  // Begin time
	  begin_time=(int)time((time_t*)NULL);

	  // Open info file and write
	  filename=CreateFilenameFromTag("%ssimulation_info_%s.out",cmd->out_tag,cmd->out_directory);
	  f_info=fopen(filename,append ? "at" : "wt");
	  dw_free(filename);
	  if (f_info) 
	    {
	      fprintf(f_info,"\nSimulation Run Info\n");
	      fprintf(f_info,"Number saved draws: %d\n",ndraws);
	      fprintf(f_info,"Total draws: %d\n",ndraws*thin);
	      fprintf(f_info,"Thinning factor: %d\n",thin);
	      fprintf(f_info,"Burnin period: %d\n",burnin);
	      fprintf(f_info,"Adaptive tuning period: %d\n",tuning);
	      fprintf(f_info,"File format: %s\n",f_free ? (f_flat ? "free and flat formats" : "free format") : (f_flat ? "flat format" : "file not written")); 
	      fprintf(f_info,"Seed: %d\n",dw_ParseInteger_String(nargs,args,"seed",0));
	      fprintf(f_info,"Begin time stamp: %s",ctime(&begin_time));
	    }

	  // Calibration of jumping parameters
	  printf("Calibrating jumping parameters - %d draws\n",tuning);
	  AdaptiveMetropolisScale(model,tuning,1000,1,(FILE*)NULL);      // tuning iterations - 1000 iterations before updating - verbose
	  end_time=(int)time((time_t*)NULL);
	  printf("Elapsed Time: %d seconds\n",(int)(end_time - begin_time));

	  // Reset initial parameters
	  ConvertFreeParametersToTheta(model,pElementV(parameters));
	  ConvertFreeParametersToQ(model,pElementV(parameters)+NumberFreeParametersTheta(model));

	  // Reset degeneracy counts
	  model->n_degenerate_draws=model->total_degenerate_draws=0;

	  // Burn-in period
	  printf("Burn-in period - %d draws\n",burnin);
	  for (check=period, count=1; count <= burnin; count++)
	    {
	      DrawAll(model);

	      if (count == check)
		{
		  check+=period;
		  printf("%d iterations completed out of %d\n",count,burnin);
		  end_time=(int)time((time_t*)NULL);
		  printf("Elapsed Time: %d seconds\n",(int)(end_time - begin_time));
		}
	    }
	  end_time=(int)time((time_t*)NULL);
	  printf("Elapsed Time: %d seconds\n",(int)(end_time - begin_time));
	  ResetMetropolisInformation(p);

#ifdef DEBUG_1
	  TMatrix *A0_Xi, *B;
	  FILE *f_debug;

	  f_debug=fopen("debug.out","wt");
	  //f_debug=stdout;

	  A0_Xi=dw_CreateArray_matrix(model->sv->nstates);
	  B=dw_CreateArray_matrix(model->sv->nstates);
	  for (i=model->sv->nstates-1; i >= 0; i--)
	    {
	      A0_Xi[i]=CreateMatrix(p->nvars,p->nvars);
	      B[i]=CreateMatrix(p->nvars*p->nlags,p->nvars);
	    }
	  fprintf(f_debug,"%lf %lf %lf\n",LogPosterior_StatesIntegratedOut(model),LogLikelihood_StatesIntegratedOut(model),LogPrior(model));
	  /* GetIRMatrices(model->sv->nstates,A0_Xi,B,p); */
	  /* for (i=0; i < model->sv->nstates; i++) */
	  /*   fprintf(f_debug," %lf",LargestRoot(B[i])); */
	  /* fprintf(f_debug,"\n"); */
	  DrawStates(model);
#endif
 
	  // Simulation
	  printf("Simulating - %d draws\n",ndraws);
	  for (check=period, count=1; count <= ndraws; count++)
	    {
#ifdef DEBUG_1
	      if (f_free)
		{
		  //printf("Previous: %lf %lf %lf\n",LogPosterior_StatesIntegratedOut(model),LogLikelihood_StatesIntegratedOut(model),LogPrior(model));

		  //DrawStates(model);
		  //printf("States: %lf %lf %lf\n",LogPosterior_StatesIntegratedOut(model),LogLikelihood_StatesIntegratedOut(model),LogPrior(model));

		  /* DrawTransitionMatrixParameters(model); */
		  /* printf("%lf %lf %lf\n",LogPosterior_StatesIntegratedOut(model),LogLikelihood_StatesIntegratedOut(model),LogPrior(model)); */

		  /* DrawTheta(model); */
		  /* printf("%lf %lf %lf\n",LogPosterior_StatesIntegratedOut(model),LogLikelihood_StatesIntegratedOut(model),LogPrior(model)); */

		  /* DrawZeta_DotProducts(model); */
		  /* ((T_VAR_Parameters*)(model->theta))->valid_parameters=1; */
		  /* ThetaChanged(model); */
		  /* printf("Zeta: %lf %lf %lf\n",LogPosterior_StatesIntegratedOut(model),LogLikelihood_StatesIntegratedOut(model),LogPrior(model)); */

		  /* DrawA0_Metropolis(model); */
		  /* ((T_VAR_Parameters*)(model->theta))->valid_parameters=1; */
		  /* ThetaChanged(model); */
		  /* printf("A0: %lf %lf %lf\n",LogPosterior_StatesIntegratedOut(model),LogLikelihood_StatesIntegratedOut(model),LogPrior(model)); */

		  /* DrawAplus(model); */
		  /* ((T_VAR_Parameters*)(model->theta))->valid_parameters=1; */
		  /* ThetaChanged(model); */
		  /* printf("%lf %lf %lf\n",LogPosterior_StatesIntegratedOut(model),LogLikelihood_StatesIntegratedOut(model),LogPrior(model)); */

		  if (p->Specification & SPEC_SIMS_ZHA)
		    {
		      Draw_psi(model);
		      Update_bplus_from_lambda_psi(p);
		      Update_Aplus_from_bplus_A0(p);
		      p->valid_parameters=1;
		      ThetaChanged(model);
		      fprintf(f_debug,"%lf %lf %lf\n",LogPosterior_StatesIntegratedOut(model),LogLikelihood_StatesIntegratedOut(model),LogPrior(model));

		      /* Draw_lambda(model); */
		      /* Update_bplus_from_lambda_psi(p); */
		      /* Update_Aplus_from_bplus_A0(p); */
		      /* p->valid_parameters=1; */
		      /* ThetaChanged(model); */
		      /* printf("lambda: %lf %lf %lf\n",LogPosterior_StatesIntegratedOut(model),LogLikelihood_StatesIntegratedOut(model),LogPrior(model)); */
		    }
		  else
		    { printf("Not Sim-Zha specification\n"); dw_exit(0); }

		  // Normalize
		  //Normalize_VAR((T_VAR_Parameters*)(model->theta));

		  //getchar();

		  /* fprintf(f_free,"%lf %lf %lf\n",LogPosterior_StatesIntegratedOut(model),LogLikelihood_StatesIntegratedOut(model),LogPrior(model)); */
		  /* GetIRMatrices(model->sv->nstates,A0_Xi,B,p); */
		  /* for (i=0; i < model->sv->nstates; i++) */
		  /*   fprintf(f_free," %lf",LargestRoot(B[i])); */
		  /* fprintf(f_free,"\n"); */
		}
#else
	      for (i=thin; i > 0; i--) DrawAll(model);

	      if (f_flat)
		{
		  WriteBaseTransitionMatricesFlat(f_flat,(char*)NULL,model,"%lf ");
		  if (nd1)
		    Write_VAR_ParametersFlat_A0_Diagonal_One(f_flat,model,"%lf ");
		  else
		    Write_VAR_ParametersFlat(f_flat,model,"%lf ");
		  fprintf(f_flat,"\n");
		}

	      if (f_free)
		{
		  fprintf(f_free,"%le %le ",LogPosterior_StatesIntegratedOut(model),LogLikelihood_StatesIntegratedOut(model));
		  ConvertThetaToFreeParameters(model,pElementV(parameters));
		  ConvertQToFreeParameters(model,pElementV(parameters)+NumberFreeParametersTheta(model));
		  dw_PrintVector(f_free,parameters,"%le ");
		}
#endif

	      if (count == check)
		{
		  check+=period;
		  printf("%d * %d iterations completed out of %d * %d\n",count,thin,ndraws,thin);
		  end_time=(int)time((time_t*)NULL);
		  printf("Elapsed Time: %d seconds\n",(int)(end_time - begin_time));
		}
	    }
	  end_time=(int)time((time_t*)NULL);
	  printf("Total Elapsed Time: %d seconds\n",(int)(end_time - begin_time));

	  if (f_info)
	    {
	      fprintf(f_info,"End time stamp: %s\n",ctime(&end_time));
	      fprintf(f_info,"Number degenerate draws of regimes: %d (%d)\n",model->n_degenerate_draws,model->total_degenerate_draws);
	      fprintf(f_info,"degenerate draws %s\n",model->reject_degenerate_draws ? "rejected" : "kept");
	      /* fprintf(f_info,"Degenerate draws by regime: "); */
	      /* for (i=0; i < model->sv->nstates; i++) fprintf(f_info,"%d ",model->states_count[i]); */
	      /* fprintf(f_info,"\n"); */
	    }

	  // clean up
	  if (f_flat) fclose(f_flat);
	  if (f_free) fclose(f_free);
	  if (f_info) fclose(f_info);
	  FreeVector(parameters);
	}

  // clean up
  Free_VARCommandLine(cmd);
}
