%	Test EKF and UKF for 'Random Wally Walk'
clear;
clc;

%%  Configuration
n = 5;          % number of states
T = 3;          % time constant for Wally

N = 200;        % total dynamic steps

qt = 0.6;       % true std of process 
rt = 0.3;       % true std of measurement

q = 0.6;        % std of process (assumed value)
r = 0.3;        % std of measurement (assumed value)
Q = q^2*eye(n); % covariance of process
R = r^2;        % covariance of measurement  

%%  Preparation

% allocate memory
sR = zeros(n,N);    % reference
sV = zeros(n,N);    % reference + process noise
zV = zeros(2,N);    % measurement
xV = zeros(n,N);    % estimate for EKF
xVu = zeros(n,N);    % estimate for UKF


%%  1. Step - Create a (ideal) Reference Simulation of our Dynamic System
%
%   This will represent our System as it would behave in an ideal world.
%   The System-Behaviour is described by the non-linear function f(x,u).
%   The System-Input is u, The System-States x.

f = @(x,u)[x(2); -1/x(5) * x(2) + u(2); x(4); -1/x(5) * x(4) + u(1); x(5)];
              % non-linear state equation for 'random wally walk'

% Simulate a walk of N steps where the input is a [0,1] random variable
u=randn(N,2);


s_init = [0;0;1;0;T];
              % initial state

s = s_init;

% Now let's calculate all N simulation steps for our discrete System
for k = 1:N
  sR(:,k)= s;                      % save actual state
  s = f(s,u(k,:));                 % update process 
end


%%  2. Step - Add Process Noise to the Dynamic System
%
%   We simulated the ideal behaviour of our System in Step 1.
%   Because of NOT ideal conditions, we have deviations between
%   ideal and non-ideal behaviour. This is to be represented here
%   with the true random process noise qt.

for k = 1:N
  sV(:,k)= sR(:,k) + [qt * randn(n-1,1); 0];    % save state of reference + process noise
end


%%  3. Step - Make Measurement's (with Measurement Noise)
%
%   Our measurement's are defined in our measurement equation h(x,u).
%   

h=@(x,u)[x(1); x(3)];
              % measurement equation
              % we will be able to make a direct measure of
              % only the x/y positions


for k = 1:N
  z = h(sV(:,k),u(k,:)) + rt * randn(2,1);   % make measurement from the reference + process noise by passing it to h(x,u)
                                             % additionally we have to add the true measurement noise rt

  zV(:,k)  = z;                              % save measurments
end


%%  4. Step - Do the Filtering with the EKF (Extended Kalman Filter)
%
%

x = s_init + [qt * randn(n-1,1); 0];
              % initial state with process noise

P = eye(n);
              % initial state covariance
              
for k = 1:N
  [x, P] = ekf(f,x,u(k,:),P,h,zV(:,k),Q,R);     % ekf 
  xV(:,k) = x;                                  % save estimate
end


%%  5. Step - Do the Filtering with the UKF (Unscented Kalman Filter)
%
%

x = s_init + [qt * randn(n-1,1); 0];
              % initial state with process noise

P = eye(n);
              % initial state covariance
              
for k = 1:N
  [x, P] = ukf(f,x,u(k,:),P,h,zV(:,k),Q,R);     % ukf
  xVu(:,k) = x;                                  % save estimate
end


%%  6. Step - Plot
%
%   Make Plots of everything done above.
%   

subplot(n,1,1)
plot(1:N, sR(1,:), 'r-', 1:N, sV(1,:), 'b-.', 1:N, zV(1,:), 'g-', 1:N, xV(1,:), 'k:', 1:N, xVu(1,:), 'm--')
legend('Reference', 'Reference + Process Noise', 'Measurement', 'EKF', 'UKF');
subplot(n,1,2)
plot(1:N, sR(2,:), 'r-', 1:N, sV(2,:), 'b-.', 1:N, xV(2,:), 'k:', 1:N, xVu(2,:), 'm--')
subplot(n,1,3)
plot(1:N, sR(3,:), 'r-', 1:N, sV(3,:), 'b-.', 1:N, zV(2,:), 'g-', 1:N, xV(3,:), 'k:', 1:N, xVu(3,:), 'm--')
subplot(n,1,4)
plot(1:N, sR(4,:), 'r-', 1:N, sV(4,:), 'b-.', 1:N, xV(4,:), 'k:', 1:N, xVu(4,:), 'm--')
subplot(n,1,5)
plot(1:N, sR(5,:), 'r-', 1:N, sV(5,:), 'b-.', 1:N, xV(5,:), 'k:', 1:N, xVu(5,:), 'm--')