%--------------------------------------------------------------------------
% 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.1;

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

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

f = [];
h = [1 0 0;
     0 1 0];

y = [0; 0];
x = [0.001; 0.001; 0.001];

x_hat   = x;
x_tilde = x;
y_tilde = y;

x_mem       = [];
y_mem       = [];
x_tilde_mem = [];
y_tilde_mem = [];
x_hat_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)];
    
    %**********************************************************************
    % prediction
    %**********************************************************************
    x_tilde_old = x_tilde;
    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)];
    
    y_tilde_old = y_tilde;
    y_tilde = h * x_tilde;
    
    diff_x = x_tilde_old - x_tilde;
    diff_y = y_tilde_old - y_tilde;
    
    %**********************************************************************
    % F = df/dx | x=x_hat     H = dh/dx | x=x_hat
    %**********************************************************************
    % F = diff_x / x_hat;
    F = [diff_x(1)/x_hat(1) diff_x(1)/x_hat(2) diff_x(1)/x_hat(3);
         diff_x(2)/x_hat(1) diff_x(2)/x_hat(2) diff_x(2)/x_hat(3);
         diff_x(3)/x_hat(1) diff_x(3)/x_hat(2) diff_x(3)/x_hat(3)];  

    % H = diff_y / x_hat;
    H = [diff_y(1)/x_hat(1) diff_y(1)/x_hat(2) diff_y(1)/x_hat(3);
         diff_y(2)/x_hat(1) diff_y(2)/x_hat(2) diff_y(2)/x_hat(3)];
    % H = h;
    
    %**********************************************************************
    % 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
    %**********************************************************************
    y_mem       = [y_mem y];
    x_tilde_mem = [x_tilde_mem x_tilde];
    y_tilde_mem = [y_tilde_mem y_tilde];
    x_hat_mem   = [x_hat_mem x_hat];

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,:));
stairs(t, y_mem(2,:));

subplot(212);
hold on
stairs(t, y_tilde_mem(1,:));
stairs(t, y_tilde_mem(2,:));

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