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

close all; clc; clear; randn('seed',0); rand('seed',1);

a_truth = 0.5;
%a_guess = 0.5; % make sure that this matches the optimal results
a_guess = 0.6;
a_std = 0.2;
%a_std = 0.0; % make sure that this matches the optimal results
Phi_truth = [ 0.9, 1; a_truth, 0.8 ]; 

n = size(Phi_truth,1);
Q = eye(n);
R = 0.1 * eye(2);

H = eye(2);

dt = 0.05;
T = 1.; % total time of simulation
N = T/dt; % number of timesteps 
ts = 0:dt:T;

addpath( '../../../FullBNT-1.0.4/Kalman/' );
addpath( '../../../FullBNT-1.0.4/KPMstats/' );

x0 = [1;-1];
P0 = zeros(n,n);
[x_truth,z] = sample_lds( Phi_truth, H, Q, R, x0, N+1 ); % x_truth(1) and z(1) correspond to the same point in time.

f1 = figure();
hx1 = plot( ts, x_truth(1,:), '-g' ); hold on; 
hz1 = plot( ts, z(1,:), 'xk' ); 
legend( [hx1,hz1], {'state 1','measurment 1'}, 'location', 'best' );
xlabel('time (sec)'); ylabel('x1');

f2 = figure();
hx2 = plot( ts, x_truth(2,:), '-g' ); hold on; 
hz2 = plot( ts, z(2,:), 'xk' ); 
legend( [hx2,hz2], {'state 2','measurment 2'}, 'location', 'best' );
xlabel('time (sec)'); ylabel('x2');

% Part (a): Perform Kalman filtering on this system:
%
if( 0 ) 
  [xhat, V, VV, loglik, innovations, xpred, Vpred, S, Sinv] = kalman_filter(z, Phi_truth, H, Q, R, x0, P0);
  
  % Plot the Kalman filtering results 
  figure(f1);
  hx1hat = plot( ts, xhat(1,:), '-or' ); 
  legend( [hx1,hz1,hx1hat], {'state 1','measurment 1','state 1 Kalman estimate'}, 'location', 'best' );
  %saveas( gcf, '../../WriteUp/Graphics/Chapter4/sect_7_prob_1_x1_plot_a.eps', 'epsc' );
  
  figure(f2);
  hx2hat = plot( ts, xhat(2,:), '-or' ); 
  legend( [hx2,hz2,hx2hat], {'state 2','measurment 2','state 2 Kalman estimate'}, 'location', 'best' );
  %saveas( gcf, '../../WriteUp/Graphics/Chapter4/sect_7_prob_1_x2_plot_a.eps', 'epsc' );
end
  
% Part (b): Perform Kalman filtering but using the incorrect Phi:
%
if( 0 )
  Phi_model = Phi_truth;
  Phi_model(2,1) = a_guess;

  [xhat, V, VV, loglik, innovations, xpred, Vpred, S, Sinv] = kalman_filter(z, Phi_model, H, Q, R, x0, P0);
  
  % Plot the Kalman filtering results 
  figure(f1);
  hx1hat = plot( ts, xhat(1,:), '-or' ); 
  legend( [hx1,hz1,hx1hat], {'state 1','measurment 1','state 1 Kalman estimate'}, 'location', 'best' );
  %saveas( gcf, '../../WriteUp/Graphics/Chapter4/sect_7_prob_1_x1_plot_b.eps', 'epsc' );
  
  figure(f2);
  hx2hat = plot( ts, xhat(2,:), '-or' ); 
  legend( [hx2,hz2,hx2hat], {'state 2','measurment 2','state 2 Kalman estimate'}, 'location', 'best' );
  %saveas( gcf, '../../WriteUp/Graphics/Chapter4/sect_7_prob_1_x2_plot_b.eps', 'epsc' );
end

% Part (c): Parameter adaptive filtering for the (2,1) component of Phi denoted as "a":
%
if( 1 ) 
  z = z(:,2:end); % drop the measurement from the initial state x_0
  
  f3 = figure;
  hx3 = plot( ts, a_truth * ones(N+1,1), '-g' ); hold on;
  legend( [hx3], {'state 3 (a)'}, 'location', 'best' );
  xlabel('time (sec)'); ylabel('a');

  % We use the extended Kalman filter on this system:
  %
  n_A = 3; % the extended state dimension
  x_0 = [ 1; -1; a_guess ]; % the extended state initial conditions
  P_0 = zeros(n_A,n_A); P_0(n_A,n_A) = a_std^2;
  H_A = zeros(2,3); H_A(1,1) = 1; H_A(2,2) = 1; 
  Lambda_A = zeros(3,2); Lambda_A(1,1) = 1; Lambda_A(2,2) = 1; 
  Q_A = Q; 
  R_A = R;

  [xhatminus,pminus,xhatplus,pplus] = sect_7_prob_1_extended_kf(x_0, P_0, Phi_truth, H_A, Lambda_A, Q_A, R_A, z, ts);

  % Now plot the a posteriori state estimates: 
  %
  figure(f1);
  hxhat1 = plot( ts, xhatplus(1,:), '-r' ); 
  %s11 = sqrt( squeeze(pplus(1,1,:)) )';
  %plot( ts, xhatplus(1,:) - 1.97 * s11, '-.r' ); 
  %plot( ts, xhatplus(1,:) + 1.97 * s11, '-.r' ); 
  legend( [hx1,hz1,hxhat1], {'x_1', 'measurement','x_1 estimate'}, 'location', 'best' );
  %saveas( gcf, '../../WriteUp/Graphics/Chapter4/sect_7_prob_1_part_c_x1_plot', 'epsc' ); 

  figure(f2);
  hxhat2 = plot( ts, xhatplus(2,:), '-r' ); 
  legend( [hx2,hz2,hxhat2], {'x_2', 'measurement','x_2 estimate'}, 'location', 'best' );
  %saveas( gcf, '../../WriteUp/Graphics/Chapter4/sect_7_prob_1_part_c_x2_plot', 'epsc' ); 

  figure(f3);
  hxhat3 = plot( ts, xhatplus(3,:), '-r' ); 
  %s33 = sqrt( squeeze(pplus(3,3,:)) )';
  %plot( ts, xhatplus(3,:) - 1.97 * s33, '-.r' ); 
  %plot( ts, xhatplus(3,:) + 1.97 * s33, '-.r' ); 
  legend( [hx3,hxhat3], {'true value a', 'estimate of a'}, 'location', 'best' );
  %saveas( gcf, '../../WriteUp/Graphics/Chapter4/sect_7_prob_1_part_c_a_plot', 'epsc' ); 
end