/*
 * 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_historical_decomposition.h"
#include "dw_matrix_array.h"
#include "dw_std.h"

#include <math.h>

/*
  Assumes:
    var   : (nobs - t0 + 1 + h) x (nepsilon + nu + 2)*ny matrix
    t0    : integer between 1 and nobs - 1
    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 dw_state_space_historical_decomposition(TMatrix var, int t0, int h, TStateModel *model)
{
  T_MSStateSpace *statespace=(T_MSStateSpace*)(model->theta);
  TVector *z_decomp, y, z, u, e, pr,
    *P,                          // P[t][s] = P(s(t)=s | Y(T)) 1 <= t <= T
    **I1,                        // I1[t][s] = E[z(t) | Y(T), s(t)=s] 1 <= t <= T
    **I2,                        // I2[t][s] = E[z(t) | Y(T), s(t+1)=s] 1 <= t <= T-1
    **SEepsilon=(TVector**)NULL, // SEepsilon[t][s] = E[epsilon(t) | Y(T), s(t)=s] 2 <= t <= T
    **SEu=(TVector**)NULL;       // SEu[t][s] = E[u(t) | Y(T), s(t)=s] 2 <= t <= T
  int t, s, i, j, nc=statespace->nepsilon+statespace->nu+2;

  if (!Smooth_MSStateSpace(model) || !ComputeSmoothedProbabilities(model)) return (TMatrix)NULL;

  y=CreateVector(statespace->ny);
  z=CreateVector(statespace->nz);

  if (statespace->nlags_encoded > 1)
    {
      P=IntegrateStates((TVector*)NULL,model->SP,statespace->zeta_modulus,statespace->nbasestates,2);
      I1=IntegrateStatesV((TVector**)NULL,P,statespace->ISEz,statespace->nbasestates,
			  statespace->zeta_modulus/statespace->nbasestates,2);
      dw_FreeArray(P);
    }
  else
    if (statespace->nlags_encoded > 0)
      I1=statespace->ISEz;
    else
      {
	printf("Not yet implemented\n");
	dw_exit(0);
      }

  P=IntegrateStates((TVector*)NULL,model->SP,statespace->nbasestates,statespace->zeta_modulus,2);

  if (statespace->nlags_encoded > 0)
    {
      I2=dw_CreateRectangularArray_vector(statespace->nobs+1,statespace->nbasestates);
      for (t=statespace->nobs-1; t > 0; t--)
  	IntegrateStatesSingleV(I2[t],model->SP[t+1],statespace->SEz[t],statespace->nbasestates,statespace->zeta_modulus,2);
    }
  else
    I2=statespace->SEz;

  // Compute smoothed measurment shocks
  SEu=ComputeMeasurmentErrors(model,I1,statespace->nbasestates);

  // Compute smoothed fundamental shocks
  SEepsilon=ComputeFundamentalShocks(model,I1,I2,statespace->nbasestates);

  if (!var) var=CreateMatrix(statespace->nobs-t0+1+h,nc*statespace->ny);
  InitializeMatrix(var,0.0);

  z_decomp=dw_CreateArray_vector(statespace->nepsilon+1);
  for (j=statespace->nepsilon; j >= 0; j--)
    InitializeVector(z_decomp[j]=CreateVector(statespace->nz),0.0);
  for (s=statespace->nbasestates-1; s >= 0; s--)
    LinearCombinationV(z_decomp[0],1.0,z_decomp[0],ElementV(P[t0],s),I1[t0][s]);

  e=(statespace->nepsilon > 0) ? CreateVector(statespace->nepsilon) : (TVector)NULL;
  u=(statespace->nu > 0) ? CreateVector(statespace->nu) : (TVector)NULL;
  pr=EquateVector((TVector)NULL,P[statespace->nobs]);

  // data
  for (t=t0; t <= statespace->nobs; )
    {
      // record data
      for (i=statespace->ny-1; i >= 0; i--)
	ElementM(var,t-t0,nc*i)=ElementV(statespace->y[t],i);

      // Compute cummulative effects of fundamental shocks on observables and record - (requires H[s]=H and a[s]=a)
      for (j=statespace->nepsilon; j >= 0; j--)
	{
	  ProductMV(y,statespace->H[0],z_decomp[j]);
	  for (i=statespace->ny-1; i >= 0; i--)
	    ElementM(var,t-t0,nc*i+j+1)=ElementV(y,i);
	}
      for (i=statespace->ny-1; i >= 0; i--)
	ElementM(var,t-t0,nc*i+1)+=ElementV(statespace->a[0],i);

      // Compute effects of measurment shocks on observables and record
      for (j=statespace->nu-1; j >= 0; j--)
	for (s=statespace->nbasestates-1; s >= 0; s--)
	  {
	    InitializeVector(u,0.0);
	    ElementV(u,j)=ElementV(P[t],s)*ElementV(SEu[t][s],j);
	    ProductMV(y,statespace->Phiy[s],e);
	    for (i=statespace->ny-1; i >= 0; i--)
	      ElementM(var,t-t0,nc*i+statespace->nepsilon+j+2)=ElementV(y,i);
	  }

      // Compute next period cummulative effects
      t++;
      // Uses F[s]=F 
      for (j=statespace->nepsilon; j >= 0; j--)
	ProductMV(z_decomp[j],statespace->F[0],z_decomp[j]);
      if (t <= statespace->nobs)
	for (s=statespace->nbasestates-1; s >= 0; s--)
	  UpdateV(1.0,z_decomp[0],ElementV(P[t],s),statespace->b[s]);
      else
	{
	  ProductMV(pr,model->sv->baseQ,pr);
	  for (s=statespace->nbasestates-1; s >= 0; s--)
	    UpdateV(1.0,z_decomp[0],ElementV(pr,s),statespace->b[s]);
	}

      // Add shocks
      if (t <= statespace->nobs)
	{
	  for (j=statespace->nepsilon-1; j >= 0; j--)
	    for (s=statespace->nbasestates-1; s >= 0; s--)
	      {
		InitializeVector(e,0.0);
		ElementV(e,j)=ElementV(P[t],s)*ElementV(SEepsilon[t][s],j);
		ProductMV(z,statespace->Phiz[s],e);
		AddVV(z_decomp[j+1],z_decomp[j+1],z);
	      }
	}
    }

  // forecasts
  for ( ; t <= statespace->nobs+h; )
    {
      // Compute cummulative effects of fundamental shocks on observables and record - (requires H[s]=H and a[s]=a)
      for (j=statespace->nepsilon; j >= 0; j--)
	{
	  ProductMV(y,statespace->H[0],z_decomp[j]);
	  for (i=statespace->ny-1; i >= 0; i--)
	    {
	      ElementM(var,t-t0,nc*i+j+1)=ElementV(y,i);
	      ElementM(var,t-t0,nc*i)+=ElementV(y,i);
	    }
	}
      for (i=statespace->ny-1; i >= 0; i--)
	{
	  ElementM(var,t-t0,nc*i+1)+=ElementV(statespace->a[0],i);
	  ElementM(var,t-t0,nc*i)+=ElementV(statespace->a[0],i);
	}

      // Compute next period cummulative effects
      t++;
      // Uses F[s]=F
      for (j=statespace->nepsilon; j >= 0; j--)
	ProductMV(z_decomp[j],statespace->F[0],z_decomp[j]);
      ProductMV(pr,model->sv->baseQ,pr);
      for (s=statespace->nbasestates-1; s >= 0; s--)
	UpdateV(1.0,z_decomp[0],ElementV(pr,s),statespace->b[s]);
    }

  FreeVector(u);
  FreeVector(y);
  FreeVector(e);
  FreeVector(z);
  FreeVector(pr);
  dw_FreeArray(z_decomp);
  if (I1 != statespace->ISEz) dw_FreeArray(I1);
  if (I2 != statespace->SEz) dw_FreeArray(I2);
  dw_FreeArray(SEepsilon);
  dw_FreeArray(SEu);
  dw_FreeArray(P);
 
  return var;
}

/*
  Assumes:
    var   : (nobs - t0 + 1 + h) x (nepsilon + 2)*nz matrix
    t0    : integer between 1 and nobs - 1
    model : pointer to valid TStateModel structure

  Returns:
    Fills var.  The element in column (nepsilon+2)*i is E[z(t) | Y(t)], the 
    element in column (nepsilon+2)*i + 1 is the exogeneous component, the
    element in column (nepsilon+nu+2)*i + 2 + j is the cummulative response of 
    the ith state variable to the jth shock.  The kth row corresponds to time 
    t0 + k.  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 dw_state_space_historical_decomposition_states(TMatrix var, int t0, int h, TStateModel *model)
{
  T_MSStateSpace *statespace=(T_MSStateSpace*)(model->theta);
  TVector *z_decomp, z, e, pr,
    *P,                          // P[t][s] = P(s(t)=s | Y(T)) 1 <= t <= T
    **I1,                        // I1[t][s] = E[z(t) | Y(T), s(t)=s] 1 <= t <= T
    **I2,                        // I2[t][s] = E[z(t) | Y(T), s(t+1)=s] 1 <= t <= T-1
    **SEepsilon=(TVector**)NULL; // SEepsilon[t][s] = E[epsilon(t) | Y(T), s(t)=s] 2 <= t <= T
  int t, s, i, j, nc=statespace->nepsilon+2;

  if (!Smooth_MSStateSpace(model) || !ComputeSmoothedProbabilities(model)) return (TMatrix)NULL;

  z=CreateVector(statespace->nz);

  if (statespace->nlags_encoded > 1)
    {
      P=IntegrateStates((TVector*)NULL,model->SP,statespace->zeta_modulus,statespace->nbasestates,2);
      I1=IntegrateStatesV((TVector**)NULL,P,statespace->ISEz,statespace->nbasestates,
			  statespace->zeta_modulus/statespace->nbasestates,2);
      dw_FreeArray(P);
    }
  else
    if (statespace->nlags_encoded > 0)
      I1=statespace->ISEz;
    else
      {
	printf("Not yet implemented\n");
	dw_exit(0);
      }

  P=IntegrateStates((TVector*)NULL,model->SP,statespace->nbasestates,statespace->zeta_modulus,2);

  if (statespace->nlags_encoded > 0)
    {
      I2=dw_CreateRectangularArray_vector(statespace->nobs+1,statespace->nbasestates);
      ComputeTransitionMatrix(model->nobs,model);
      for (t=statespace->nobs-1; t > 0; t--)
	IntegrateStatesSingleV(I2[t],model->SP[t+1],statespace->SEz[t],statespace->nbasestates,statespace->zeta_modulus,2);
    }
  else
    I2=statespace->SEz;

  // Compute smoothed fundamental shocks
  SEepsilon=ComputeFundamentalShocks(model,I1,I2,statespace->nbasestates);

  if (!var) var=CreateMatrix(statespace->nobs-t0+1+h,nc*statespace->nz);
  InitializeMatrix(var,0.0);

  z_decomp=dw_CreateArray_vector(statespace->nepsilon+1);
  for (j=statespace->nepsilon; j >= 0; j--)
    InitializeVector(z_decomp[j]=CreateVector(statespace->nz),0.0);
  for (s=statespace->nbasestates-1; s >= 0; s--)
    LinearCombinationV(z_decomp[0],1.0,z_decomp[0],ElementV(P[t0],s),I1[t0][s]);

  e=(statespace->nepsilon > 0) ? CreateVector(statespace->nepsilon) : (TVector)NULL;
  pr=EquateVector((TVector)NULL,P[statespace->nobs]);

  for (t=t0; t <= statespace->nobs; )
    {
      // compute and record E[z[t] | Y[t])
      for (i=statespace->nz-1; i >= 0; i--)
	for (s=statespace->nbasestates-1; s >= 0; s--)
	  ElementM(var,t-t0,nc*i)+=ElementV(P[t],s)*ElementV(I1[t][s],i);

      // record cummulative effects of fundamental shocks on states
      for (j=statespace->nepsilon; j >= 0; j--)
	for (i=statespace->nz-1; i >= 0; i--)
	  ElementM(var,t-t0,nc*i+j+1)=ElementV(z_decomp[j],i);

      // Compute next period cummulative effects
      t++;
      // Uses F[s]=F and b[s]=b 
      for (j=statespace->nepsilon; j >= 0; j--)
	ProductMV(z_decomp[j],statespace->F[0],z_decomp[j]);
      if (t <= statespace->nobs)
	for (s=statespace->nbasestates-1; s >= 0; s--)
	  UpdateV(1.0,z_decomp[0],ElementV(P[t],s),statespace->b[s]);
      else
	{
	  ProductMV(pr,model->sv->baseQ,pr);
	  for (s=statespace->nbasestates-1; s >= 0; s--)
	    UpdateV(1.0,z_decomp[0],ElementV(pr,s),statespace->b[s]);
	}

      if (t <= statespace->nobs)
	{
	  for (j=statespace->nepsilon-1; j >= 0; j--)
	    for (s=statespace->nbasestates-1; s >= 0; s--)
	      {
		InitializeVector(e,0.0);
		ElementV(e,j)=ElementV(P[t],s)*ElementV(SEepsilon[t][s],j);
		ProductMV(z,statespace->Phiz[s],e);
		AddVV(z_decomp[j+1],z_decomp[j+1],z);
	      }
	}
    }

  // forecasts
  for ( ; t <= statespace->nobs+h; )
    {
      // record cummulative effects of fundamental shocks on states
      for (j=statespace->nepsilon; j >= 0; j--)
	for (i=statespace->nz-1; i >= 0; i--)
	  {
	    ElementM(var,t-t0,nc*i+j+1)=ElementV(z_decomp[j],i);
	    ElementM(var,t-t0,nc*i)+=ElementV(z_decomp[j],i);
	  }

      // Compute next period cummulative effects
      t++;
      // Uses F[s]=F and b[s]=b 
      for (j=statespace->nepsilon; j >= 0; j--)
	ProductMV(z_decomp[j],statespace->F[0],z_decomp[j]);
      ProductMV(pr,model->sv->baseQ,pr);
      for (s=statespace->nbasestates-1; s >= 0; s--)
	UpdateV(1.0,z_decomp[0],ElementV(pr,s),statespace->b[s]);
    }

  FreeVector(e);
  FreeVector(z);
  FreeVector(pr);
  dw_FreeArray(z_decomp);
  if (I1 != statespace->ISEz) dw_FreeArray(I1);
  if (I2 != statespace->SEz) dw_FreeArray(I2);
  dw_FreeArray(SEepsilon);
  dw_FreeArray(P);
 
  return var;
}

TVector* SmoothedShocks(TStateModel *model)
{
  T_MSStateSpace *statespace=(T_MSStateSpace*)(model->theta);
  TVector *P,                          // P[t][s] = P(s(t)=s | Y(T)) 1 <= t <= T
          **I1,                        // I1[t][s] = E[z(t) | Y(T), s(t)=s] 1 <= t <= T
          **I2,                        // I2[t][s] = E[z(t) | Y(T), s(t+1)=s] 1 <= t <= T-1
          **SEepsilon=(TVector**)NULL, // SEepsilon[t][s] = E[epsilon(t) | Y(T), s(t)=s] 1 <= t <= T
          *rtrn=(TVector*)NULL;        // rtrn[t] = E[epsilon(t) | Y(T)] 1 <= t <= T
  int t, k;

  if (!Smooth_MSStateSpace(model) || !ComputeSmoothedProbabilities(model)) return (TVector*)NULL;

  if (statespace->nlags_encoded > 1)
    {
      P=IntegrateStates((TVector*)NULL,model->SP,statespace->zeta_modulus,statespace->nbasestates,2);
      I1=IntegrateStatesV((TVector**)NULL,P,statespace->ISEz,statespace->nbasestates,
			  statespace->zeta_modulus/statespace->nbasestates,2);
      dw_FreeArray(P);
    }
  else
    if (statespace->nlags_encoded > 0)
      I1=statespace->ISEz;
    else
      {
	printf("Not yet implemented\n");
	dw_exit(0);
      }

  P=IntegrateStates((TVector*)NULL,model->SP,statespace->nbasestates,statespace->zeta_modulus,2);

  if (statespace->nlags_encoded > 0)
    {
      I2=dw_CreateRectangularArray_vector(statespace->nobs+1,statespace->nbasestates);
      ComputeTransitionMatrix(model->nobs,model);
      for (t=statespace->nobs-1; t > 0; t--)
	IntegrateStatesSingleV(I2[t],model->SP[t+1],statespace->SEz[t],statespace->nbasestates,statespace->zeta_modulus,2);
    }
  else
    I2=statespace->SEz;

  // Compute smoothed fundamental shocks
  SEepsilon=ComputeFundamentalShocks(model,I1,I2,statespace->nbasestates);

  rtrn=dw_CreateArray_vector(statespace->nobs+1);
  for (t=statespace->nobs; t > 1; t--)
    for (rtrn[t]=ProductVS((TVector)NULL,SEepsilon[t][0],ElementV(P[t],0)), k=statespace->nbasestates-1; k > 0; k--)
      UpdateV(1.0,rtrn[t],ElementV(P[t],k),SEepsilon[t][k]);

  dw_FreeArray(SEepsilon);
  dw_FreeArray(P);
  if (I1 != statespace->ISEz) dw_FreeArray(I1);
  if (I2 != statespace->SEz) dw_FreeArray(I2);

  return rtrn;
}

TVector* SmoothedStates(TStateModel *model)
{
  T_MSStateSpace *statespace=(T_MSStateSpace*)(model->theta);
  TVector *P,                          // P[t][s] = P(s(t)=s | Y(T)) 1 <= t <= T
          **I1,                        // I1[t][0] = E[z(t) | Y(T)] 1 <= t <= T
          *rtrn=(TVector*)NULL;        // rtrn[t] = E[z(t) | Y(T)] 1 <= t <= T
  int t;

  if (!Smooth_MSStateSpace(model) || !ComputeSmoothedProbabilities(model)) return (TVector*)NULL;

  if (statespace->nlags_encoded > 0)
    {
      P=IntegrateStates((TVector*)NULL,model->SP,statespace->zeta_modulus,statespace->nbasestates,2);
      I1=IntegrateStatesV((TVector**)NULL,P,statespace->ISEz,1,statespace->zeta_modulus,2);
      dw_FreeArray(P);
    }
  else
    I1=statespace->ISEz;

  rtrn=dw_CreateArray_vector(statespace->nobs+1);
  for (t=statespace->nobs; t >= 1; t--)
    rtrn[t]=EquateVector((TVector)NULL,I1[t][0]);

  if (I1 != statespace->ISEz) dw_FreeArray(I1);

  return rtrn;
}

/*******************************************************************************/
/******************************** Test Routines ********************************/
/*******************************************************************************/
void CheckHistoricalDecomposition(TMatrix var, int verbose, TStateModel *model)
{
  T_MSStateSpace *statespace=(T_MSStateSpace*)(model->theta);
  PRECISION sum, max=0.0, dif, pmax=0.0, pdif;
  int i, j, t, tmax, imax, ptmax, pimax, nc=ColM(var)/statespace->ny;
  printf("Checking sums in historical decomposition\n");
  for (t=RowM(var)-1; t >= 0; t--)
    {
      if (verbose) printf("t=%d ",t);
      for (i=statespace->ny-1; i >= 0; i--)
	{
	  for (sum=0.0, j=nc-1; j > 0; j--) sum+=ElementM(var,t,i*nc+j);
	  dif=fabs(sum-ElementM(var,t,i*nc));
	  pdif=(ElementM(var,t,i*nc) != 0) ? 100.0*fabs(dif/ElementM(var,t,i*nc)) : -1.0;
	  if (dif > max) { max=dif; tmax=t; imax=i; }
	  if (pdif > pmax) { pmax=pdif; ptmax=t; pimax=i; } 
	  if (verbose) printf("%lf ",dif);
	}
      if (verbose) printf("\n");
    }
  printf("Max error: %le - t=%d  variable=%d\n",max,tmax,imax);
  printf("Max %% error: %lf - t=%d  variable=%d\n",pmax,ptmax,pimax);
}
