/*
 * 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_counterfactual.h"
#include "dw_MSStateSpace.h"

/*
   Assumes:
     data   : (nobs+1-tau) x (ny+1) matrix or null pointer
     tau    : initial time.
     z      : ? x (nz+1) matrix  state variables z[t].
     shocks : ? x (nepsilon+1) matrix of fundamental shocks (epsilon[t]).
     prob   : ? x (nbasestates+1) matrix of marginal regime probabilities.
     model  : pointer to valid TStateModel structure.

   Results:
     Computes counterfactual data

   Returns:
     The matrix data upon success and null upon failure.  If data is null, then 
     it is created.

   Notes:
     The first column of z, shocks, and prob contain the time and must be 
     consecutive integers.  The matrix z must contain time tau, and shocks and 
     prob must contain times tau+1 through nobs.

     ASSUMES THAT ONLY CONSTANT TERMS OR VARIANCES ARE TIME VARYING.
*/
TMatrix dw_state_space_counterfactual(TMatrix data, int tau, TMatrix states, TMatrix shocks, TMatrix prob, TStateModel *model)
{
  T_MSStateSpace *statespace=(T_MSStateSpace*)(model->theta);
  TVector z1, z2, y, e=(TVector)NULL, tmp;
  int i, s, t, tt, iz=ElementM(states,0,0), is=ElementM(shocks,0,0), ip=ElementM(prob,0,0);

  // check inputs
  if ((iz > tau) || (ElementM(states,RowM(states)-1,0) < tau) || (ColM(states) != statespace->nz+1) 
      || (is > tau+1) || (ElementM(shocks,RowM(shocks)-1,0) < model->nobs) || (ColM(shocks) != statespace->nepsilon+1)
      || (ip > tau+1) || (ElementM(prob,RowM(prob)-1,0) < model->nobs)|| (ColM(prob) != statespace->nbasestates+1))
    return (TMatrix)NULL;

  // create forecast if necessary
  if (!data && !(data=CreateMatrix(statespace->nobs+1-tau,statespace->ny+1))) return (TMatrix)NULL;

  // debug
  if (!Smooth_MSStateSpace(model) || !ComputeSmoothedProbabilities(model))
    {
      printf("Unable to smooth\n");
      return (TMatrix)NULL;
    }

  // forecast
  z1=CreateVector(statespace->nz);
  z2=CreateVector(statespace->nz);
  y=CreateVector(statespace->ny);
  e=CreateVector(statespace->nepsilon);
  for (i=DimV(z1)-1; i >= 0; i--)
    ElementV(z1,i)=ElementM(states,tau-iz,i+1);
  ProductMV(y,statespace->H[0],z1);
  for (s=statespace->nbasestates-1; s >= 0; s--)
    LinearCombinationV(y,1.0,y,ElementM(prob,tau-ip,s+1),statespace->a[s]);
  ElementM(data,0,0)=tau;
  for (i=statespace->ny-1; i >= 0; i--)
    ElementM(data,0,i+1)=ElementV(y,i);
  for (tt=1, t=tau+1; t <= statespace->nobs; tt++, t++)
    {
      ProductMV(z2,statespace->F[0],z1);
      for (i=DimV(e)-1; i >= 0; i--)
	ElementV(e,i)=ElementM(shocks,t-is,i+1);
      for (s=statespace->nbasestates-1; s >= 0; s--)
	{
	  LinearCombinationV(z2,1.0,z2,ElementM(prob,t-ip,s+1),statespace->b[s]);
	  UpdateProductMV(1.0,z2,ElementM(prob,t-ip,s+1),statespace->Phiz[s],e);
	}
      ProductMV(y,statespace->H[0],z2);
      for (s=statespace->nbasestates-1; s >= 0; s--)
	LinearCombinationV(y,1.0,y,ElementM(prob,t-ip,s+1),statespace->a[s]);
      ElementM(data,tt,0)=t;
      for (i=statespace->ny-1; i >= 0; i--)
	ElementM(data,tt,i+1)=ElementV(y,i);
      tmp=z1;
      z1=z2;
      z2=tmp;
    }
  FreeVector(y);
  FreeVector(z2);
  FreeVector(z1);
  FreeVector(e);

  return data;
}
