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

#ifndef __MARKOV_SWITCHING_STATE_SPACE_MODEL__
#define __MARKOV_SWITCHING_STATE_SPACE_MODEL__

#ifdef __cplusplus
#include "MarkovSwitchingModel.hh"
extern "C" {
#endif

#include "dw_switch.h"
#include "dw_switchio.h"
#include "dw_matrix.h"

/* Impulse response flags */
#define IR_STATES       0x0001
#define IR_OBSERVABLES  0x0002

#ifndef __cplusplus
typedef struct MakovSwitchingModel MarkovSwitchingModel;
typedef struct GensysForm_tag GensysForm;
#endif

typedef struct T_MSStateSpace_tag
{
  int ny, 
    nz, 
    nu, 
    nepsilon, 
    nobs, 
    nstates,             // nstates = nbasestates^(nlags_encoded+1)
    nbasestates,         // h is equal to nbasestates in the comments below
    nlags_encoded,       // rho is equal to nlag_encoded in the comments below
    zeta_modulus;        // zeta_modulus = nbasestates^nlags_encoded

  // Data
  TVector *y;            // y[t] must be defined for 1 <= t <= nobs

  // Reduced form parameters
  TVector *a;
  TMatrix *H;
  TVector *b;
  TMatrix *F;
  TMatrix *R;
  TMatrix *V;
  TMatrix *G;
  TMatrix *Phiy;
  TMatrix *Phiz;

  int NonZeroG;

  // Filtering
  int t0;          // Kalman forward recursion updated through t0

  TVector **Ez;    // Ez[t][k] ~ E[z(t) | Y(t), xi(t)=k, theta, q]

  TMatrix **Ezz;   // Ezz[t][k] ~ E[(z(t) - Ez(t,xi(t)))(z(t) - Ez(t,xi(t)))' | Y(t), xi(t)=k, theta, q] 

  TVector **IEz;   // IEz[t][j] ~ E[z(t) | Y(t), zeta(t)=j, theta, q]

  TMatrix **IEzz;  // IEzz[t][j] ~ E[(z(t) - Ez(t,zeta(t)))(z(t) - Ez(t,zeta(t)))' | Y(t), zeta(t)=j, theta, q] 

  TVector **Ez1;   // Ez1[t][k] ~ E[z(t) | Y(t-1), xi(t)=k, theta, q]

  TMatrix **Ezz1;  // Ezz1[t][k] ~ E[(z(t) - Ez1(t,xi(t)))(z(t) - Ez1(t,xi(t)))' | Y(t-1), xi(t)=k, theta, q] 

  TVector **Ey1;   // Ey1[t][k] ~ E[y(t) | Y(t-1), xi(t)=k, theta, q]

  TMatrix **Eyy1;  // Eyy1[t][k] ~ E[(y(t) - Ey1(t,xi(t)))(y(t) - Ey1(t,xi(t)))' | Y(t-1), xi(t)=k, theta, q] 

  TVector **ey1;   // ey1[t][k] ~ y(t) - E[y(t) | Y(t-1), xi(t)=k, theta, q]

  TVector *L;      // L[t](k) ~ P[y(t) | Y(t-1), xi(t)=k, theta, q]

  // Smoothing
  int smoothed;          // non-zero if Kalman smoothing completed

  TVector **SEz;         // SEz[t][k] ~ E[z(t) | Y(T), xi(t+1)=k, theta, q] for 1 <= t < nobs
  TVector **ISEz;        // ISEz[t][j] ~ E[z(t) | Y(T), zeta(t)=j, theta, q] for 1 <= t <= nobs

  // Tao's data structure
  TStateModel *TZstatemodel;
  
  // Jake's Markov Model
  void *MSModel;
  void *gensys_struct; 

} T_MSStateSpace; 

T_MSStateSpace* Create_MSStateSpace(int ny, int nz, int nu, int nepsilon, int nbasestates, int nlags_encoded, int nobs);
T_MSStateSpace* Create_MSStateSpaceFromStateModel(TStateModel *TZstatemodel, int nlags_encoded); 
void Free_MSStateSpace(T_MSStateSpace *statespace);
ThetaRoutines* CreateThetaRoutines_MSStateSpace(void);

TStateModel* CreateStateModel_MSStateSpace(TStateModel *basemodel, int nlags_encoded);

char* ModelType_MSStateSpace(void);
double LogPrior_MSStateSpace(TStateModel *model);
void StatesChanged_MSStateSpace(TStateModel *model);
void ThetaChanged_MSStateSpace(TStateModel *model);
void TransitionMatrixParametersChanged_MSStateSpace(TStateModel *model);
void Invalidate_MSStateSpace(T_MSStateSpace *statespace);
int Initialize_MSStateSpace(TStateModel *model);
void Invalidate_MSStateSpace(T_MSStateSpace *statespace);
int NumberFreeParametersTheta_MSStateSpace(TStateModel *model);
void ConvertFreeParametersToTheta_MSStateSpace(TStateModel *model, PRECISION *buffer);
void ConvertThetaToFreeParameters_MSStateSpace(TStateModel *model, PRECISION *buffer);

PRECISION LogConditionalLikelihood_MSStateSpace(int s, int t, TStateModel *model);

TVector pEyt_MSStateSpace(TVector y, int T, TVector *data, int s, int t, struct TStateModel_tag *model);

// Kalman Filtering
int Filter(int T, TStateModel *model);
int ConditionalFilter(int t0, int t1, PRECISION ***y, TVector *Ez_i, TMatrix *Ezz_i, TVector Pxi_i, TMatrix Q, 
		      TVector *Pxi, TVector *Pxi1, TVector **IEz, TMatrix **IEzz, TVector **Ez1, TMatrix **Ezz1, 
		      TVector **Ez, TMatrix **Ezz, T_MSStateSpace *statespace);

int Smooth_MSStateSpace(TStateModel *model);
int SmoothProbabilities_MSStateSpace(int t0, int t1, TVector *SPxi, TVector *Pxi, TVector *Pxi1, TMatrix Q);
void SmoothMean_MSStateSpace(int t0, int t1, TVector **SEz, TVector **ISEz, TVector **Ez1, TMatrix **Ezz1, 
			     TVector **IEz, TMatrix **IEzz, TVector *SPxi, T_MSStateSpace *statespace);
TVector** ComputeFundamentalShocks(TStateModel *model, TVector **Ez, TVector **Ez1, int a);
TVector** ComputeMeasurmentErrors(TStateModel *model, TVector **Ez, int a);

int WriteFilteredProbabilities(FILE *f, char *filename, TStateModel *model);

TMatrix Forecast_MSStateSpace(TMatrix forecast, int t0, int data_periods, int forecast_periods, TVector *shocks, TStateModel *model);
TVector* CummulativeHistoricalDecomposition(int t0, TStateModel *model);
TVector* CummulativeHistoricalDecompositionStates(int t0, TStateModel *model);
TVector* Check_HistoricalDecompositon_Shocks(TVector *dev, TStateModel *model);

int dw_state_space_variance_decomposition(TMatrix VD, int *S,  TStateModel *model, int flag);
int dw_state_space_impulse_response(TMatrix IR, int *S, TStateModel *model, int flag);

// Checks
void Check_FilteredInfo(TStateModel *model, PRECISION tol);
void Check_SmoothedInfo(TStateModel *model, PRECISION tol);

// Routines to be written
int Filter_ConditionalRegimes(int T, int *S, TStateModel *model);

// Jakes Routines
void draw_free_parameters_from_prior_markov_model(void* msm, double **buffer);
double get_log_prior_from_markov_model(void* msm);
void states_changed_in_markov_model(void* msm);
void theta_changed_in_markov_model(void* msm);
void transition_matrix_parameters_changed_im_markov_model(void *msm);
void invalidate_markov_model(void *msm);
int number_free_parameters_theta_markov_model(void *msm);
void convert_free_parameter_to_theta_markov_model(void *msm, double *buffer);
void convert_theta_to_free_parameters_markov_model(void *msm, double *buffer);
  
#ifdef __cplusplus
}
#endif  
  
#endif

/********************************************************************************
State space form:

    y(t) = a(s(t)) + H(s(t))*z(t) + Phiy(s(t))*u(t)               (measurement equation)
    z(t) = b(s(t)) + F(s(t))*z(t-1) + Phiz(s(t))*epsilon(t)       (state equation)


    uhat(t) = Phiy(s(t))*u(t)
    epsilonhat(t) = Phiz(s(t))*epsilon(t)
    E[u(t)u(t)'] = I  ==>  E[uhat(t)uhat(t)'] = R(s(t)) = Phiy(s(t))*Phiy'(s(t))
    E[epsilon(t)epsilon(t)'] = I  ==>  E[epsilonhat(t)epsilonhat(t)'] = V(s(t)) = Phiz(s(t))*Phiz'(s(t))
    E[epsilonhat(t)uhat(t)'] = Phiz(s(t))*E[epsilon(t)uhat(t)']*Phiy(s(t))' = G(s(t))
    u(t) and epsilon(t) are independent of s(t)

    y(t)       : ny vector
    u(t)       : ny vector
    z(t)       : nz vector
    epsilon(t) : nz vector
    a(i)       : ny vector
    H(i)       : ny x nz matrix
    b(i)       : nz vector
    F(i)       : nz x nz matrix
    R(i)       : ny x ny matrix
    V(i)       : nz x nz matrix
    G(i)       : nz x ny matrix
    Phiy(i)    : ny x nu matrix
    Phiz(i)    : nz x nepsilon matrix

    0 <= i < h


Fundemental relation:
 
        x1        mu1   sigma11  sigma12
    if      ~  N      ,                 
        x2        mu2   sigma21  sigma22

    then x1 | x2  ~  N(m1 + sigma12 inverse(sigma22) (x2 - mu2), sigma11 - sigma12 inverse(sigma22) sigma21)


Recursive filtering relations:

    zeta(t) = (s(t), ... ,s(t-rho+1))  
    xi(t)   = (s(t), ... ,s(t-rho+1),s(t-rho)) = (s(t),zeta(t-1)) = (zeta(t),s(t-rho))

  The index rho >= 0 is fixed.  If rho == 0, then zeta(t) is empty.  Both xi and zeta 
  are encoded as integers.  Under this encodeding, 

          xi(t) = zeta(t)*h + s(t-rho) = s(t)*(h^rho) + zeta(t-1)

    Pxi(t,xi(t))           ~= P[xi(t) | Y(t), theta, q]

    Ps(t,zeta(t),s(t-rho)) ~= P[s(t-rho) | Y(t), zeta(t), theta, q]

    Ez(t,xi(t))            ~= E[z(t) | Y(t), xi(t), theta, q]

    Ezz(t,xi(t))           ~= E[(z(t) - Ez(t,xi(t)))(z(t) - Ez(t,xi(t)))' | Y(t), xi(t), theta, q] 

    IEz(t,zeta(t))         ~= E[z(t) | Y(t), zeta(t), theta, q]

    IEzz(t,zeta(t))        ~= E[(z(t) - Ez(t,zeta(t)))(z(t) - Ez(t,zeta(t)))' | Y(t), zeta(t), theta, q] 

    Ez1(t,xi(t))           ~= E[z(t) | Y(t-1), xi(t), theta, q]

    Ezz1(t,xi(t))          ~= E[(z(t) - Ezl(t,xi(t)))(z(t) - Ezl(t,xi(t)))' | Y(t-1), xi(t), theta, q] 

    Eyl(t,xi(t))           ~= E[y(t) | Y(t-1), xi(t), theta, q]
   
    Eyyl(t,xi(t)))         ~= E[(y(t) - Ey(t,xi(t)))(y(t) - Ey(t,xi(t)))' | Y(t-1), xi(t), theta, q] 

  Filtering approximations:

    Ez(t,xi(t)) ~ given

    Ezz(t,xi(t)) ~ given

    Pxi(t,xi(t)) = ProbabilityStateConditionalCurrent(xi(t),t,model)

    Ps(t,(zeta(t),s(t-rho))) = Pxi(t,(zeta(t),s(t-rho)))/Sum{Pxi(t,(zeta(t),i)); 0 <= i < h}

    IEz(t,zeta(t)) = Sum{Ps(t,zeta(t),i)*Ez(t,(zeta(t),i)); 0 <= i < h}

    IEzz(t,zeta(t)) = Sum{Ps(t,zeta(t),i)*Ezz(t,(zeta(t),i)); 0 <= i < h}

    Ezl(t+1,(s(t+1),zeta(t))) = b(s(t+1)) + F(s(t+1))*IEz(t,zeta(t))

    Ezzl(t+1,(s(t+1),zeta(t))) = V(s(t+1)) + F(s(t+1))*IEzz(t,zeta(t))*F(s(t+1))'

    Eyl(t+1,(s(t+1),zeta(t))) = a(s(t+1)) + H(s(t+1))*Ez1(t+1,(s(t+1),zeta(t)))
   
    Eyyl(t+1,(s(t+1),zeta(t))) = R(s(t+1)) + H(s(t+1))*Ezz1(t+1,(s(t+1),zeta(t)))*H'(s(t+1)) 
                                     + H(s(t+1))*G(s(t+1)) + G'(s(t+1))*H'(s(t+1))

    Ez(t+1,(s(t+1),zeta(t))) = Ez1(t+1),(s(t+1),zeta(t))) + (Ezz1(t+1,(s(t+1),zeta(t)))*H'(s(t+1)) + G(s(t+1)))
                                   *Inv(Eyy1(t+1,(s(t+1),zeta(t))))*(y(t+1) - Ey1(t+1,(s(t+1),zeta(t))))

    Ezz(t+1,(s(t+1),zeta(t))) = Ezz1(t+1,(s(t+1),zeta(t))) - (Ezz1(t+1,(s(t+1),zeta(t)))*H'(s(t+1)) + G(s(t+1)))
                                     *Inv(Eyy1(t+1,(s(t+1),zeta(t))))*(Ezz1(t+1,(s(t+1),zeta(t)))*H'(s(t+1)) + G(s(t+1)))'

    If the values for Ez(t,xi(t)), Ezz(t,xi(t)), and Pxi(t,xi(t)) were exact, then the values for  
    Ps(t,zeta(t),s(t-rho)), IEz(t,zeta(t)), and IEzz(t,zeta(t)) could be shown to be exact using 
    properties of conditional and marginal probablities.  The values of Ezl(t+1,xi(t+1)), 
    Ezzl(t+1,xi(t+1)), Eyl(t+1,xi(t+1)), and  Eyyl(t+1,xi(t+1)) could be shown to be exact using the
    state and measurement equations.  However, the values of Ez(t+1,xi(t+1)) and Ezz(t+1,xi(t+1)) 
    would be exact only if the joint distribution of z(t+1) and y(t+1), conditional on Y(t) and 
    xi(t+1), were normal.  This would follow from the fundemental relation.  Furthermore, the 
    distribution of z(t+1), conditional on Y(t+1) and xi(t+1) would also be normal.  Unfortunately, 
    even if the distribution of z(t), conditional on Y(t) and xi(t) were normal, the distribution of 
    z(t) conditional on Y(t) and zeta(t) would not be normal, and hence the joint distribution of 
    z(t+1) and y(t+1), conditional on Y(t) and xi(t+1), would not be normal.  Thus the above 
    filtering recursion is only approximate.

    Under the approximation that y(t) conditional on Y(t-1) and xi(t) is normal with mean 
    Ey1(t,xi(t)) and variance Eyy1(t,xi(t), we can recursively compute the function 
    LogConditionalLikelihood(xi(t),t,model). This recursion will call 
    ProbabilityStateConditionalCurrent(xi(tau),tau,model), and hence ForwardRecursion(tau,model), 
    for tau <= t-1.  Since ForwardRecursion(tau,model) only makes calls to 
    LogConditionalLikelihood(xi(tau),tau,model) with tau <= t-1, this recursion will be well 
    defined.


Likelihood Function:

  The likelihood function P(y(t) | Y(t-1), xi(t), theta, q) is approximately normal with mean Ey1(t,xi(t)
  and variance Eyy1(t,xi(t)).  Thus the log likelihood is approximately

    -0.5*(log(2*pi)*ny + LogAbsDeterminate(Eyy1[t][k]) + ey1[t][k]'*Inverse(Eyy1[t][k])*ey1[t][k])

  where k = xi(t) and ey1[t][k] = y(t) - Ey1[t][k].


Recursive smoothing relations:

  


********************************************************************************/
