/*
 * 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 "sbvar_historical_decomposition.h"
#include "VARbase.h"
#include "VARio.h"
#include "dw_parse_cmd.h"
#include "dw_error.h"
#include "dw_std.h"

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

static void CheckHistoricalDecomposition(int *S, TMatrix HD, TStateModel *model);
static void CheckCumulativeHistoricalDecomposition(TMatrix CHD, TStateModel *model);

/*
  Assumes:
    S     : integer vector of length nobs + 1 + h) with 0 <= S[t] < nstates
    t0    : integer between 1 and nobs
    h     : non-negative integer
    model : pointer to valid TStateModel structure

  Returns:
    Fills var.  The element in column (nepsilon+nu+2)*i is the ith variable, the 
    element in column (nepsilon+nu+2)*i + 1 is the exogeneous component, the
    element in column (nepsilon+nu+2)*i + 2 + j is the cummulative response of 
    the ith variable to the jth shock.  The fundamental shocks epsilon(t) are 
    ordered first.  The kth row corresponds to time t0 + k.  For k > nobs-t0+1, 
    the element in column (nepsilon+nu+2)*i is the forecast of the ith variable.
    The initial condition is at time t0 -- E[z(t0) | Y(T)]

  Notes:
    This works when the only time varying parameters are Phiy and Phyz, which 
    includes the constant parameter case.  Also, nlags_encoded must be bigger 
    zero.
*/
TMatrix sbvar_cumulative_historical_decomposition(int *S, int t0, int h, TStateModel *model)
{
  TVector *cx, *cy;
  T_VAR_Parameters *p=(T_VAR_Parameters*)(model->theta);
  TMatrix CHD, HD, *A0, *Aplus, *Zeta;
  int i, j, t;
  
  CHD=CreateMatrix(p->nobs-t0+h+1,1+p->nvars*(p->nvars+2));

  HD=sbvar_historical_decomposition(S,model);

  cx=dw_CreateArray_vector(p->nvars+1);
  for (i=0; i < p->nvars; i++)
    InitializeVector(cx[i]=CreateVector(p->npre),0.0);
  cx[p->nvars]=EquateVector((TVector)NULL,p->X[t0]);
  cy=dw_CreateArray_vector(p->nvars+1);
  for (i=0; i < p->nvars; i++)
    InitializeVector(cy[i]=CreateVector(p->nvars),0.0);
  cy[p->nvars]=EquateVector((TVector)NULL,p->Y[t0]);
  A0=MakeA0_All((TMatrix*)NULL,p);
  Aplus=MakeAplus_All((TMatrix*)NULL,p);
  Zeta=MakeZeta_All((TMatrix*)NULL,p);

  ElementM(CHD,0,0)=t0;
  for (j=p->nvars-1; j >= 0; j--)
    {
      ElementM(CHD,0,j*(p->nvars+2)+1)=ElementV(p->Y[t0],j);
      ElementM(CHD,0,j*(p->nvars+2)+2)=ElementV(cy[p->nvars],j);
      for (i=p->nvars-1; i >= 0; i--)
	ElementM(CHD,0,j*(p->nvars+2)+3+i)=ElementV(cy[i],j);
    }

  for (t=t0+1; t <= p->nobs; t++)
    {
      for (i=p->nvars; i >= 0; i--)
	{
	  memmove(pElementV(cx[i])+p->nvars,pElementV(cx[i]),p->nvars*(p->nlags-1)*sizeof(PRECISION));
	  memcpy(pElementV(cx[i]),pElementV(cy[i]),p->nvars*sizeof(PRECISION));
	}

      i=p->nvars;
      ProductVM(cy[i],cx[i],Aplus[S[t]]);
      ProductInverseVM(cy[i],cy[i],A0[S[t]]);

      for (i--; i >= 0; i--)
	{
	  ProductVM(cy[i],cx[i],Aplus[S[t]]);
	  ElementV(cy[i],i)+=ElementM(HD,t,i+1)/sqrt(ElementM(Zeta[S[t]],i,i));
	  ProductInverseVM(cy[i],cy[i],A0[S[t]]);
	}

      ElementM(CHD,t-t0,0)=t;
      for (j=p->nvars-1; j >= 0; j--)
	{
	  ElementM(CHD,t-t0,j*(p->nvars+2)+1)=ElementV(p->Y[t],j);
	  ElementM(CHD,t-t0,j*(p->nvars+2)+2)=ElementV(cy[p->nvars],j);
	  for (i=p->nvars-1; i >= 0; i--)
	    ElementM(CHD,t-t0,j*(p->nvars+2)+3+i)=ElementV(cy[i],j);
	}
    }

  for ( ; t <= p->nobs+h; t++)
    {
      for (i=p->nvars; i >= 0; i--)
	{
	  memmove(pElementV(cx[i]),pElementV(cx[i])+p->nvars,p->nvars*(p->nlags-1)*sizeof(PRECISION));
	  memcpy(pElementV(cx[i]),pElementV(cy[i]),p->nvars*sizeof(PRECISION));
	}

      for (i=p->nvars; i >= 0; i--)
	{   
	  ProductVM(cy[i],cx[i],Aplus[S[t]]);
	  ProductInverseVM(cy[i],cy[i],A0[S[t]]);
	}

      for (j=p->nvars-1; j >= 0; j--)
	{
	  ElementM(CHD,t-t0-1,j*(p->nvars+2))=ElementV(p->Y[t],j);
	  ElementM(CHD,t-t0-1,j*(p->nvars+2)+1)=ElementV(cy[p->nvars],j);
	  for (i=p->nvars-1; i >= 0; i--)
	    ElementM(CHD,t-t0-1,j*(p->nvars+2)+2+i)=ElementV(cy[i],j);
	}
    }

  dw_FreeArray(Zeta);
  dw_FreeArray(Aplus);
  dw_FreeArray(A0);
  dw_FreeArray(cy);
  dw_FreeArray(cx);
  FreeMatrix(HD);

  CheckCumulativeHistoricalDecomposition(CHD,model);

  return CHD;
}


/*
   S is a vector of integers of length at least nobs with 0 <= S[t] < nstates and model 
   is a valid pointer to a properly initialized TStateMode structure.

   Returns a matrix of size nvars+2 by nobs-1.  The first column is the date number, the
   next nvars columns are the historical shocks given the path of regimes S, and the
   last column is S.
*/
TMatrix sbvar_historical_decomposition(int *S, TStateModel *model)
{
  TVector epsilon;
  TMatrix HD, *A0, *Aplus, *Zeta;
  T_VAR_Parameters *p=(T_VAR_Parameters*)(model->theta);
  int i, t;

  InitializeMatrix(HD=CreateMatrix(p->nobs+1,p->nvars+2),0.0);

  epsilon=CreateVector(p->nvars);  
  A0=MakeA0_All((TMatrix*)NULL,p);
  Aplus=MakeAplus_All((TMatrix*)NULL,p);
  Zeta=MakeZeta_All((TMatrix*)NULL,p);

  for (t=1; t <= p->nobs; t++)
    {
      ProductVM(epsilon,p->Y[t],A0[S[t]]);
      UpdateProductVM(1.0,epsilon,-1.0,p->X[t],Aplus[S[t]]);
      ElementM(HD,t,0)=t;
      for (i=p->nvars-1; i >= 0; i--) ElementM(HD,t,i+1)=ElementV(epsilon,i)*sqrt(ElementM(Zeta[S[t]],i,i));
      ElementM(HD,t,p->nvars+1)=S[t];
    }

  dw_FreeArray(Zeta);
  dw_FreeArray(Aplus);
  dw_FreeArray(A0);
  FreeVector(epsilon);

  CheckHistoricalDecomposition(S,HD,model);

  return HD;
}

static void CheckHistoricalDecomposition(int *S, TMatrix HD, TStateModel *model)
{
  TVector epsilon, y;
  TMatrix *A0, *Aplus, *Zeta;
  T_VAR_Parameters *p=(T_VAR_Parameters*)(model->theta);
  int i, k, t;

  epsilon=CreateVector(p->nvars);  
  y=CreateVector(p->nvars);
  A0=MakeA0_All((TMatrix*)NULL,p);
  Aplus=MakeAplus_All((TMatrix*)NULL,p);
  Zeta=MakeZeta_All((TMatrix*)NULL,p);
  for (k=p->nstates-1; k >= 0; k--)
    for (i=p->nvars-1; i >= 0; i--)
      ElementM(Zeta[k],i,i)=1.0/sqrt(ElementM(Zeta[k],i,i));

  for (t=1; t <= p->nobs; t++)
    {
      for (i=p->nvars-1; i >= 0; i--) ElementV(epsilon,i)=ElementM(HD,t,i+1);
      ProductVM(y,epsilon,Zeta[S[t]]);
      UpdateProductVM(1.0,y,1.0,p->X[t],Aplus[S[t]]);
      UpdateProductVM(1.0,y,-1.0,p->Y[t],A0[S[t]]);
      if (Norm(y) > 1.0e-10)
	{
	  printf("historical shocks do not return data: t = %d  S[%d] = %d\n y = ",t,t,S[t]);
	  dw_PrintVector(stdout,y,"%lf ");
	}
    }

  printf("Finished check historical decomposition\n");

  dw_FreeArray(Zeta);
  dw_FreeArray(Aplus);
  dw_FreeArray(A0);
  FreeVector(epsilon);
} 

static void CheckCumulativeHistoricalDecomposition(TMatrix CHD, TStateModel *model)
{
  T_VAR_Parameters *p=(T_VAR_Parameters*)(model->theta);
  PRECISION y;
  int i, j, t;

  for (t=0; t < RowM(CHD); t++)
    for (j=p->nvars-1; j >= 0; j--)
      {
	y=ElementM(CHD,t,j*(p->nvars+2)+1);
	for (i=p->nvars; i >= 0; i--)
	  y-=ElementM(CHD,t,j*(p->nvars+2)+2+i);
	if (fabs(y) > 1.0e-10)
	  printf("cumulative historical decomposition does not sum to data: t = %d, variable = %d  - %lf\n",t,j,y);
      }

  printf("Finished check cumulative historical decomposition\n");
} 


int* sbvar_most_likely_regime(int t0, int horizon, TStateModel *model)
{
  int *S;
  T_VAR_Parameters *p=(T_VAR_Parameters*)(model->theta);
  int i, t;
  TVector pr;

  S=dw_CreateArray_int(p->nobs-t0+1+horizon);
  ComputeSmoothedProbabilities(model);
  for (t=t0; t <= p->nobs; t++)
    for (S[t-t0]=0, i=p->nstates-1; i > 0; i--)
      if (ElementV(model->SP[t],i) > ElementV(model->SP[t],S[t-t0])) S[t-t0]=i;
  if (horizon > 0)
    {
      pr=CreateVector(p->nstates);
      ComputeTransitionMatrix(p->nobs,model);
      EquateVector(pr,model->SP[p->nobs]);
      for ( ; t <= p->nobs+horizon; t++)
	{
	  ProductMV(pr,model->sv->Q,pr);
	  for (S[t-t0]=0, i=p->nstates-1; i > 0; i--)
	    if (ElementV(pr,i) > ElementV(pr,S[t-t0])) S[t-t0]=i;
	}
      FreeVector(pr);
    }

  return S;
}


/*
  Additional command line parameters

    -cumulative
       Produces the cumulative historical decomposition.  The default is
       to produce the historical shock decomposition.

    -t0
       The starting point for the cumulative historical decomposition.  The default
       value is 1, which is the start of the data.

    -horizon
       Produces forecasts of the cumulative historical decomposition.  The default 
       value is 0.

         
*/
void dw_sbvar_historical_decomposition_command_line(int nargs, char **args)
{
  char *out_filename;
  TStateModel *model;
  int *S;
  TMatrix HD;
  FILE *f_out;
  TVARCommandLine *cmd;
  int horizon, t0;

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

	  // Output tag
	  OutTag_VARCommandLine(cmd);

	  out_filename=CreateFilenameFromTag("%ssbvar_historical_%s.out",cmd->out_tag,cmd->out_directory);

	  if (dw_FindArgument_String(nargs,args,"cumulative") == -1)
	    {
	      if (f_out=fopen(out_filename,"wt"))
		{
		  printf("Constructing historical decomposition\n");

		  S=sbvar_most_likely_regime(0,0,model);
		  HD=sbvar_historical_decomposition(S,model);

		  dw_PrintMatrix(f_out,HD,"%lf ");
		  fprintf(f_out,"\n");
		  Write_VAR_ParametersFlat(f_out,model,"%le ");

		  fclose(f_out);
		  dw_FreeArray(S);
		  FreeMatrix(HD);

		  dw_free(out_filename);
		  out_filename=CreateFilenameFromTag("%ssbvar_historical_%s_headers.out",cmd->out_tag,cmd->out_directory);
		  if (f_out=fopen(out_filename,"wt"))
		    {
		      Write_VAR_ParametersFlat_Headers(f_out,model);
		      fclose(f_out);
		    }
		}
	      dw_free(out_filename);
	    }
	  else
	    {
	      out_filename=CreateFilenameFromTag("%ssbvar_cumulative_historical_%s.out",cmd->out_tag,cmd->out_directory);

	      if (f_out=fopen(out_filename,"wt"))
		{
		  printf("Constructing cumulative historical decomposition\n");

		  horizon=dw_ParseInteger_String(nargs,args,"horizon",0);
		  t0=dw_ParseInteger_String(nargs,args,"t0",1);

		  S=sbvar_most_likely_regime(0,horizon,model);
		  HD=sbvar_cumulative_historical_decomposition(S,t0,horizon,model);

		  dw_PrintMatrix(f_out,HD,"%lf ");

		  fclose(f_out);
		  dw_FreeArray(S);
		  FreeMatrix(HD);

		  /* dw_free(out_filename); */
		  /* out_filename=CreateFilenameFromTag("%ssbvar_historical_%s_headers.out",cmd->out_tag,cmd->out_directory); */
		  /* if (f_out=fopen(out_filename,"wt")) */
		  /*   { */
		  /*     Write_VAR_ParametersFlat_Headers(f_out,model); */
		  /*     fclose(f_out); */
		  /*   } */
		}
	      dw_free(out_filename);
	    }
	}

  Free_VARCommandLine(cmd);
}

