/*
 * 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.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"

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

/*
   Assumes:
     IR    : (h+1) x (nz*nepsilon) 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 state variables with respect to 
     the fundamental 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*nz) is the response of
     the ith state variable to the jth fundamental shock at horizon k.  Horizon 0 
     is the contemporaneous response.
*/
/* int dw_state_space_impulse_response_states(TMatrix IR, int *S, int h, TStateModel *model) */
/* { */
/*   TVector *z, e; */
/*   T_MSStateSpace *statespace=(T_MSStateSpace*)(model->theta); */
/*   int i, j, k; */

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

/*   // check that there are fundamental shocks */
/*   if (statespace->nepsilon == 0) return 0; */

/*   // allocate memory */
/*   z=dw_CreateArray_vector(statespace->nepsilon); */

/*   // contemporaneous response to fundamental shocks */
/*   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); */
/*       for (i=statespace->nz-1; i >= 0; i--) */
/* 	ElementM(IR,0,i+j*statespace->nz)=ElementV(z[j],i); */
/*       ElementV(e,j)=0.0; */
/*     } */
/*   FreeVector(e); */

/*   // horizon k response to fundamental shocks */
/*   for (k=1; k <= h; k++) */
/*     for (j=statespace->nepsilon-1; j >= 0; j--) */
/*       { */
/* 	ProductMV(z[j],statespace->F[S[k]],z[j]); */
/* 	for (i=statespace->nz-1; i >= 0; i--) */
/* 	  ElementM(IR,k,i+j*statespace->nz)=ElementV(z[j],i); */
/*       } */

/*   // free memory and return success */
/*   dw_FreeArray(z); */
/*   return 1; */
/* } */

/*
   Assumes:
     IR    : (h+1) x (ny*(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*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_observables(TMatrix IR, int *S, int h, TStateModel *model) */
/* { */
/*   TVector *z, y, e; */
/*   T_MSStateSpace *statespace=(T_MSStateSpace*)(model->theta); */
/*   int i, j, k; */

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

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

/*   // allocate memory and initialize */
/*   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]); */
/* 	  for (i=statespace->ny-1; i >= 0; i--) */
/* 	    ElementM(IR,0,i+j*statespace->ny)=ElementV(y,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); */
/* 	  for (i=statespace->ny-1; i >= 0; i--) */
/* 	    ElementM(IR,0,i+(statespace->nepsilon+j)*statespace->ny)=ElementV(y,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]); */
/* 	    for (i=statespace->ny-1; i >= 0; i--) */
/* 	      ElementM(IR,k,i+j*statespace->ny)=ElementV(y,i); */
/* 	  } */
/*       dw_FreeArray(z); */
/*     } */

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

/*
   Assumes
    f_out : valid FILE pointer
    s : constant base regime for the impluse response.
    percentiles : vector of numbers between 0 and 1.
    posterior_draws : each row contains log posterior, log likelihood, and posterior draw or null
    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 state variables to the fundamental shocks given
    one is always in regime s.

   Returns:
    One upon success and zero otherwise.

   Notes:
     For each percentile, the element in position (k,i+j*nz) is the response of
     the ith state 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_states_percentile_regime(FILE *f_out, int s, TVector percentiles, TMatrix posterior_draws, 
						             int thin, int h, TStateModel *model)
{
  T_MSStateSpace *statespace;
  int rtrn=0, ntheta, nq, *S, i;
  TVector x;
  TMatrix IR;
  TMatrixHistogram *histogram;

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

  statespace=(T_MSStateSpace*)(model->theta);
  ntheta=NumberFreeParametersTheta(model);
  nq=NumberFreeParametersQ(model);

  // s must be in the proper range
  if ((s < 0) || (statespace->nbasestates <= s)) return 0;

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

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

  if (posterior_draws)
    {
      x=CreateVector(ColM(posterior_draws));
      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(IR,S,model,IR_STATES))
	    goto ERROR_EXIT;

	  // Accumulate impulse response
	  AddMatrixObservation(IR,histogram);
	}
      FreeVector(x);
    }
  else
    {
      // 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(IR,S,model,IR_STATES))
	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);

  return rtrn;
}

/*
   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 or null
    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*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_observables_percentile_regime(FILE *f_out, int s, TVector percentiles, TMatrix posterior_draws, 
			                                          int thin, int h, TStateModel *model)
{
  T_MSStateSpace *statespace;
  int rtrn=0, ntheta, nq, *S, i;
  TVector x;
  TMatrix IR;
  TMatrixHistogram *histogram;

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

  statespace=(T_MSStateSpace*)(model->theta);
  ntheta=NumberFreeParametersTheta(model);
  nq=NumberFreeParametersQ(model);

  // s must be in the proper range
  if ((s < 0) || (statespace->nbasestates <= s)) return 0;

  // allocate memory
  S=(int*)dw_malloc((h+1)*sizeof(int));
  IR=CreateMatrix(h+1,statespace->ny*(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;

  if (posterior_draws)
    {
      x=CreateVector(ColM(posterior_draws));
      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(IR,S,model,IR_OBSERVABLES))
	    goto ERROR_EXIT;

	  // Accumulate impulse response
	  AddMatrixObservation(IR,histogram);
	}
      FreeVector(x);
    }
  else
    {
      // 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(IR,S,model,IR_OBSERVABLES))
	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);


  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 state variables to the fundamental shocks.

   Returns:
    One upon success and zero otherwise.

   Notes:
     For each percentile, the element in position (k,i+j*nz) is the response of
     the ith state 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_states_percentile_ergodic(FILE *f_out, TVector percentiles, TMatrix posterior_draws, 
						              int thin, TVector init_prob, int h, TStateModel *model)
{
  T_MSStateSpace *statespace;
  int rtrn=0, ntheta, nq, *S, use_ergodic, i, j, m;
  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;

  statespace=(T_MSStateSpace*)(model->theta);
  ntheta=NumberFreeParametersTheta(model);
  nq=NumberFreeParametersQ(model);

  // allocate memory
  x=CreateVector(ColM(posterior_draws));
  prob=CreateVector(model->sv->nstates);
  S=(int*)dw_malloc((h+1)*sizeof(int));
  IR=CreateMatrix(h+1,statespace->nz*statespace->nepsilon);
  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
      m=DrawDiscrete(init_prob);
      S[0]=model->sv->lag_index[m][0];
      for (j=1; j <= h; j++)
	{
	  ColumnVector(prob,model->sv->Q,m);
	  m=DrawDiscrete(prob);
	  S[j]=model->sv->lag_index[m][0];
	}

      // Compute impulse response
      if (!dw_state_space_impulse_response(IR,S,model,IR_STATES))
	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;
}

/*
   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_observables_percentile_ergodic(FILE *f_out, TVector percentiles, TMatrix posterior_draws, 
							           int thin, TVector init_prob, int h, TStateModel *model)
{
  T_MSStateSpace *statespace;
  int rtrn=0, ntheta, nq, *S, use_ergodic, i, j, m;
  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;

  statespace=(T_MSStateSpace*)(model->theta);
  ntheta=NumberFreeParametersTheta(model);
  nq=NumberFreeParametersQ(model);

  // allocate memory
  x=CreateVector(ColM(posterior_draws));
  prob=CreateVector(model->sv->nstates);
  S=(int*)dw_malloc((h+1)*sizeof(int));
  IR=CreateMatrix(h+1,statespace->ny*(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
      m=DrawDiscrete(init_prob);
      S[0]=model->sv->lag_index[m][0];
      for (j=1; j <= h; j++)
	{
	  ColumnVector(prob,model->sv->Q,m);
	  m=DrawDiscrete(prob);
	  S[j]=model->sv->lag_index[m][0];
	}

      // Compute impulse response
      if (!dw_state_space_impulse_response(IR,S,model,IR_OBSERVABLES))
	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;
}

