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_x) + [zeros(n_p) h.*eye(n_v);
                zeros(n_p)  zeros(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))]';


%% Residual function

R = @(X) residual(X, Y, F, G, N, n_x, n_y, std_w, std_v);

%% Initial guess: simulate system deterministically

x_init_true = [5 -7 1.7 5 5 10]' + 1;  % true inital state

X_init = zeros(n_x, N+1);
X_init(:, 1) = x_init_true;

for n=1:N
    X_init(:, n+1) = F(X_init(:, n));
end

X_init = reshape(X_init, (N+1)*n_x, 1);    % reshape to column vector

%% Solve estimation problem

[X_opt_col, ~, R_opt, ~, ~, ~, J] = lsqnonlin(R, X_init);        

X_opt = reshape(X_opt_col, n_x, N+1);      % n_x x (N+1) format 

J = full(J);   
%w_opt = R_opt(1:(N+1)*n_x, 1);
%v_opt = R_opt((N+1)*n_x+1:end, 1);

%% Check noise assumptions

figure(3);
histogram(R_opt, 31); 
title('Histogram of residual');

%% Covariance estimate
beta = 1/(N*n_x + N*n_y - (N+1)*n_x)*(R_opt'*R_opt);  % should be approximately 1
sigma_estimate = beta*(J'*J)^(-1);

%% Predicted position at timestep N+1 and its covariance estimate 

p_ls = X_opt(1:3, N+1);
sigma_p_ls = sigma_estimate((N*n_x+1):(N*n_x+n_p), (N*n_x+1):(N*n_x+n_p));

%% 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_ls(dims{i},1);
    [vec, val] = eig(sigma_p_ls(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_opt = zeros(n_y, N);

for n=1:N
   Y_opt(:, n) = G(X_opt(:, 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_opt(1, :), X_opt(2, :), X_opt(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_opt(1, :), Y_opt(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_opt(3, :), Y_opt(4, :), 'b.-');
legend('Y measured', 'Y estimated');

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


%% Residual
function [R] = residual(X, Y, F, G, N, n_x, n_y, std_w, std_v)

    R = zeros(n_x*N + n_y*N, 1);

    x_next = X(1:n_x, 1);
    
    W = [(1./std_w).*eye(N*n_x)  zeros(N*n_x, N*n_y);
         zeros(N*n_y, N*n_x)    (1./std_v).*eye(N*n_y)];

    for n=1:N
        x_current = x_next;
        x_next = X((n*n_x+1):(n+1)*n_x, 1);

        i1 = (n-1)*n_x + 1;
        i2 = N*n_x + (n-1)*n_y + 1;

        R(i1:(i1+n_x-1), 1) = x_next - F(x_current); 
        R(i2:(i2+n_y-1), 1) = Y(:, n) - G(x_current);
    end
 
    R = W*R;
end





