% Extended Kalman Filter with Animation and Range Visualization
clear; clc; close all;

% Simulation parameters
dt = 0.1;
T = 30;
N = T / dt;

% Landmark position
xL = 10; yL = 0;

% Control inputs
v = 1.0;
omega = 0.2;
u = [v; omega];

% Initial states
x_true = [0; 0; pi/4];
x_est = [0.5; -0.5; pi/3];
P = 0.1 * eye(3);

% Noise covariances
Qscale = 0.05;
Q = diag([0.01, 0.01, 0.005])*Qscale;
R = 0.5^2;

% Data logs
X_true = zeros(3, N);
X_est = zeros(3, N);
Z = zeros(1, N);

% Setup figure
figure('Color','white');
axis equal;
grid on;
xlim([-5 20]); ylim([-10 15]);
xlabel('x'); ylabel('y');
title('EKF Range-Only Localization');
hold on;
plot(xL, yL, 'rx', 'MarkerSize', 10, 'LineWidth', 2);
legend_entries = {'Landmark', 'True Position', 'Estimated Position'};

NIS = zeros(1, N);

% Animation loop
for k = 1:N
    % === True Motion ===
    theta = x_true(3);
    x_true(1) = x_true(1) + v*dt*cos(theta);
    x_true(2) = x_true(2) + v*dt*sin(theta);
    x_true(3) = wrapToPi(x_true(3) + omega*dt);

    % === Measurement ===
    dx = x_true(1) - xL;
    dy = x_true(2) - yL;
    r = sqrt(dx^2 + dy^2);
    z = r + sqrt(R)*randn();

    % === EKF Prediction ===
    theta_est = x_est(3);
    x_pred = [
        x_est(1) + v*dt*cos(theta_est);
        x_est(2) + v*dt*sin(theta_est);
        x_est(3) + omega*dt
    ];
    x_pred(3) = wrapToPi(x_pred(3));

    F = [
        1, 0, -v*dt*sin(theta_est);
        0, 1,  v*dt*cos(theta_est);
        0, 0, 1
    ];
    P_pred = F * P * F' + Q;

    % === EKF Update ===
    dx = x_pred(1) - xL;
    dy = x_pred(2) - yL;
    r_hat = sqrt(dx^2 + dy^2);
    H = [dx / r_hat, dy / r_hat, 0];
    S = H * P_pred * H' + R;
    K = P_pred * H' / S;
    
    beta = z - r_hat;
    NIS(k) = beta^2 / S;

    x_est = x_pred + K * (beta);
    x_est(3) = wrapToPi(x_est(3));
    P = (eye(3) - K * H) * P_pred;

    % Store data
    X_true(:,k) = x_true;
    X_est(:,k) = x_est;
    Z(k) = z;

    % === Plotting ===
    cla; % Clear axes
    plot(xL, yL, 'rx', 'MarkerSize', 10, 'LineWidth', 2); hold on;

    % Range measurement circle
    viscircles([xL, yL], z, 'Color', [0.7 0.7 1], 'LineStyle', '--', 'LineWidth', 1);

    % True and estimated trajectories
    plot(X_true(1,1:k), X_true(2,1:k), 'g-', 'LineWidth', 2); 
    plot(X_est(1,1:k), X_est(2,1:k), 'b--', 'LineWidth', 1.5);
    plot(x_true(1), x_true(2), 'go', 'MarkerSize', 6, 'MarkerFaceColor', 'g');
    plot(x_est(1), x_est(2), 'bo', 'MarkerSize', 6, 'MarkerFaceColor', 'b');

    % Uncertainty ellipse (2D position only)
    plot_uncertainty_ellipse(x_est(1:2), P(1:2,1:2), 0.9, 'b');

    legend(legend_entries, 'Location', 'northeastoutside');
    grid on
    xlim([-15 15]); ylim([-15 15]);
    pause(0.05);

end


function plot_uncertainty_ellipse(mu, Sigma, confidence, color)
    if nargin < 4
        color = 'b';
    end
    if nargin < 3
        confidence = 0.9;  % 90% confidence by default
    end

    % Chi-square quantile for desired confidence level (2D)
    s = sqrt(chi2inv(confidence, 2));

    % Compute ellipse
    [V, D] = eig(Sigma);
    t = linspace(0, 2*pi, 100);
    a = s * sqrt(D(1,1));
    b = s * sqrt(D(2,2));
    ellipse = V * [a*cos(t); b*sin(t)] + mu;

    plot(ellipse(1,:), ellipse(2,:), 'Color', color, 'LineWidth', 1.2);
end
