% 
% 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; 

fh = figure; hold on; 

% do a representative number of these problems: 
% 
%rand('seed',0); randn('seed',0); 

% initial conditions/distribution assumed for the state x(0): 
mu_0 = 0; 
sigma2_0 = 100; 

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

% the measurement noise variance: 
sigma2_v = 1.; 

% the dynamic model noise variance: 
sigma2_w = 1;
C_w      = sigma2_w * eye(2); 

% dynamic equation coefficient: 
F = 0; 

% measurement coefficient: 
H = 1; 

% dynamic noise coefficient:
G = 0.5 * [ 1, 1 ]; 

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

% storage: 
% 
zhat            = zeros(1,Ntimesteps); 
zi              = zeros(1,Ntimesteps); 
xip1_given_i    = zeros(1,Ntimesteps); 
cip1_given_i    = zeros(1,Ntimesteps); 
s               = zeros(1,Ntimesteps); 
k               = zeros(1,Ntimesteps); 
xi_given_i      = zeros(1,Ntimesteps); 
ci_given_i      = zeros(1,Ntimesteps); 

w_truth         = zeros(1,Ntimesteps);             % the true disturbances that make up x(i) = 0.5 * ( w(i) + w(i-1) ) 
w_truth(1)      = mvnrnd(0,sigma2_w);
x_truth         = zeros(1,Ntimesteps);             % holds the true state values
x_truth(1)      = 0.5 * w_truth(1);                % x_truth(0) = 0.5 * ( w(0) + w(-1) ) = 0.5 * w(0) 

% two ways to seed/start the Kalman filter: (i=0 case) 
if( 1 ) 
  % -- 1: explicitly seed the Kalman filter with a zero expectation: 
  % 
  xip1_given_i(1) = mu_0;              % i=0 case 
  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 ;
  k(ti)          = cip1_given_i(ti) * (H.') / ( s(ti)+eps ); 
  zi(ti)         = H * x_truth(ti) + mvnrnd(0,sigma2_v);          % 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 propagation (predicted):
  % 
  xip1_given_i(ti+1) = F * xi_given_i(ti); 
  cip1_given_i(ti+1) = F * ci_given_i(ti) * (F.') + G * C_w * (G.'); 
   
  % and the actual state propogation:
  if( ti<Ntimesteps )
    w_truth(ti+1) = mvnrnd(0,sigma2_w);
    x_truth(ti+1) = 0.5 * ( w_truth(ti+1) + w_truth(ti) ); 
  end

end

figure(fh); 
%plot( xip1_given_i ); hold on; 
ha = plot( xi_given_i, '-k' ); hold on; 
ht = plot( x_truth, '-g' ); 
hm = plot( zi, 'or' ); axis tight; 
legend( [ha,ht,hm], {'Kalman estimate', 'Truth', 'Measurements' } ); 

saveas( gcf, '../../WriteUp/Graphics/Chapter4/ex_9_kf_examples', 'epsc' );