% === Problem 3a: Effect of Noise on Parameter Estimation ===

clear; close all;

% True system parameters
m = 0.75;
L = 1.25;
c = 0.15;
g = 9.81;

mL2_true = m * L^2;
mgL_true = m * g * L;
theta_true = [mL2_true; c; mgL_true];

% Load clean data
data = readtable('output/problem1_data.csv');
t = data.t;
q_clean = data.q;
u = data.u;
Ts = t(2) - t(1);
N = length(t);

% Derivatives from clean q
dq_clean = zeros(N,1);
ddq_clean = zeros(N,1);
for k = 2:N-1
    dq_clean(k) = (q_clean(k+1) - q_clean(k-1)) / (2*Ts);
    ddq_clean(k) = (q_clean(k+1) - 2*q_clean(k) + q_clean(k-1)) / Ts^2;
end

% LS estimation on clean data
X_clean = [ddq_clean, dq_clean, q_clean];
theta_hat_clean = (X_clean' * X_clean) \ (X_clean' * u);
rel_error_clean = abs((theta_hat_clean - theta_true) ./ theta_true) * 100;

% Reconstruct q̂_clean
q_hat_clean = zeros(N, 1);
dq_hat_clean = zeros(N, 1);
q_hat_clean(1) = q_clean(1);
dq_hat_clean(1) = dq_clean(1);
for k = 2:N
    ddq_hat_clean_k = (1/theta_hat_clean(1)) * ...
        (u(k-1) - theta_hat_clean(2)*dq_clean(k-1) - theta_hat_clean(3)*q_clean(k-1));
    dq_hat_clean(k) = dq_hat_clean(k-1) + Ts * ddq_hat_clean_k;
    q_hat_clean(k) = q_hat_clean(k-1) + Ts * dq_hat_clean(k-1);
end

% === Loop over noise levels ===
noise_levels = [0.001, 0.0025];

for i = 1:length(noise_levels)
    noise_std = noise_levels(i);
    q_noisy = q_clean + noise_std * randn(size(q_clean));

    % Derivatives from noisy q
    dq_noisy = zeros(N,1);
    ddq_noisy = zeros(N,1);
    for k = 2:N-1
        dq_noisy(k) = (q_noisy(k+1) - q_noisy(k-1)) / (2*Ts);
        ddq_noisy(k) = (q_noisy(k+1) - 2*q_noisy(k) + q_noisy(k-1)) / Ts^2;
    end

    % LS estimation on noisy data
    X_noisy = [ddq_noisy, dq_noisy, q_noisy];
    theta_hat_noisy = (X_noisy' * X_noisy) \ (X_noisy' * u);
    rel_error_noisy = abs((theta_hat_noisy - theta_true) ./ theta_true) * 100;

    % Reconstruct q̂_noisy
    q_hat_noisy = zeros(N,1);
    dq_hat_noisy = zeros(N,1);
    q_hat_noisy(1) = q_noisy(1);
    dq_hat_noisy(1) = dq_noisy(1);
    for k = 2:N
        ddq_hat_noisy_k = (1/theta_hat_noisy(1)) * ...
            (u(k-1) - theta_hat_noisy(2)*dq_noisy(k-1) - theta_hat_noisy(3)*q_noisy(k-1));
        dq_hat_noisy(k) = dq_hat_noisy(k-1) + Ts * ddq_hat_noisy_k;
        q_hat_noisy(k) = q_hat_noisy(k-1) + Ts * dq_hat_noisy(k-1);
    end

    % Print results
    fprintf('\n--- Noise std = %.4f ---\n', noise_std);
    fprintf('Clean:  mL2=%.4f (%.2f%%), c=%.4f (%.2f%%), mgL=%.4f (%.2f%%)\n', ...
        theta_hat_clean(1), rel_error_clean(1), ...
        theta_hat_clean(2), rel_error_clean(2), ...
        theta_hat_clean(3), rel_error_clean(3));
    fprintf('Noisy:  mL2=%.4f (%.2f%%), c=%.4f (%.2f%%), mgL=%.4f (%.2f%%)\n', ...
        theta_hat_noisy(1), rel_error_noisy(1), ...
        theta_hat_noisy(2), rel_error_noisy(2), ...
        theta_hat_noisy(3), rel_error_noisy(3));

    % === Combined plot ===
    figure('Name', sprintf('Noise std = %.4f', noise_std), 'Position', [100, 100, 1000, 800]);

    subplot(2,1,1);
    plot(t, q_clean, 'b', ...
         t, q_hat_clean, 'g--', ...
         t, q_hat_noisy, 'r:');
    legend('Actual q(t)', 'Estimated (clean)', 'Estimated (noisy)');
    title(sprintf('Estimated Output (σ = %.4f)', noise_std));
    ylabel('q(t) [rad]');
    grid on;

    subplot(2,1,2);
    bar([rel_error_clean, rel_error_noisy]);
    set(gca, 'XTickLabel', {'mL^2', 'c', 'mgL'});
    legend({'No Noise', sprintf('With Noise (σ=%.4f)', noise_std)}, 'Location', 'northwest');
    ylabel('Relative Error [%]');
    title('Estimation Error');
    grid on;

    % Save figure
    filename = sprintf('output/Prob3a_NoiseStd%.4f.png', noise_std);
    saveas(gcf, filename);
end