/*
 * 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_state_space_impulse_response_special.h"
#include "dw_MSStateSpace.h"
#include "dw_metropolis_simulation.h"
#include "dw_matrix_array.h"
#include "dw_histogram.h"
#include "dw_parse_cmd.h"
#include "dw_math.h"
#include "dw_std.h"

#define size_special 1
int ComputeSpecial(TVector special, TVector z, TVector y)
{
  if (!z) return 1;
  // (1.0/u_y)*(y_t - c_y*c_t - i_y*i_t - g_y*g_t)
  ElementV(special,0) = 6.113390270546515*(ElementV(z,0) - 0.163575357656760*ElementV(z,1) 
					   - 0.157481335414194*ElementV(z,2) - 0.18*ElementV(z,16));
  return 0;
}


/*
   Assumes:
     IR    : (h+1) x (size_special*(nepsilon+nu)) matrix
     S     : integer array of length h+1
     h     : non-negative integer
     model : valid pointer to TStateModel/T_MSStateSpace structure

   Returns:
     One upon success and zero otherwise. 

   Results:
     Fills IR with the impulse responses of the observable variables with respect 
     to the fundamental and measurement error shocks for horizons 0 through h 
     inclusive.

   Notes:
     The current value of the parameters implicit in model are used to form the
     impulse responses.  The element in position (k,i+j*size_special) is the 
     response of the ith observable variable to the jth shock at horizon k.  
     Shocks 0 through nepsilon-1 are fundamental and shocks nepsilon through 
     nepsilon+nu-1 are measurement error.  Horizon 0 is the contemporaneous 
     response.
*/
int dw_state_space_impulse_response_special(TMatrix IR, int *S, int h, TStateModel *model)
{
  TMatrix rtrn;
  TVector *z, y, e, special;
  T_MSStateSpace *statespace=(T_MSStateSpace*)(model->theta);
  int i, j, k, err;

  // check that there are fundamental or measurement shocks
  if (statespace->nepsilon + statespace->nu == 0) return 0;

  // initialize structure
  if ((statespace->t0 < 0) && !Initialize_MSStateSpace(model)) return 0;

  // allocate memory and initialize
  special=CreateVector(size_special);
  y=CreateVector(statespace->ny);
  InitializeMatrix(IR,0.0);

  // contemporaneous response to fundamental shocks
  if (statespace->nepsilon > 0)
    {
      z=dw_CreateArray_vector(statespace->nepsilon);
      InitializeVector(e=CreateVector(statespace->nepsilon),0.0);
      for (j=statespace->nepsilon-1; j >= 0; j--)
	{
	  ElementV(e,j)=1.0;
	  z[j]=ProductMV((TVector)NULL,statespace->Phiz[S[0]],e);
	  ProductMV(y,statespace->H[S[0]],z[j]);

	  // begin special computation
	  err=ComputeSpecial(special,z[j],y);

	  if (!err)
	    for (i=size_special-1; i >= 0; i--)
	      ElementM(IR,0,i+j*size_special)=ElementV(special,i);
	  ElementV(e,j)=0.0;
	}
      FreeVector(e);
    }

  // contemporaneous response to measument error shocks
  if (statespace->nu > 0)
    {
      InitializeVector(e=CreateVector(statespace->nu),0.0);
      for (j=statespace->nu-1; j >= 0; j--)
	{
	  ElementV(e,j)=1.0;
	  ProductMV(y,statespace->Phiy[S[0]],e);

	  // begin special computation
          err=ComputeSpecial(special,(TVector)NULL,y);

	  if (!err)
	    for (i=size_special-1; i >= 0; i--)
	      ElementM(IR,0,i+(statespace->nepsilon+j)*size_special)=ElementV(special,i);
	  ElementV(e,j)=0.0;
	}
      FreeVector(e);
    }

  // horizon k response to fundamental shocks
  if (statespace->nepsilon > 0)
    {
      for (k=1; k <= h; k++)
	for (j=statespace->nepsilon-1; j >= 0; j--)
	  {
	    ProductMV(z[j],statespace->F[S[k]],z[j]);
	    ProductMV(y,statespace->H[S[k]],z[j]);

	    // begin  special computation
	    err=ComputeSpecial(special,z[j],y);

	    if (!err)
	      for (i=size_special-1; i >= 0; i--)
		ElementM(IR,k,i+j*size_special)=ElementV(special,i);
	  }
      dw_FreeArray(z);
    }

  // free memory and return success
  FreeVector(y);
  FreeVector(special);
  return 1;
}

/*
   Assumes
    f_out : valid FILE pointer
    s : constant regime for the impluse response.
    percentiles : vector of numbers between 0 and 1 inclusive
    posterior_draws : each row contains log posterior, log likelihood, and posterior draw
    thin : thinning factor
    h : non-negative integer
    model : point to valid TStateModel/T_MSStateSpace structure

   Results:
    Computes and prints to the file f_out the requested percentiles of the 
    impulse responses of the observables variables to the fundamental and 
    measurement error shocks.

   Returns:
    One upon success and zero otherwise.

   Notes:
     For each percentile, the element in position (k,i+j*size_special) is the 
     response of the ith observable variable to the jth shock at horizon k.  
     Shocks 0 through nepsilon-1 are fundamental and shocks nepsilon through 
     nepsilon+nu-1 are measurement error.  Horizon 0 is the contemporaneous 
     response.
*/
int dw_state_space_impulse_response_special_percentile_regime(FILE *f_out, int s, TVector percentiles, TMatrix posterior_draws, 
			                                          int thin, int h, TStateModel *model)
{
  T_MSStateSpace *statespace=(T_MSStateSpace*)(model->theta);
  int rtrn=0, ntheta=NumberFreeParametersTheta(model), nq=NumberFreeParametersQ(model), *S, i, j, k;
  TVector x;
  TMatrix IR;
  TMatrixHistogram *histogram;

  // quick check of passed parameters
  if (!f_out || !percentiles || !posterior_draws || (thin <= 0) || (h < 0) || !model) return 0;

  // allocate memory
  x=CreateVector(ColM(posterior_draws));
  S=(int*)dw_malloc((h+1)*sizeof(int));
  IR=CreateMatrix(h+1,size_special*(statespace->nepsilon+statespace->nu));
  histogram=CreateMatrixHistogram(RowM(IR),ColM(IR),40,HISTOGRAM_VARIABLE);

  // set S to s
  for (i=0; i <= h; i++) S[i]=s;

  for (i=RowM(posterior_draws)-1; i >= 0; i-=thin)
    {
      // Get free parameters and push them into model
      RowVector(x,posterior_draws,i);
      if (ntheta > 0) ConvertFreeParametersToTheta(model,pElementV(x)+2);
      if (nq > 0) ConvertFreeParametersToQ(model,pElementV(x)+2+ntheta);
      
      // Get transition matrix -- this should be modified if transition matrix is time varying
      if ((model->sv->t1 < model->nobs) && !sv_ComputeTransitionMatrix(model->nobs,model->sv,model))
	goto ERROR_EXIT;

      // Compute impulse response
      if (!dw_state_space_impulse_response_special(IR,S,h,model))
	goto ERROR_EXIT;

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

  for (i=0; i < DimV(percentiles); i++)
    {
      MatrixPercentile(IR,ElementV(percentiles,i),histogram);
      dw_PrintMatrix(f_out,IR,"%lg ");
      fprintf(f_out,"\n");
    }

  rtrn=1;

ERROR_EXIT:
  FreeMatrixHistogram(histogram);
  FreeMatrix(IR);
  dw_free(S);
  FreeVector(x);

  return rtrn;
}

/*
   Assumes
    f_out : valid FILE pointer
    percentiles : vector of numbers between 0 and 1 inclusive
    posterior_draws : each row contains log posterior, log likelihood, and posterior draw
    thin : thinning factor
    init_prob : vector or null pointer
    h : non-negative integer
    model : point to valid TStateModel/T_MSStateSpace structure

   Results:
    Computes and prints to the file f_out the requested percentiles of the 
    impulse responses of the observables variables to the fundamental and 
    measurement error shocks.

   Returns:
    One upon success and zero otherwise.

   Notes:
     For each percentile, the element in position (k,i+j*ny) is the response of
     the ith observable variable to the jth shock at horizon k.  Shocks 0 through 
     nepsilon-1 are fundamental and shocks nepsilon through nepsilon+nu-1 are 
     measurement error.  Horizon 0 is the contemporaneous response.
*/
int dw_state_space_impulse_response_special_percentile_ergodic(FILE *f_out, TVector percentiles, TMatrix posterior_draws, 
							       int thin, TVector init_prob, int h, TStateModel *model)
{
  T_MSStateSpace *statespace=(T_MSStateSpace*)(model->theta);
  int rtrn=0, ntheta=NumberFreeParametersTheta(model), nq=NumberFreeParametersQ(model), *S, use_ergodic, i, j, k;
  TVector x, prob;
  TMatrix IR;
  TMatrixHistogram *histogram;

  // quick check of passed parameters
  if (!f_out || !percentiles || !posterior_draws || (thin <= 0) || (h < 0) || !model) return 0;

  // allocate memory
  x=CreateVector(ColM(posterior_draws));
  prob=CreateVector(model->sv->nstates);
  S=(int*)dw_malloc((h+1)*sizeof(int));
  IR=CreateMatrix(h+1,size_special*(statespace->nepsilon+statespace->nu));
  histogram=CreateMatrixHistogram(RowM(IR),ColM(IR),40,HISTOGRAM_VARIABLE);

  // Use ergodic initial probability
  if (!init_prob)
    {
      use_ergodic=1;
      init_prob=CreateVector(model->sv->nstates);
    }
  else
    use_ergodic=0;

  for (i=RowM(posterior_draws)-1; i >= 0; i-=thin)
    {
      // Get free parameters and push them into model
      RowVector(x,posterior_draws,i);
      if (ntheta > 0) ConvertFreeParametersToTheta(model,pElementV(x)+2);
      if (nq > 0) ConvertFreeParametersToQ(model,pElementV(x)+2+ntheta);
      
      // Get transition matrix -- this should be modified if transition matrix is time varying
      if ((model->sv->t1 < model->nobs) && !sv_ComputeTransitionMatrix(model->nobs,model->sv,model))
	goto ERROR_EXIT;

      // Get ergodic distribution if necessary
      if (use_ergodic) Ergodic(init_prob,model->sv->Q);

      // Draw regimes
      S[0]=DrawDiscrete(init_prob);
      for (j=1; j <= h; j++)
	{
	  ColumnVector(prob,model->sv->Q,S[j-1]);
	  S[j]=DrawDiscrete(prob);
	}

      // Compute impulse response
      if (!dw_state_space_impulse_response_special(IR,S,h,model))
	goto ERROR_EXIT;

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

  for (i=0; i < DimV(percentiles); i++)
    {
      MatrixPercentile(IR,ElementV(percentiles,i),histogram);
      dw_PrintMatrix(f_out,IR,"%lg ");
      fprintf(f_out,"\n");
    }

  rtrn=1;

ERROR_EXIT:
  if (use_ergodic) dw_free(init_prob);
  FreeMatrixHistogram(histogram);
  FreeMatrix(IR);
  dw_free(S);
  FreeVector(prob);
  FreeVector(x);

  return rtrn;
}

/*
   Command line parameters

     -t : filename tag
     -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)
     -horizon : The horizon over which to compute the impulse responses. (default = 12)
     -thin : Thinning factor.  Only 1/thin of the draws in posterior draws file are used. (default = 1)
     -ergodic : Print impuluse responses using ergodic distribution as initial condition for shocks.
     -regimes : Print impulse responses as if each regime were permanent.
     -regime s : Print impulse responses to regime s as if this regime were permanent.

*/
void dw_state_space_impulse_response_special_command_line(int nargs, char **args, TStateModel *model)
{
  FILE *f_in, *f_out;
  char *model_id, *in_filename, *out_filename, *fmt, *tag;
  TMatrix posterior_draws;
  TVector percentiles=(TVector)NULL, init_prob=(TVector)NULL;
  int s, i, j, n, nparameters=NumberFreeParametersTheta(model) + NumberFreeParametersQ(model), h, thin, states, observables;

  // MSStateSpace model?
  if (!(model_id=ModelType(model)) || strcmp(model_id,"T_MSStateSpace"))
    {
      printf("dw_state_space_forecast_command_line(): Model type = %s\n",model_id ? model_id : "null");
      return;
    }

  // Get tag from command line
  tag=dw_ParseString(nargs,args,'t',(char*)NULL);

  // Read posterior draws
  printf("Reading posterior draws\n");
  posterior_draws=dw_ReadPosteriorDraws((FILE*)NULL,(char*)NULL,tag,nparameters);
  if (!posterior_draws)
    {
      printf("dw_state_space_impulse_response_command_line(): Unable to read posterior draws\n");
      printf("  tag: %s\n",tag ? tag : "null");
      dw_exit(0);
    }

  // Setup percentiles, horizon, thinning factor
  if ((i=dw_FindArgument_String(nargs,args,"percentiles")) == -1)
    {
      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);
	    printf("dw_state_space_impulse_response_command_line(): Error parsing percentiles\n");
	    dw_exit(0);
	  }
      }
    else
      {
	printf("dw_state_space_impulse_response_command_line(): Error parsing percentiles\n");
	dw_exit(0);
      }
  h=dw_ParseInteger_String(nargs,args,"horizon",8);
  thin=dw_ParseInteger_String(nargs,args,"thin",1);

  // ergodic impulse responses
  if (dw_FindArgument_String(nargs,args,"ergodic") != -1)
    {
      // Open output file for impulse responses of observables
      fmt="ir_special_percentiles_ergodic_%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);

      // Compute percentiles for impulse responses of observables
      printf("Constructing percentiles for impulse responses of special - ergodic\n");
      dw_state_space_impulse_response_special_percentile_ergodic(f_out,percentiles,posterior_draws,thin,init_prob,h,model);
      fclose(f_out);
    }

  if (dw_FindArgument_String(nargs,args,"regimes") != -1)
    for (s=0; s < model->sv->nstates; s++)
      {
	// Open output file for impulse responses of observables
	fmt="ir_special_percentiles_regime_%d_%s.prn";
	sprintf(out_filename=(char*)dw_malloc(strlen(fmt) + strlen(tag) + 10),fmt,s,tag);
	f_out=fopen(out_filename,"wt");
	dw_free(out_filename);

	// Compute percentiles for impulse responses of observables
	printf("Constructing percentiles for impulse responses of special - regime %d\n",s);
	dw_state_space_impulse_response_special_percentile_regime(f_out,s,percentiles,posterior_draws,thin,h,model);
	fclose(f_out);
      }
  else
    if (((s=dw_ParseInteger_String(nargs,args,"regime",-1)) >= 0) && (s < model->sv->nstates))
      {
	// Open output file for impulse responses of observables
	fmt="ir_special_percentiles_regime_%d_%s.prn";
	sprintf(out_filename=(char*)dw_malloc(strlen(fmt) + strlen(tag) + 10),fmt,s,tag);
	f_out=fopen(out_filename,"wt");
	dw_free(out_filename);

	// Compute percentiles for impulse responses of observables
	printf("Constructing percentiles for impulse responses of special - regime %d\n",s);
	dw_state_space_impulse_response_special_percentile_regime(f_out,s,percentiles,posterior_draws,thin,h,model);
	fclose(f_out);
      }

  FreeMatrix(posterior_draws);
  FreeVector(percentiles);
}
