close all;
clear;

%% Load measurements

data = load('volleyball_cameras.mat');

Y = data.Y_measured;   % measurements, n_y x N              [1, ..., N]
X_true = data.X_true;  % true state trajectory, n_x x (N+1) [1, ..., N+1]

%% Parameter
N = size(Y, 2);

h = 0.05;            % step size 

std_w = 0.01;        % standard deviation of state noise
std_v = 5*10^(-4);   % standard deviation of measurement noise

e_l = [1 0 0]';          % line of sight
e_u = [0 0 1]';          % up vector
e_s = cross(e_u, e_l);   % right eye to left eye

d_eyes = 1;    % dist between eyes

L = 0.02;      % dist to screen

g = 9.81;      % gravity

P_R = [2  0.5 1.8]';        % position of right eye
P_L = P_R + d_eyes.*e_s;    % position of left eye

n_y = 4;                % two 2-dimensional cameras

n_p = 3;                % 3-dimensional position 
n_v = 3;                % 3-dimensional velocity

n_x = n_p + n_v;        % state dimension

%% State model 
 
A = [eye(n_p) h.*eye(n_v);
     zeros(n_p)  eye(n_v)];

b = [0 0 -0.5*g*h^2 0 0 -h.*g]';

F = @(x) A*x + b;

%% Measurement model

alpha_R = @(p) L./((p-P_R)'*e_l); 
alpha_L = @(p) L./((p-P_L)'*e_l);

u_R = @(p)  e_u'*(alpha_R(p).*(p - P_R));
v_R = @(p) -e_s'*(alpha_R(p).*(p - P_R));

u_L = @(p)  e_u'*(alpha_L(p).*(p - P_L));
v_L = @(p) -e_s'*(alpha_L(p).*(p - P_L));

G = @(x) [v_R(x(1:3,1)) u_R(x(1:3,1)) v_L(x(1:3,1)) u_L(x(1:3,1))]';

%% Initialization

x_init_true = X_true(:, 1);        % true inital state

x_predict = x_init_true + 0.5;    % disturb the true inital state
P_predict = 1*eye(n_x);

V = std_v^2*eye(n_y);
W = std_w^2*eye(n_x);

X_est = zeros(n_x, N);

%% EKF steps

for i=1:N
            
    y = Y(:, i);
    
    g   = G(x_predict);
    J_G = Jacobian(G, x_predict);

    [x_current, P_current] = update(y, x_predict, P_predict, g, J_G, V);

    X_est(:, i) = x_current;

    [x_predict, P_predict] = predict(x_current, P_current, A, b, W);
end

%% Predicted position at timestep N+1 and its covariance estimate 
p_predict = x_predict(1:3);
sigma_p_predict = P_predict(1:3, 1:3);

%% Plot confidence ellipsoids

figure(4);
xy = [cos(linspace(0,2*pi,50)); sin(linspace(0,2*pi,50))];
dims = { [1, 2], [1, 3], [2, 3]};

xlabels = {"x", "x", "y"};
ylabels = {"y", "z", "z"};

plots = {};

for i=1:3
    
    plots{i} = subplot(1,3,i); hold on; grid on; axis equal;

    p_2d = p_predict(dims{i},1);
    [vec, val] = eig(sigma_p_predict(dims{i}, dims{i})); 

    xy_ellipse = vec*sqrt(val)*xy;

    plot(xy_ellipse(1,:), xy_ellipse(2,:));
    plot(0, 0, 'o', 'MarkerSize', 8);
    
    xlabel(xlabels{i}); ylabel(ylabels{i});
end

linkaxes([plots{1}, plots{2}, plots{3}], 'x')
linkaxes([plots{1}, plots{2}, plots{3}], 'y')

%% Compute 2D projection of X_opt

Y_est = zeros(n_y, N);

for n=1:N
   Y_est(:, n) = G(X_est(:, n));
end

%% Plot estimated and true 3D trajectories

figure(1); hold on; grid on;
plot3(X_true(1, :), X_true(2, :), X_true(3, :), 'kx-');

plot3(X_est(1, :), X_est(2, :), X_est(3, :), 'b-');

plot3(P_R(1), P_R(2), P_R(3), 'ro', 'MarkerSize', 5); 
plot3(P_L(1), P_L(2), P_L(3), 'bo', 'MarkerSize', 5);

quiver3(P_R(1), P_R(2), P_R(3), 5*e_l(1), 5*e_l(2), 5*e_l(3), 'r')
quiver3(P_L(1), P_L(2), P_L(3), 5*e_l(1), 5*e_l(2), 5*e_l(3), 'b')

view([-100, 7]);

%% Plot estimated 2D projections and measurements

figure(2); 
s1 = subplot(2,1,2); hold on; grid on; 
title('Right eye');
plot(Y(1, :), Y(2,:), 'ro');
plot(Y_est(1, :), Y_est(2, :), 'r.-');
legend('Y measured', 'Y estimated');

s2= subplot(2,1,1); hold on; grid on;
title('Left eye');
plot(Y(3, :), Y(4,:), 'bo');
plot(Y_est(3, :), Y_est(4, :), 'b.-');
legend('Y measured', 'Y estimated');

linkaxes([s1, s2], 'x')
linkaxes([s1, s2], 'y')

%% Computation of Jacobian via finite differences

function [J]= Jacobian(F, x)
    
    epsilon = sqrt(eps(x));  % epsilon = sqrt(precision)
   
    Fx = F(x);
    n_f = size(Fx, 1);
    n_x = size(x, 1);
    J = zeros(n_f, n_x);
    
    I = eye(n_x);
 
    for i=1:n_x
        x_perturbed = x + epsilon(i).*I(:,i);
        J(:, i) = (F(x_perturbed)-Fx)/epsilon(i);
    end
end

%% Predict and Update function for linear state and nonlinear output model
function [x_predict, P_predict] = predict(x_estimate, P_estimate, A, b, W)
  % Predict step for linear output model
  % x_estimate  estimated state   
  % P_estimate  (approximate) covariance of the estimated state
  % A, b        linear state model, x+ = Ax + b
  % W           covariance matrix of state noise
  
  x_predict = A*x_estimate + b;
  
  P_predict = A*P_estimate*A' + W;

end

function [x_estimate, P_estimate]  = update(y, x_predict, P_predict, g, J_G, V)
  % Update step for nonlinear output model
  % y           measurement
  % x_predict   predicted state   
  % P_predict   (approximate) covariance of predicted state
  % g           predicted output, g = G(x_predict)
  % J_G         Jacobian of G at x_predict
  % V           covariance matrix of measurement noise
  C = J_G;
  y_tilde = y - g + C * x_predict;
  
  P_estimate = inv(inv(P_predict) + C'*(V\C));
  x_estimate = x_predict  + P_estimate*C'*(V\(y_tilde-C*x_predict));
  
end

