close all;

% ## General parameters
% Fix random seed  so everyone gets the same results.
rng(230)
% Sampling time:
dt = 0.2;
% System noise:
R = diag([0.02^2, 0.02^2, 0.02^2]);
% Measurement noise (1D measurements):
Q = 0.002^2;
% Construct vector with control input:
uhc = repmat([0.1; 0.105], 1, 63);
ust = repmat([0.1; 0.1], 1, 30);
u = [uhc, ust, uhc, ust];

% ## Compute the true trajectory:
%    Evaluate kinematic model according to control input
%    and add some sampled system noise.
currx = [0; 0; 0];
x_true = currx;
z_meas = [];
for i = 1:size(u,2)
    % Evaluate kinematic Model:
    currx = motion_diff(currx, u(:,i));
    % Add system noise:
    currx = currx + mvnrnd([0; 0; 0], R)';
    % Simulate a measurement:
    z_meas = [z_meas, measurement(currx) + mvnrnd(zeros(size(Q,1),1),Q)'];
    % Store result for plotting
    x_true = [x_true, currx];
end

% ## Evaluate EKF without measurements: 
%    Estimates are based on control input alone.
curr_mean = [0; 0; 0];
curr_cov = diag([0, 0, 0]);

nomeas_mean = curr_mean;
nomeas_cov = reshape(curr_cov, 9, 1);
for i = 1:size(u,2)
    % Update mean and Jacobian G:
    [curr_mean G] = motion_diff(curr_mean, u(:,i));
    
    % Update covariance matrix
    curr_cov = G*curr_cov*(G') + R;
    
    % Store result for plotting
    nomeas_mean = [nomeas_mean, curr_mean];
    nomeas_cov = [nomeas_cov, reshape(curr_cov, 9, 1)];
end % for each u
plot_trajectory_cov(x_true, nomeas_mean, nomeas_cov, dt);

% ## Full EKF: Estimates based on control input and measurements.
curr_mean = [0; 0; 0];
curr_cov = diag([0. 0, 0]);

ekf_mean = curr_mean;
ekf_cov = reshape(curr_cov, 9, 1);
for i = 1:size(u,2)
    % Update step:
    [curr_mean G] = motion_diff(curr_mean, u(:,i));
    curr_cov = G*curr_cov*(G') + R;
    
    % Correction step:
    % Expected measurement:
    [z_exp, H] = measurement(curr_mean);
   
    % Kalman gain:
    K = curr_cov * (H') * inv(H*curr_cov*(H') + Q)
    curr_mean = curr_mean + K*(z_meas(:,i) - z_exp);
    curr_cov = (eye(3) - K*H)*curr_cov;
    
    % Store result for plotting
    ekf_mean = [ekf_mean, curr_mean];
    ekf_cov = [ekf_cov, reshape(curr_cov, 9, 1)];
end % for each u

plot_trajectory_cov(x_true, ekf_mean, ekf_cov, dt);