%% Robot2DLocalizationGUI
% Demonstrates a 2D Position-Velocity Kalman Filter in a GUI with real-time playback.
% - Prediction step uses accelerometer data (kalman_predict).
% - Correction step uses position measurements (kalman_update).
% - Includes live controls for process/measurement noise and displays
%   true/estimated trajectories with an uncertainty ellipse.

function Robot2DLocalizationGUI()
    %% Create UI figure and layout
    f = figure('Name','2D Kalman Filter Fusion','NumberTitle','off','Position',[100, 100, 1000, 600]);
    % Axes for plot (left side)
    ax = axes(f, 'Units', 'pixels', 'Position', [50 50 600 500]);
    hold(ax, 'on');
    grid(ax, 'on');
    axis(ax, 'equal');
    xlabel(ax, '\xi_x'); 
    ylabel(ax, '\xi_y');

    % UI Controls (right side)
    uicontrol(f, 'Style','text','Position',[700 500 160 20],'String','Process Noise Std');
    procSlider = uicontrol(f, 'Style','slider','Position',[700 480 160 20],'Min',0.01,'Max',1,'Value',0.1);
    uicontrol(f, 'Style','text','Position',[700 440 160 20],'String','Measurement Noise Std');
    measSlider = uicontrol(f, 'Style','slider','Position',[700 420 160 20],'Min',0.1,'Max',2,'Value',0.5);

    uicontrol(f, 'Style','pushbutton','String','Run','Position',[700 370 160 30],'Callback',@runKalman);

    %% Nested function to run the Kalman Filter with real-time playback
    function runKalman(~, ~)
        cla(ax); % Clear any previous plots in axes
        dt = 0.5;
        T = 10;
        N = T / dt;

        % Model matrices
        A = [1 0 dt 0;
             0 1 0 dt;
             0 0 1 0;
             0 0 0 1];

        B = [0.5*dt^2 0;
             0 0.5*dt^2;
             dt 0;
             0 dt];

        H = [1 0 0 0;
             0 1 0 0];

        % Slider values
        process_noise_std = get(procSlider,'Value');
        measurement_noise_std = get(measSlider,'Value');
        Q = (process_noise_std^2) * eye(4);
        R = (measurement_noise_std^2) * eye(2);

        % Allocate memory
        x_true = zeros(4, N);
        x_est  = zeros(4, N);
        P = eye(4);
        accel_true = zeros(2, N);
        z_meas = zeros(2, N);

        % Initial states
        x_true(:,1) = [0; 0; 1; 0];
        x_est(:,1) = [4; 0; 1; 0]; % same or set your own guess
        P = P * 1.5;

        %% Generate True Motion and Measurements (all at once)
        for k = 2:N
            angle = 2*pi*(k*dt)/T;
            accel_true(:,k) = [-cos(angle); -sin(angle)];

            w = mvnrnd(zeros(4,1), Q)';
            x_true(:,k) = A*x_true(:,k-1) + B*accel_true(:,k) + w;

            v = mvnrnd(zeros(2,1), R)';
            z_meas(:,k) = H*x_true(:,k) + v;
        end

        %% Prepare plot handles for real-time animation
        truePlot = plot(ax, x_true(1,1), x_true(2,1), 'b-', 'LineWidth', 2);
        estPlot  = plot(ax, x_est(1,1),  x_est(2,1),  'r--','LineWidth',1.5);
        measPlot = scatter(ax, z_meas(1,1), z_meas(2,1), 10, 'k', 'filled', 'MarkerFaceAlpha', 0.2);
        plot_uncertainty_ellipse(ax, x_est(1:2,1), P(1:2,1:2), 'r');

        % Update the title
        title(ax, sprintf('Real-Time Kalman Filter\nProcess Noise=%.2f, Measurement Noise=%.2f',...
            process_noise_std, measurement_noise_std));

        %% Real-Time Kalman Filter loop
        for k = 2:N
            % Predict step
            [x_pred, P_pred] = kalman_predict(x_est(:,k-1), P, accel_true(:,k), A, B, Q);

            % Update step
            [x_est(:,k), P] = kalman_update(x_pred, P_pred, z_meas(:,k), H, R);

            % Update line data for the first k points
            set(truePlot, 'XData', x_true(1,1:k), 'YData', x_true(2,1:k));
            set(estPlot,  'XData', x_est(1,1:k),  'YData', x_est(2,1:k));
            set(measPlot, 'XData', z_meas(1,1:k),  'YData', z_meas(2,1:k));

            % Draw uncertainty ellipse at the current estimate
            plot_uncertainty_ellipse(ax, x_est(1:2,k), P(1:2,1:2), 'r');

            drawnow;
            pause(dt); % small pause to see the animation
        end
    end
end

%% Helper function: Plot Uncertainty Ellipse
function plot_uncertainty_ellipse(ax, mu, Sigma, color)
    hold(ax,'on');
    [V, D] = eig(Sigma);
    if any(diag(D) < 0)
        return; % Skip if covariance is not positive semidefinite
    end
    t = linspace(0, 2*pi, 100);
    a = (V * sqrt(D)) * [cos(t); sin(t)];
    plot(ax, mu(1) + a(1,:), mu(2) + a(2,:), color, 'LineWidth', 0.5);
    legend(ax,'True Position','Estimated Position','Measurements');
end
