/*
 * 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_switch.h"
#include "dw_switchio.h"
#include "VARbase.h"
#include "VARio.h"
#include "dw_sbvar_command_line.h"
#include "dw_matrix_rand.h"
#include "dw_parse_cmd.h"
#include "dw_ascii.h"
#include "dw_error.h"
#include "dw_histogram.h"
#include "dw_math.h"
#include "dw_std.h"

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

TMatrix ForecastHistogramToPercentile(TMatrix F, int data, int horizon, int T, PRECISION percentile, TMatrixHistogram *histogram, TStateModel *model, int flag)
{
  T_VAR_Parameters *p=(T_VAR_Parameters*)(model->theta);
  TMatrix G;
  int i, j, d=flag ? 1 : 0;
  if (!F)
    {
      if (!(F=CreateMatrix(data+horizon,p->nvars+d))) return (TMatrix)NULL;
    }
  else
    if ((RowM(F) != data+horizon) || (ColM(F) != p->nvars+d))
      {
	dw_Error(SIZE_ERR);
	return (TMatrix)NULL;
      }
  MatrixPercentile(G=CreateMatrix(horizon,p->nvars),percentile,histogram);
  for (i=0; i < data; i++)
    {
      if (d) ElementM(F,i,0)=T-data+i+1;
      for (j=0; j < p->nvars; j++) ElementM(F,i,j+d)=ElementV(p->Y[T-data+i+1],j);
    }
  for (i=0; i < horizon; i++)
    {
      if (d) ElementM(F,i+data,0)=T+i+1;
      for (j=0; j < p->nvars; j++) ElementM(F,i+data,j+d)=ElementM(G,i,j);
    }
  FreeMatrix(G);
  return F;
}

/*
   Assumes
    f_out : valid FILE pointer
    percentiles : vector of numbers between 0 and 1 inclusive
    draws : number of draws of shocks and regimes to make for each posterior draw
    posterior_file : FILE pointer to file containing posterior draws.  If null, current parameters are used.
    T : last observation to treat as data.  Usually equals model->nobs.
    h : non-negative integer
    model : point to valid TStateModel structure

   Results:
    Computes and prints to the file f_out the requested percentiles for forecasts 
    of the observables.

   Returns:
    One upon success and zero otherwise.

*/
TMatrixHistogram* forecast_percentile(int draws, FILE *posterior_file, int thin, int T, int h, TStateModel *model, int flag)
{
  T_VAR_Parameters *p;
  int done=0, *S, i=0, j, k, m, n=1000;
  TVector init_prob, prob, *shocks, initial;
  TMatrix forecast;
  TMatrixHistogram *histogram;

 // quick check of passed parameters
  if ((draws <= 0) || (T < 0) || (h <= 0) || !model) return (TMatrixHistogram*)NULL;

  p=(T_VAR_Parameters*)(model->theta);

  if (T > p->nobs) return (TMatrixHistogram*)NULL;

  // allocate memory
  S=(int*)dw_malloc(h*sizeof(int));
  forecast=CreateMatrix(h,p->nvars);
  histogram=CreateMatrixHistogram(h,p->nvars,100,HISTOGRAM_VARIABLE);
  initial=CreateVector(p->npre);
  shocks=dw_CreateArray_vector(h);
  for (i=h-1; i >= 0; i--) shocks[i]=CreateVector(p->nvars);
  init_prob=CreateVector(p->nstates);
  prob=CreateVector(p->nstates);

  // Initial value
  EquateVector(initial,p->X[T]);

  i=0;
  while (!done)
    {
      // Read parameters and push them into model
      if (!posterior_file)
	done=1;
      else
	{
	  dw_NextLine(posterior_file,thin-1);
	  if (!GetPosteriorDraw(posterior_file,model,flag))
	    {
	      done=2;
	      printf("total posterior draws processed - %d\n",i);
	    }
	  else
	    if (i++ == n)
	      {
		printf("%d posterior draws processed\n",i);
		n+=1000;
	      }
	}

      if (done != 2)
	{
	    // Get filtered probability at time T
	    for (j=p->nstates-1; j >= 0; j--)
	      ElementV(init_prob,j)=ProbabilityStateConditionalCurrent(j,T,model);

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

		// Draw shocks
		for (j=h-1; j >= 0; j--) dw_NormalVector(shocks[j]); // InitializeVector(shocks[i],0.0);

		// Compute forecast
		if (!forecast_base(forecast,h,initial,shocks,S,model))
		  goto ERROR_EXIT;

		// Accumulate impulse response
		AddMatrixObservation(forecast,histogram);
	      }
	  }
    }

ERROR_EXIT:
  FreeMatrix(forecast);
  dw_free(S);
  FreeVector(initial);
  FreeVector(prob);
  FreeVector(init_prob);
  dw_FreeArray(shocks);

  return histogram;
}

/*
   Assumes
    f_out : valid FILE pointer
    percentiles : vector of numbers between 0 and 1 inclusive
    draws : number of draws of shocks to make for each posterior draw
    posterior_file : FILE pointer to file containing posterior draws.  If null, current parameters are used.
    s : base state
    T : last observation to treat as data.  Usually equals model->nobs.
    h : non-negative integer
    model : point to valid TStateModel/T_MSStateSpace structure

   Results:
    Computes and prints to the file f_out the requested percentiles for forecasts 
    of the observables.

   Returns:
    One upon success and zero otherwise.

   Notes:
    The regime at time T is drawn from the filtered probabilities at time t, and
    is set to s there after. 

*/
TMatrixHistogram* forecast_percentile_regime(int draws, FILE *posterior_file, int thin, int s, int T, int h, TStateModel *model, int flag)
{
  T_VAR_Parameters *p;
  int done=0, *S, i=0, j, k, n=1000;
  TVector init_prob, prob, *shocks, initial;
  TMatrix forecast;
  TMatrixHistogram *histogram;

 // quick check of passed parameters
  if ((draws <= 0) || (T < 0) || (h <= 0) || !model) return 0;

  p=(T_VAR_Parameters*)(model->theta);

  if (T > p->nobs) return 0;

  // allocate memory
  S=(int*)dw_malloc(h*sizeof(int));
  for (i=0; i < h; i++) S[i]=s;
  forecast=CreateMatrix(h,p->nvars);
  histogram=CreateMatrixHistogram(h,p->nvars,100,HISTOGRAM_VARIABLE);
  initial=CreateVector(p->npre);
  shocks=dw_CreateArray_vector(h);
  for (i=h-1; i >= 0; i--) shocks[i]=CreateVector(p->nvars);
  init_prob=CreateVector(p->nstates);
  prob=CreateVector(p->nstates);

  // Initial value
  EquateVector(initial,p->X[T]);

  i=0;
  while (!done)
    {
      // Read parameters and push them into model
      if (!posterior_file)
	done=1;
      else
	{
	  dw_NextLine(posterior_file,thin-1);
	  if (!GetPosteriorDraw(posterior_file,model,flag))
	    {
	      done=2;
	      printf("total posterior draws processed - %d\n",i);
	    }
	  else
	    if (i++ == n)
	      {
		printf("%d posterior draws processed\n",i);
		n+=1000;
	      }
	}

      if (done != 2)
	{
	  for (k=draws; k > 0; k--)
	    {
	      // Draw shocks
	      for (j=h-1; j >= 0; j--) dw_NormalVector(shocks[j]); // InitializeVector(shocks[i],0.0);

	      // Compute forecast
	      if (!forecast_base(forecast,h,initial,shocks,S,model))
		goto ERROR_EXIT;

	      // Accumulate impulse response
	      AddMatrixObservation(forecast,histogram);
	    }
	}
    }

ERROR_EXIT:
  FreeMatrix(forecast);
  dw_free(S);
  FreeVector(initial);
  FreeVector(prob);
  FreeVector(init_prob);
  dw_FreeArray(shocks);

  return histogram;
}


/*
  Additional command line parameters

   -horizon <integer>
      If this argument exists, then the forecast horizon is given by the passed 
      integer.  The default value is 12.

   -data <integer>
      If this argument exists, then the number of data points included in the
      output is given by the passed integer.  The default value is 0.

   -error_bands 
      Output error bands.  (default = off - only median is computed)

   -percentiles n p_1 p_2 ... p_n
      Percentiles to compute. The first parameter after percentiles must be the 
      number of percentiles and the following values are the actual percentiles. 
      default = 3  0.16  0.50  0.84   if error_bands flag is set
              = 1  0.50               otherwise

   -parameter_uncertainty 
      Apply parameter uncertainty when computing error bands or median.

   -shocks_per_parameter <integer> 
      Number of shocks and regime paths to draw for each parameter draw.  The 
      default value is 1 if parameter_uncertainty is set and 10,000 otherwise.
     
   -thin 
      Thinning factor.  Only 1/thin of the draws in posterior draws file are 
      used. The default value is 1.

   -regimes 
      Produces forecasts as if each regime were permanent. (default = off)

   -regime <integer>
      Produces forecasts as if regime were permanent.  Regime numbers are zero
      based.  (default = off)

   -mean
      Produces mean forecast.  (default = off)

   -flat
      Assumes that posterior draws is a flat file format.  The default is that 
      the posterior draws are in a free file format.  

   -nodate
      If not defined, then the first column of the output is a date number.  A 
      date number of one is the first date in the data series, a date number of
      nobs is the last date in the data series, and date numbers larger than this
      correspond to forecasts.  (default = off)

*/
void dw_sbvar_forecast_command_line(int nargs, char **args)
{
  char *post=(char*)NULL, *out_filename, *fmt="%%sforecasts_percentiles_regime_%d_%%s.out", *buffer;
  TStateModel *model;
  T_VAR_Parameters *p;
  TVector percentiles=(TVector)NULL;
  int dates, type, data, s, horizon, thin, draws, i, j, n, T;
  FILE *f_out, *posterior_file;
  TMatrixHistogram *histogram;
  TMatrix forecast;

  TVARCommandLine *cmd=(TVARCommandLine*)NULL;

  // 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);

	  // horizon
	  horizon=dw_ParseInteger_String(nargs,args,"horizon",12);

	  // Output tag
	  OutTag_VARCommandLine(cmd);

	  // Type
	  type=(dw_FindArgument_String(nargs,args,"flat") >= 0) ? F_FLAT : F_FREE;

	  // Dates
	  dates=(dw_FindArgument_String(nargs,args,"nodate") >= 0) ? 0 : 1;

	  // Data output
	  data=dw_ParseInteger_String(nargs,args,"data",0);

	  // Last data point
	  T=p->nobs;


	  //============================= Compute forecasts ============================= 

	  /* // Mean forecast */
	  /* if (dw_FindArgument_String(nargs,args,"mean") != -1) */
	  /*   { */
	  /*     fmt="forecasts_mean_%s.prn"; */
	  /*     sprintf(out_filename=(char*)dw_malloc(strlen(fmt) + strlen(tag) - 1),fmt,tag); */
	  /*     f_out=fopen(out_filename,"wt"); */
	  /*     dw_free(out_filename); */
	  /*     printf("Constructing mean forecast\n"); */
	  /*     if (F=dw_state_space_mean_unconditional_forecast((TVector*)NULL,h,statespace->nobs,model)) */
	  /* 	for (i=0; i < h; i++) */
	  /* 	  dw_PrintVector(f_out,F[i],"%le "); */
	  /*     fclose(f_out); */
	  /*   } */

	  // Parameter uncertainty
	  if (dw_FindArgument_String(nargs,args,"parameter_uncertainty") != -1)
	    {
	      // Open posterior draws file
	      post=CreateFilenameFromTag((type == F_FREE) ? "%ssimulation_%s.out" : "%sdraws_%s.out",cmd->in_tag ? cmd->in_tag : "notag",cmd->in_directory);
	      if (!(posterior_file=fopen(post,"rt")))
		{
		  dw_free(post);
		  printf("dw_sbvar_forecast command line(): Unable to open posterior draws file: %s\n",post);
		  return;
		}
	      else
		dw_free(post);

	      // Get thinning factor from command line
	      if ((thin=dw_ParseInteger_String(nargs,args,"thin",1)) <= 0) thin=1;

	      // Get shocks_per_parameter from command line
	      draws=dw_ParseInteger_String(nargs,args,"shocks_per_parameter",1);
	    }
	  else
	    {
	      // Using posterior estimate
	      posterior_file=(FILE*)NULL;

	      // thinning factor not used
	      thin=1;

	      // Get shocks_per_parameter from command line
	      draws=dw_ParseInteger_String(nargs,args,"shocks_per_parameter",10000);
	    }

	  // Setup percentiles
	  if ((i=dw_FindArgument_String(nargs,args,"percentiles")) == -1)
	    if (dw_FindArgument_String(nargs,args,"error_bands") == -1)
	      {
		percentiles=CreateVector(1);
		ElementV(percentiles,0)=0.5;
	      }
	    else
	      {
		percentiles=CreateVector(3);
		ElementV(percentiles,0)=0.16; ElementV(percentiles,1)=0.5; ElementV(percentiles,2)=0.84;
	      }
	  else
	    if ((i+1 < nargs) && dw_IsInteger(args[i+1]) && ((n=atoi(args[i+1])) > 0) && (i+1+n < nargs))
	      {
		percentiles=CreateVector(n);
		for (j=0; j < n; j++)
		  if (!dw_IsFloat(args[i+2+j])|| ((ElementV(percentiles,j)=atof(args[i+2+j])) <= 0.0) || (ElementV(percentiles,j) >= 1.0)) 
		    break;
		if (j < n)
		  {
		    FreeVector(percentiles);
		    dw_UserError("dw_sbvar_forecast command line():  Error parsing percentiles\n");
		    return;
		  }
	      }
	    else
	      {
		dw_UserError("dw_sbvar_forecast command line():  Error parsing percentiles\n");
		return;
	      }

	  if (dw_FindArgument_String(nargs,args,"regimes") != -1)
	    for (s=0; s < p->nstates; s++)
	      {
		if (posterior_file) rewind(posterior_file);
		// s must not contain more than 10 digits 
		sprintf(buffer=(char*)dw_malloc(strlen(fmt) + 10),fmt,s);
		out_filename=CreateFilenameFromTag(buffer,cmd->out_tag,cmd->out_directory);
		dw_free(buffer);
		if (f_out=fopen(out_filename,"wt"))
		  {		
		    printf("Constructing percentiles for forecasts - regime %d\n",s);
		    if (histogram=forecast_percentile_regime(draws,posterior_file,thin,s,T,horizon,model,type))
		      {
			for (i=0; i < DimV(percentiles); i++)
			  if (forecast=ForecastHistogramToPercentile((TMatrix)NULL,data,horizon,T,ElementV(percentiles,i),histogram,model,dates))
			    {
			      dw_PrintMatrix(f_out,forecast,"%lg ");
			      fprintf(f_out,"\n");
			      FreeMatrix(forecast);
			    }
			FreeMatrixHistogram(histogram);
		      }
		    fclose(f_out);
		  }
		dw_free(out_filename);
	      }
	  else
	    if (dw_FindArgument_String(nargs,args,"regime") != -1)
              if (((s=dw_ParseInteger_String(nargs,args,"regime",-1)) >= 0) && (s < p->nstates))
		{
		  // s must not contain more than 10 digits 
		  sprintf(buffer=(char*)dw_malloc(strlen(fmt) + 10),fmt,s);
		  out_filename=CreateFilenameFromTag(buffer,cmd->out_tag,cmd->out_directory);
		  dw_free(buffer);
		  if (f_out=fopen(out_filename,"wt"))
		    {		
		      printf("Constructing percentiles for forecasts - regime %d\n",s);
		      if (histogram=forecast_percentile_regime(draws,posterior_file,thin,s,T,horizon,model,type))
			{
			  for (i=0; i < DimV(percentiles); i++)
			    if (forecast=ForecastHistogramToPercentile((TMatrix)NULL,data,horizon,T,ElementV(percentiles,i),histogram,model,dates))
			      {
				dw_PrintMatrix(f_out,forecast,"%lg ");
				fprintf(f_out,"\n");
				FreeMatrix(forecast);
			      }
			  FreeMatrixHistogram(histogram);
			}
		      fclose(f_out);
		    }
		  dw_free(out_filename);
		}
	      else
		{
		  printf("dw_sbvar_forecast command line(): regime number must be between 0 and %d inclusive\n",p->nstates-1);
		  return;
		}
	    else
	      {
		out_filename=CreateFilenameFromTag("%sforecasts_percentiles_%s.out",cmd->out_tag,cmd->out_directory);
		if (f_out=fopen(out_filename,"wt"))
		  {
		    printf("Constructing percentiles for forecasts - %d draws of shocks/regimes per posterior value\n",draws);
		    if (histogram=forecast_percentile(draws,posterior_file,thin,T,horizon,model,type))
		      {
			for (i=0; i < DimV(percentiles); i++)
			  if (forecast=ForecastHistogramToPercentile((TMatrix)NULL,data,horizon,T,ElementV(percentiles,i),histogram,model,dates))
			    {
			      dw_PrintMatrix(f_out,forecast,"%lg ");
			      fprintf(f_out,"\n");
			      FreeMatrix(forecast);
			    }
			FreeMatrixHistogram(histogram);
		      }
		    fclose(f_out);
		  }
		dw_free(out_filename);
	      }

	  if (posterior_file) fclose(posterior_file);
	  FreeVector(percentiles);
	  //=============================================================================
	}

  Free_VARCommandLine(cmd);
}
