% EE278B, Autumn 2011
% Bernd Bandemer/Abbas El Gamal, November 2011
% Vector Kalman filter example, prediction case

function p5
    A = [1 1; 0 0.9];
    
    n = 100;  % time horizon
    P1 = 1000;
    P2 = 0;  % initial state covariance matrix P is diag([P1,P2])
    Q1 = 0;
    Q2 = 10;  % state noise covariance matrix Q is diag([Q1,Q2])
    N1 = 1000;
    N2 = 1;  % observation noise covariance matrix N is diag([N1,N2])
        
    % generate a realization
    % because it is more convenient with Matlab's own indexing, 
    % we start time at index 1 (unlike lecture notes, where time starts at 0)

    X = zeros( 2, n );  
    Y = zeros( 2, n );
    
    X(:,1) = [sqrt(P1)*randn(1); sqrt(P2)*randn(1) ];
    for i=1:n-1
        Ui = [ sqrt(Q1)*randn(1); sqrt(Q2)*randn(1) ];
        X(:,i+1) = A*X(:,i) + Ui;
    end
    for i=1:n
        Y(:,i) = X(:,i) + [sqrt(N1)*randn; sqrt(N2)*randn];
    end
    
    % run the Kalman filter
    Xhatp = vectorKalmanPredictor(...
        repmat(A, [1 1 n]), ...
        diag([P1,P2]), ...
        repmat(diag([Q1,Q2]), [1 1 n]), ...
        repmat(diag([N1,N2]), [1 1 n]), ...
        Y );
    
    % make plots
    close all;
    
    figure; 
    subplot(211); hold on; grid on;
    plot(X(1,:), 'r+-');
    plot(Xhatp(1,:), 'bx-');
    xlim([0,n+1]);
    ylabel('State 1: Xhat_{i|i-1}, X_{i}');
    legend('Actual state','Predicted state','Location','Best');
    
    subplot(212); hold on;  grid on;
    plot(X(2,:), 'r+-');
    plot(Xhatp(2,:), 'bx-');
    xlim([0,n+1]);
    xlabel('time i');
    ylabel('State 2: Xhat_{i|i-1}, X_{i}');
    
    print -dpdf vectorKalman_fig1
    
    figure; 
    subplot(211); hold on;  grid on;
    plot(Xhatp(1,1:n)-X(1,:), 'b-');
    ylabel('State 1: Xhat_{i|i-1} - X_{i}');
    xlim([0,n]);
    
    subplot(212); hold on;  grid on;
    plot(Xhatp(2,1:n)-X(2,:), 'b-');
    xlim([0,n]);
    xlabel('time i');
    ylabel('State 2: Xhat_{i|i-1} - X_{i}');
    
    print -dpdf vectorKalman_fig2
end

function [Xhatp,SigmaXhatp] = vectorKalmanPredictor(As,P1,Qs,Ns,Ys)
% inputs: 
%   A (s x s x n) - system matrix for times 1:n
%   P1 (s x s) - covariance matrix of state X_1
%   Q (s x s x n) - covariance matrix of state noise at times 1:n
%   N (s x s x n) - covariance matrix of observation noise at times 1:n
%   Y (s x n) - observations for times 1:n
% outputs: 
%   Xhatp (s x n+1) - predicted state. 
%       hatXp(:,i+1) is Xhat(i+1|i)
%       hatXp(:,1) is Xhat(1|0)
%       hatXp(:,2) is Xhat(2|1), etc.
%   SigmaXhatp (s x s x n) - MSE matrix of the prediction
%       SigmaXhatp(:,:,i+1) is Sigma(i+1|i)
    
% Declare variables
[s, ~, n] = size(As);
Xhatp = zeros(s,n+1);
SigmaXhatp = zeros(s,s,n+1);
    
% Initialization
Xhatp(:,1) = zeros(s,1);
SigmaXhatp(:,:,1) = P1;

% Apply Kalman prediction
for i = 2:(n+1)
    A = As(:,:,i-1);
    Q = Qs(:,:,i-1);
    N = Ns(:,:,i-1);
    Y = Ys(:,i-1);
    
    SigmaPrev = SigmaXhatp(:,:,i-1);
    K = A*SigmaPrev/(SigmaPrev+N);
    
    Xprev = Xhatp(:,i-1);
    Xhatp(:,i) = A*Xprev + K*(Y-Xprev);
    SigmaXhatp(:,:,i) = A*SigmaPrev*(eye(s)-(SigmaPrev+N)\SigmaPrev)*A'+Q;
end

SigmaXhatp = SigmaXhatp(:,:,1:(end-1));

end