% 
% 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;0]; 
sigma2_0 = 100 * eye(2); 

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

% the measurement noise variance: 
sigma2_v = 1.0; % 0.2; 0.5; %1.0;

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

% dynamic equation coefficient: 
F = [ 0.5, 0.5; 1, 0 ]; 

% measurement coefficient: 
H = [ 1, 0 ];

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

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

% storage: 
% 
state_dim       = 2; 
measurement_dim = 1; 
zhat            = zeros(measurement_dim,Ntimesteps); 
zi              = zeros(measurement_dim,Ntimesteps); 
xip1_given_i    = zeros(state_dim,Ntimesteps); 
cip1_given_i    = zeros(state_dim,state_dim,Ntimesteps); 
s               = zeros(1,Ntimesteps); 
k               = zeros(state_dim,Ntimesteps); 
xi_given_i      = zeros(state_dim,Ntimesteps); 
ci_given_i      = zeros(state_dim,state_dim,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(state_dim,Ntimesteps);     % holds the true state values
x_truth(1,1)    = 0.5 * w_truth(1);                % x_truth(0) = 0.5 * ( w(0) + w(-1) ) = 0.5 * w(0) 
x_truth(2,1)    = 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(1,ti+1) = 0.5 * ( w_truth(ti+1) + w_truth(ti) ); 
    x_truth(2,ti+1) = x_truth(1,ti); 
  end

end

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

tt=sprintf('noise variance = %10.2f',sigma2_v); title( tt ); 
fn=sprintf('../../WriteUp/Graphics/Chapter4/ex_10_kf_examples_nv_%.2f.eps',sigma2_v); 
saveas( gcf, fn, 'epsc' );