
#include "VARbase.h"
#include "VARio.h"
#include "dw_sbvar_command_line.h"
#include "sbvar_estimate.h"
#include "dw_switch.h"
#include "dw_switchio.h"
#include "dw_matrix_rand.h"
#include "dw_error.h"
#include "dw_ascii.h"
#include "dw_std.h"
#include "string.h"

TMatrix SimulateData_VAR(TStateModel *model)
{
  T_VAR_Parameters *p=(T_VAR_Parameters*)(model->theta);
  TVector initial, prob, *shocks;
  TMatrix forecast;
  int m, i, j;

  // Allocate memory
  prob=CreateVector(p->nstates);
  initial=EquateVector((TVector)NULL,p->X[p->nobs]);
  forecast=CreateMatrix(p->nobs,p->nvars);
  shocks=dw_CreateArray_vector(p->nobs);
  for (i=p->nobs-1; i >= 0; i--) shocks[i]=CreateVector(p->nvars);

  // Get filtered probability at time T
  for (j=p->nstates-1; j >= 0; j--)
    ElementV(prob,j)=ProbabilityStateConditionalCurrent(j,p->nobs,model);

  // Draw time T regime
  p->S[0]=m=DrawDiscrete(prob);
              
  // Draw regimes from time T+1 through T+h inclusive
  for (j=0; j < p->nobs; j++)
    {
      ColumnVector(prob,model->sv->Q,m);
      p->S[j]=m=DrawDiscrete(prob);
    }

  // Draw shocks
  for (j=p->nobs-1; j >= 0; j--) dw_NormalVector(shocks[j]);

  // Compute forecast
  forecast_base(forecast,p->nobs,initial,shocks,p->S+1,model);

  // Clean up and return
  FreeVector(initial);
  FreeVector(prob);
  dw_FreeArray(shocks);

  return forecast;
}

TMatrix ActualData_VAR(TStateModel *model)
{
  T_VAR_Parameters *p=(T_VAR_Parameters*)(model->theta);
  TMatrix Y;
  int i, j;

  // Allocate memory
  Y=CreateMatrix(p->nobs,p->nvars);

  for (i=p->nobs-1; i >= 0; i--)
    for (j=p->nvars-1; j >= 0; j--)
      ElementM(Y,i,j)=ElementV(p->Y[i+1],j);

  return Y;
}

int main(int nargs, char **args)
{
  TStateModel *model, *model_constant;
  TVARCommandLine *cmd=(TVARCommandLine*)NULL, *cmd_constant=(TVARCommandLine*)NULL;
  TEstimateInfo* info;
  T_VAR_Parameters *p, *theta;
  TMatrix X, Y, A0, Aplus;
  int **translation_table;
  FILE *f_out;
  char *filename, *header;
  int i, j, s;

  // Setup model
  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
      {
	p=(T_VAR_Parameters*)(cmd->model->theta);

	if (p->npre - p->nvars*p->nlags > 1)
	  printf("Only constant is allowed as an exogenous variable\n");
	else
	  {
	    // Set new out tag
	    if (!cmd->out_tag) 
	      if (!cmd->in_tag)
		cmd->out_tag=dw_DuplicateString("constant_notag");
	      else
		{
		  cmd->out_tag=(char*)dw_malloc((strlen(cmd->in_tag)+10)*sizeof(char));
		  strcat(strcpy(cmd->out_tag,"constant_"),cmd->in_tag);
		}

	    filename=CreateFilenameFromTag("%sinit_%s.dat",cmd->out_tag,cmd->in_directory);
	    if (f_out=fopen(filename,"wt"))
	      {
		// Simulate Data
		X=CreateMatrix(p->nobs,p->npre);
		Y=SimulateData_VAR(cmd->model);
		for (j=p->npre-1; j >= 0; j--) ElementM(X,0,j)=ElementV(p->X[p->nobs],j);

		// The below lines will use the actual data
		// Y=ActualData_VAR(cmd->model);		
		// for (j=p->npre-1; j >= 0; j--) ElementM(X,0,j)=ElementV(p->X[1],j);

		for (i=1; i < p->nobs; i++)
		  {
		    for (j=p->nvars*(p->nlags-1)-1; j >= 0; j--) ElementM(X,i,j+p->nvars)=ElementM(X,i-1,j);
		    for (j=p->nvars-1; j >= 0; j--) ElementM(X,i,j)=ElementM(Y,i-1,j);
		    for (j=p->nvars*p->nlags; j < p->npre; j++) ElementM(X,i,j)=ElementM(X,i-1,j);
		  }

		// Create new TStateModel
		theta=CreateTheta_VAR(p->Specification,p->nvars,p->nlags,p->npre - p->nvars*p->nlags,p->nstates,p->nobs,p->coef_states,p->var_states,p->U,p->V,p->W,Y,X);
		if (p->Specification & SPEC_SIMS_ZHA)
		  SetPriors_VAR_SimsZha(theta,p->A0_prior,p->Aplus_prior,p->zeta_a_prior,p->zeta_b_prior,p->lambda_prior);
		else
		  SetPriors_VAR(theta,p->A0_prior,p->Aplus_prior,p->zeta_a_prior,p->zeta_b_prior);
		model=CreateStateModel(p->nobs,DuplicateMarkovStateVariable(cmd->model->sv),CreateRoutines_VAR(),theta);

		// Write specification
		Write_VAR_Specification(f_out,(char*)NULL,model);

		// Write true parameters
		header="Truth: ";
		fprintf(f_out,"//== %sFree parameters ==//\n",header);
		WriteFreeParameters(f_out,(char*)NULL,cmd->model);
		fprintf(f_out,"\n");
		WriteBaseTransitionMatrices(f_out,(char*)NULL,header,cmd->model);
		Write_VAR_Parameters(f_out,(char*)NULL,header,cmd->model);

		// Create constant parameter model and estimate
		dw_InitializeArray_int(translation_table=dw_CreateRectangularArray_int(p->nvars,1),0);
		theta=CreateTheta_VAR(p->Specification,p->nvars,p->nlags,p->npre - p->nvars*p->nlags,1,p->nobs,translation_table,translation_table,p->U,p->V,p->W,Y,X);
		dw_FreeArray(translation_table);
		if (p->Specification & SPEC_SIMS_ZHA)
		  SetPriors_VAR_SimsZha(theta,p->A0_prior,p->Aplus_prior,p->zeta_a_prior,p->zeta_b_prior,p->lambda_prior);
		else
		  SetPriors_VAR(theta,p->A0_prior,p->Aplus_prior,p->zeta_a_prior,p->zeta_b_prior);
		model_constant=CreateStateModel(p->nobs,CreateMarkovStateVariable_ConstantState(0),CreateRoutines_VAR(),theta);
		InitializeParameters_VAR((T_VAR_Parameters*)model_constant->theta);
		ThetaChanged(model_constant);
		cmd_constant=Base_VARCommandLine(nargs,args);
		cmd_constant->model=model_constant;
		if (cmd_constant->out_tag) dw_free(cmd_constant->out_tag);
		cmd_constant->out_tag=dw_DuplicateString("constant_simulated_data");
		info=SetupEstimateFromVARCommandLine(nargs,args,cmd_constant);
		FindMode_VAR_csminwel(info,(FILE*)NULL);

		// Insert constant parameters
		p=(T_VAR_Parameters*)(model->theta);
		// A0
		A0=MakeA0((TMatrix)NULL,0,(T_VAR_Parameters*)(model_constant->theta));
		for (j=p->nvars-1; j >= 0; j--)
		  for (s=p->n_coef_states[j]-1; s >= 0; s--)
		    for (i=p->nvars-1; i >= 0; i--)
		      ElementV(p->A0[j][s],i)=ElementM(A0,i,j);
		FreeMatrix(A0);
		// Aplus
		Aplus=MakeAplus((TMatrix)NULL,0,(T_VAR_Parameters*)(model_constant->theta));
		for (j=p->nvars-1; j >= 0; j--)
		  for (s=p->n_coef_states[j]-1; s >= 0; s--)
		    for (i=p->npre-1; i >= 0; i--)
		      ElementV(p->Aplus[j][s],i)=ElementM(Aplus,i,j);
		FreeMatrix(Aplus);  
		// Zeta
		for (j=p->nvars-1; j >= 0; j--)
		  for (s=p->n_var_states[j]-1; s >= 0; s--)
		    p->Zeta[j][s]=1.0;
		// b0, bplus, lambda, and psi
		Update_b0_bplus_from_A0_Aplus(p);
		if (p->Specification & SPEC_SIMS_ZHA) Update_lambda_psi_from_bplus(p);
		// Flags
		p->valid_parameters=1;
		// Transition matrix
		DefaultTransitionMatrixParameters(model);

		// Write initial values
		header="Initial: ";
		WriteBaseTransitionMatrices(f_out,(char*)NULL,header,model);
		Write_VAR_Parameters(f_out,(char*)NULL,header,model);

		// Clean-up
		FreeMatrix(X);
		FreeMatrix(Y);
		FreeStateModel(model);
		FreeEstimateInfo(info);
	      }
	    dw_free(filename);
	  }
      }

  Free_VARCommandLine(cmd);
  return 0;
}
