% 
% Written by:
% -- 
% John L. Weatherwax                2009-02-24
% 
% email: wax@alum.mit.edu
% 
% Please send comments and especially bug reports to the
% above email address.
% 
%-----

%close all; drawnow; 
clc; 
clear; 

rand('seed',0); randn('seed',0); 

% initial conditions/distribution assumed for the state x(0): 
mu_0 = [ 0; 0 ]; 
sigma2_0 = 100*eye(2); 

% compute the ACTUAL initial value (TRUTH) of our state x(i):
x0 = mvnrnd(mu_0,sigma2_0); 
x0 = x0.'; 

% the dynamic model propogation matrix 
F = [ 1/2, 1/4; 1, 0 ]; 

% the dynamic model noise variance: 
C_w = [ 1, 0 ; 0, 0 ]; 

% the measurement "spread" matrix 
H = [ 1, 0 ]; 

% the measurement noise variance: 
sigma2_v = 1; 

% My indexing convention: 
%   Matlab       Kalman Filtering Index
%     i=1   <=>         i=0
% 
%
Ntimesteps = 300;

% storage: 
% 
state_dim       = 2; 
measurement_dim = 1; 
zhat            = zeros(measurement_dim,Ntimesteps);       % predicted measurements 
zi              = zeros(measurement_dim,Ntimesteps);       % actual measurments 
xip1_given_i    = zeros(state_dim,Ntimesteps);             % predicted state 
cip1_given_i    = zeros(state_dim,state_dim,Ntimesteps);   % predicted covariance 
s               = zeros(1,Ntimesteps);                     % innovation matrix 
k               = zeros(state_dim,Ntimesteps);             % kalman gain matrix 
xi_given_i      = zeros(state_dim,Ntimesteps);             % a-posteriori state
ci_given_i      = zeros(state_dim,state_dim,Ntimesteps);   % a-posteriori covariance

x_truth         = zeros(state_dim,Ntimesteps);             % the true state value
x_truth(:,1)    = x0; 

for bi=1:3, 
  % two ways to seed/start the Kalman filter: (i=0 case) 

  if( 0 ) 
    % -- 1: explicitly seed the Kalman filter with a zero expectation: 
    % 
    xip1_given_i(:,1) = mu_0;                              
    cip1_given_i(:,:,1) = sigma2_0;
  else
    % -- 2: take a random initial condition (around zero): 
    % 
    xip1_given_i(:,1) = mvnrnd(mu_0,sigma2_0).'; 
    cip1_given_i(:,:,1) = sigma2_0;
  end

  for ti=1:Ntimesteps,

    % Estimate mean/covariance taking the observed measurement at ti into account 
    %
    zhat(:,ti)         = H * xip1_given_i(:,ti);
    s(ti)              = H * cip1_given_i(:,:,ti) * (H.') + sigma2_v * bi/2; 
    k(:,ti)            = cip1_given_i(:,:,ti) * (H.') / ( s(ti)+eps ); 
    zi(:,ti)           = H * x_truth(:,ti) + mvnrnd(0,sigma2_v * bi/2);          % simulate a measurement at timestep "ti" ... from the "truth"
    xi_given_i(:,ti)   = xip1_given_i(:,ti) + k(:,ti) * ( zi(:,ti) - zhat(:,ti) ); 
    ci_given_i(:,:,ti) = cip1_given_i(:,:,ti) - k(:,ti) * s(ti) * (k(:,ti).');

    % Perform state propogation (predicted):
    % 
    xip1_given_i(:,ti+1)   = F * xi_given_i(:,ti); 
    cip1_given_i(:,:,ti+1) = F * ci_given_i(:,:,ti) * (F.') + C_w; 

    % and the actual state propogation:
    if( ti<Ntimesteps )
      x_truth(:,ti+1) = F * x_truth(:,ti) + mvnrnd([0,0],C_w).';  
    end

  end

  %figure; 
  if( bi==1 )
    off=0; 
  elseif( bi==2 )
    off=8;
  else
    off=16;
  end
  %he = plot( xip1_given_i(1,:)+off, '-r' ); hold on; ht = plot( x_truth(1,:)+off, '-g' ); 
  he = plot( xi_given_i(1,:)+off, '-r' ); hold on; ht = plot( x_truth(1,:)+off, '-g' ); 
  hm = plot( zi+off, '.' ); 
  legend( [ he, ht, hm ], {'estimate', 'truth', 'measurement'}, 'location', 'north' );
  %title( 'truth and predicted value of x(1)' ); 
  axis tight; 
end
%saveas( gcf, '../../WriteUp/Graphics/Chapter4/ex_8_kf_example', 'epsc' ); 
saveas( gcf, '../../WriteUp/Graphics/Chapter4/ex_8_kf_examples', 'epsc' );