%   
% 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',10); rand('seed',10);

a_truth = 0.5;
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 );

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');

f3 = figure();
hx3 = plot( ts, a_truth * ones(size(ts)), '-g' ); hold on;
xlabel('time (sec)'); ylabel( 'estimate of a' );
legend( [hx3], {'true value a'}, 'location', 'southwest' );
%saveas( gcf, '../../WriteUp/Graphics/Chapter4/sect_7_prob_4_plot_x3', 'epsc' ); 

% Run the parallel Kalman filters:
%
Phi_1 = Phi_truth; Phi_1(2,1) = 0.6;
[xplus_1, Pplus_1, VV, loglik, innovations_1, xminus_1, Vpred, S_1] = kalman_filter(z, Phi_1, H, Q, R, zeros(n,1), eye(n,n) ); 
Phi_2 = Phi_truth; Phi_2(2,1) = 0.5;
[xplus_2, Pplus_2, VV, loglik, innovations_2, xminus_2, Vpred, S_2] = kalman_filter(z, Phi_2, H, Q, R, zeros(n,1), eye(n,n) ); 
Phi_3 = Phi_truth; Phi_3(2,1) = 0.4;
[xplus_3, Pplus_3, VV, loglik, innovations_3, xminus_3, Vpred, S_3] = kalman_filter(z, Phi_3, H, Q, R, zeros(n,1), eye(n,n) ); 

% Evaluate the likelihood of each parameter setting: pr[z_k|\hat{x}[p_i,(-)]]
%
like_1 = zeros(1,N);
like_2 = zeros(1,N);
like_3 = zeros(1,N);
for k=1:N+1
  like_1(k) = mvnpdf( innovations_1(:,k), zeros(n,1), S_1(:,:,k) );
  like_2(k) = mvnpdf( innovations_2(:,k), zeros(n,1), S_2(:,:,k) );
  like_3(k) = mvnpdf( innovations_3(:,k), zeros(n,1), S_3(:,:,k) );
end
all_likes = [ like_1; like_2; like_3 ];

I = 3;

% the prior probability of each hypothesis:
Pr0 = ones(I,1)/I;
Pr0 = [ 0.9, 0.05, 0.05 ]'; 

% recursivly compute the posteriori probabilty of each parameter:
Pr = zeros(I,N+1);
for k=1:(N+1)
  if( k==1 )
    prior = Pr0;
  else
    prior = Pr(:,k-1);
  end
  
  Pr(:,k) = all_likes(:,k) .* prior / sum(all_likes(:,k) .* prior);
  
end

f4 = figure(); hold on; colors='rgbkcm';
all_hs = [];
for ii=1:I
  all_hs = [ all_hs, plot( Pr(ii,:), ['-',colors(ii)] ) ];
end
legend( all_hs, {'Phi\_1','Phi\_truth','Phi\_3'}, 'location', 'best' );
xlabel('time'); ylabel('probability');
%saveas( gcf, '../../WriteUp/Graphics/Chapter4/sect_7_prob_4_plot_probs', 'epsc' ); 

% Next plot the a posteriori state estimates (best guess after averaging over all parameters): 
%
mean_xhatplus = zeros(n,N+1);
for k=1:(N+1)
  mean_xhatplus(:,k) = Pr(1,k) * xplus_1(:,k) + Pr(2,k) * xplus_2(:,k) + Pr(3,k) * xplus_3(:,k);
end

figure(f1);
hxhat1 = plot( ts, mean_xhatplus(1,:), '-r' ); 
legend( [hx1,hz1,hxhat1], {'true x1', 'measurement','average estimate'}, 'location', 'southwest' );
%%%saveas( gcf, '../../WriteUp/Graphics/Chapter4/sect_7_prob_4_plot_x1', 'epsc' ); 

figure(f2);
hxhat2 = plot( ts, mean_xhatplus(2,:), '-r' ); 
legend( [hx2,hz2,hxhat2], {'true x2', 'measurement','average estimate'}, 'location', 'southwest' );
%%%saveas( gcf, '../../WriteUp/Graphics/Chapter4/sect_7_prob_4_plot_x2', 'epsc' ); 

% Next plot the estimate value of omega_n:
%
mean_21_element = zeros(1,N+1);
for k=1:(N+1)
  mean_21_element(k) = Pr(1,k) * (0.6) + Pr(2,k) * (0.5) + Pr(3,k) * (0.4);
end

figure(f3);
hxhat3 = plot( ts, mean_21_element, '-r' ); 
legend( [hx3,hxhat3], {'true value (2,1) element', 'mean estimate of (2,1) element'}, 'location', 'best' );
%saveas( gcf, '../../WriteUp/Graphics/Chapter4/sect_7_prob_4_plot_x3', 'epsc' );