%--------------------------------------------------------------------------
% recursive "Extended Kalman Filter" computation
%--------------------------------------------------------------------------

clear all
close all
clc

enc_cnt = 12;

dt       = 16e-3;
duration = 30;

wb    = 1;
rpm   = 8;
omega = rpm * pi/30;

q = 2^-15;
r = 0.01;

Q = q * eye(3);
R = r * eye(2);

P_hat   = zeros(3);
P_tilde = zeros(3);

H = [1 0 0;
     0 1 0];

y = [0; 0];
x = [0; 0; 0];
     
x_hat   = x;
x_tilde = x;
y_tilde = y;

x_hat_mem   = [];
x_tilde_mem = [];
y_mem       = [];
y_tilde_mem = [];


for t = 0 : dt : duration
    
    %**********************************************************************
    % simulate encoder
    %**********************************************************************
    angle = t * omega;
    tmp   = mod(angle, 2*pi);
    theta = tmp - mod( tmp, ((2*pi) / enc_cnt) );
    
    y = [cos(theta);
         sin(theta)];

     
    %**********************************************************************
    % F = df/dx | x=x_hat
    %**********************************************************************
    F = [cos(wb * x_hat(3) * dt) -sin(wb * x_hat(3) * dt) 0;
         sin(wb * x_hat(3) * dt)  cos(wb * x_hat(3) * dt) 0;
                              0                        0  1];
                          
    %**********************************************************************
    % prediction
    %**********************************************************************
%   x_tilde = [x_hat(1)*cos(wb*dt*x_hat(3)) - x_hat(2)*sin(wb*dt*x_hat(3));
%              x_hat(1)*sin(wb*dt*x_hat(3)) + x_hat(2)*cos(wb*dt*x_hat(3));
%                                                                x_hat(3)];
    
    x_tilde = F * x_hat;
    y_tilde = H * x_tilde;

    
    %**********************************************************************
    % kalman gain
    %**********************************************************************
    K = P_tilde * H' * inv( H * P_tilde * H' + R );
    
    %**********************************************************************
    % estimation
    %**********************************************************************
    x_hat = x_tilde + K * (y - y_tilde);
    
    %**********************************************************************
    % prediction error covariance matrix
    %**********************************************************************
    P_tilde = F * P_hat * F' + Q;

    %**********************************************************************
    % estimation error covariance matrix
    %**********************************************************************
    P_hat = P_tilde - K * H * P_tilde;
    
    %**********************************************************************
    % save values for plotting
    %**********************************************************************
    x_hat_mem   = [x_hat_mem x_hat];
    x_tilde_mem = [x_tilde_mem x_tilde];
    
    y_mem       = [y_mem y];
    y_tilde_mem = [y_tilde_mem y_tilde];

end

t = 0 : dt : duration;


figure

subplot(211);
hold on
stairs(t, x_tilde_mem(1,:), 'b');
stairs(t, x_tilde_mem(2,:), 'r');

subplot(212);
hold on
stairs(t, x_hat_mem(1,:), 'b');
stairs(t, x_hat_mem(2,:), 'r');


figure

subplot(211);
hold on
stairs(t, y_mem(1,:), 'b');
stairs(t, y_tilde_mem(1,:), 'r');

subplot(212);
hold on
stairs(t, y_mem(2,:), 'b');
stairs(t, y_tilde_mem(2,:), 'r');

figure
subplot(211); stairs(t, x_tilde_mem(3,:), 'b');
subplot(212); stairs(t, x_hat_mem(3,:), 'r');