/**********************************************************
FILTER.SRC

This is a Multivariate Kalman filtering routine that 
handles missing observations. It is a linear filter, and
it has been augmented with a fixed interval smoother.
It is based on that described in Harvey (1989) "Forecasting
Structural Time Series Models and the Kalman Filter". It 
generally complies with the notation established in that
text.

Author: Geoffrey Shuetrim
Geoffrey.Shuetrim@apra.gov.au

Change Log:
15/9/99 Christopher Thompson improved the treatment of the
initial state vector and its precision matrix.

This code comes with no guarantees for performance. Use it
entirely at your own risk.

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


/**********************************************************
Initialise the global system matrices and filtration
output matrices
**********************************************************/
_y = zeros(_n,_obs);
_Z = zeros(_k*_n,_obs);
_d = zeros(_n,_obs);
_H = zeros(_n*_n,_obs);

_T = zeros(_k*_k,_obs);
_c = zeros(_k,_obs);
_R = zeros(_k*_s,_obs);
_Q = zeros(_s*_s,_obs);

_a  = zeros(_k,_obs);
_a0 = zeros(_k,1);
_ap = _a;
_as = _a;

_P  = zeros(_k*_k,_obs);
_P0 = zeros(_k*_k,1);
_Pp = _P;
_Ps = _P;

_LL = zeros(1,_obs);

@ Make forward declarations of procecdures @
external proc SysMat;

_infinity = 1e5;

/**********************************************************
Define the filtration

The process for each observation except the first:
(Note that the first observation is the prior!!!)

1.	Form the system matrices for the current period

2.	Get the state and precision matrix predictions for
	the current observation, based on date up to the
	previous observation.

3.	Compute the prediction errors and the prediction 
	error variances.

4.	Compute the updated state and precision matrices 
	and store them. Also store the prediction error 
	variances and the prediction errors themselves.

Note that all of the filtration output is stored in global 
matrices which makes the filter calls faster at the cost
of some robustness of the code!  This means that there
are no return parameters to worry about.

**********************************************************/
proc(0) = Filter;

  local t; @ loop counter @

  @ local versions of the relevant system matrices @
  local Ppt,Tt,Ht,Qt,Rt,Zt,Xt,Ft,Ft_inverse,vt;
  
  Zt = reshape(_Z[.,1],_n,_k);
  Ht = reshape(_H[.,1],_n,_n);
  Tt = reshape(_T[.,1],_k,_k);
  Rt = reshape(_R[.,1],_k,_s);
  Qt = reshape(_Q[.,1],_s,_s);
  
  _ap[.,1] = Tt * _a0 + _c[.,1];
  Ppt = Tt * _P0 * Tt' + Rt * Qt * Rt';
  _Pp[.,1] = vecr(Ppt);
  
  local Wt,mask;		
  mask = _y[.,1] ./= miss(0,0);
  Wt = selif(diagrv(zeros(_n,_n),mask),mask); 	
  
  Ft = Wt * (Zt * Ppt * Zt' + Ht) * Wt';
  Ft_inverse = invpd(Ft);
  vt = Wt * (missrv(_y[.,1],0) - Zt * _ap[.,1] - _d[.,1]);
  
  _a[.,1] = _ap[.,1] + Ppt * Zt' * Wt' * Ft_inverse * vt;
  _P[.,1] = vecr(Ppt - Ppt' * Zt' * Wt' * Ft_inverse * Wt * Zt * Ppt);
  _LL[.,1] = -0.5*(sumc(mask)*ln(2*pi)+ln(det(Ft))+vt'*Ft_inverse*vt);
  
  for t (2,cols(_y),1);
    
    Zt = reshape(_Z[.,t],_n,_k);
    Ht = reshape(_H[.,t],_n,_n);
    Tt = reshape(_T[.,t],_k,_k);
    Rt = reshape(_R[.,t],_k,_s);
    Qt = reshape(_Q[.,t],_s,_s);
    
    _ap[.,t] = Tt * _a[.,t-1] + _c[.,t];
    Ppt = Tt * reshape(_P[.,t-1],_k,_k) * Tt' + Rt * Qt * Rt';
    _Pp[.,t] = vecr(Ppt);
    
    mask = _y[.,t] ./= miss(0,0);
    Wt = selif(diagrv(zeros(_n,_n),mask),mask); 
    
    if not scalmiss(Wt); @ we have some current data @
      
      Ft = Wt * (Zt * Ppt * Zt' + Ht) * Wt';
      Ft_inverse = invpd(Ft);
      vt = Wt * ( missrv(_y[.,t],0) - Zt * _ap[.,t] - _d[.,t]);
      
      _a[.,t] = _ap[.,t] + Ppt * Zt' * Wt' * Ft_inverse * vt;
      _P[.,t] = vecr(Ppt - Ppt' * Zt' * Wt' * Ft_inverse * Wt * Zt * Ppt);
      _LL[.,t] = -0.5 * (sumc(mask)*ln(2*pi) + ln(det(Ft)) + vt' * Ft_inverse * vt);
      
    else; @ we have no data for the current period @
      
      _a[.,t] = _ap[.,t];
      _P[.,t] = _Pp[.,t];
      _LL[.,t] = 0;
      
    endif;
    
  endfor;
  
endp;

/**********************************************************
Define the fixed interval smoother

Note: Assumes that the filtering has been completed!
**********************************************************/
proc(0) = Smoother;

  _as[.,cols(_y)] = _a[.,cols(_y)];
  _Ps[.,cols(_y)] = _P[.,cols(_y)];
  
  @ local versions of the relevant system matrices @
  local Pstar,Pt,Ppt_lead;
  
  local t; @ loop counter @
  
  for t (cols(_y)-1,1,-1);
    
    Pt = reshape(_P[.,t],_k,_k);
    Ppt_lead = reshape(_Pp[.,t+1],_k,_k);
    Pstar = Pt * reshape(_T[.,t+1],_k,_k)' * invpd(Ppt_lead);
    
    _Ps[.,t] = vecr(Pt + Pstar * (reshape(_Ps[.,t+1],_k,_k) - Ppt_lead) * Pstar');
    
    _as[.,t] = _a[.,t] + Pstar * (_as[.,t+1] - _ap[.,t+1]);
    
  endfor;
  
endp;

/**********************************************************
Loglikelihood computation procedure
**********************************************************/
proc(1) = LL(p,d);

  call SysMat(p);
  call Filter;
  retp(_LL');
  
endp;
