%
% Problem 1a: Gradient estimation for mass-spring-damper system (Exercise 1a)
%
clear

% True system parameters
m_true = 1.315;
b_true = 0.225;
k_true = 0.725;

% Simulation parameters
Ts = 0.001;
T_total = 30;
t = 0:Ts:T_total;
N = length(t);

% Gamma setup
gamma = 0.33;
use_normalization = true;  % Set to false to disable normalization

fprintf('Using gamma = %.4f\n', gamma);
if use_normalization
    fprintf('Using normalized gradient update.\n');
else
    fprintf('Using unnormalized gradient update.\n');
end

% Define both input cases
input_cases = {...
    struct('name', 'constant', 'u', 2.5 * ones(1, N)), ...
    struct('name', 'sine', 'u', 2.5 * sin(t)) ...
};

fprintf('True m: %.4f, b: %.4f, k: %.4f\n', m_true, b_true, k_true);

for case_idx = 1:length(input_cases)
    input_case = input_cases{case_idx};
    u = input_case.u;

    % Preallocate state variables
    x = zeros(1, N);
    dx = zeros(1, N);
    ddx = zeros(1, N);

    % Initial conditions
    x(1) = 0;
    dx(1) = 0;

    % Simulate true system using RK4
    for k = 1:N-1
        f = @(x_, dx_, u_) (1/m_true) * (u_ - b_true * dx_ - k_true * x_);
        k1 = f(x(k), dx(k), u(k));
        k2 = f(x(k) + Ts/2 * dx(k), dx(k) + Ts/2 * k1, u(k));
        k3 = f(x(k) + Ts/2 * dx(k), dx(k) + Ts/2 * k2, u(k));
        k4 = f(x(k) + Ts * dx(k), dx(k) + Ts * k3, u(k));
        ddx(k) = k1;  % store first derivative (for later reuse)
        dx(k+1) = dx(k) + Ts/6 * (k1 + 2*k2 + 2*k3 + k4);
        x(k+1) = x(k) + Ts * dx(k);
    end

    % Approximate acceleration using finite differences
    ddx(1:end-1) = diff(dx) / Ts;
    ddx(end) = ddx(end-1);  % replicate last value

    % Gradient estimation setup
    theta_hat = zeros(3, N);            % rows: [m; b; k]
    theta_hat(:, 1) = [1; 1; 1];        % initial guesses

    % Run gradient estimator
    for k = 1:N-1
        u_vec = [ddx(k); dx(k); x(k)];
        y = u(k);
        y_hat = theta_hat(:, k)' * u_vec;
        e = y - y_hat;
        if use_normalization
            norm_factor = 1 + norm(u_vec)^2;
            theta_hat(:, k+1) = theta_hat(:, k) + Ts * gamma * (e / norm_factor) * u_vec;
        else
            theta_hat(:, k+1) = theta_hat(:, k) + Ts * gamma * e * u_vec;
        end
    end

    % Print final estimated values to console
    fprintf('\nFinal estimates for input: %s\n', input_case.name);
    fprintf('Estimated m: %.4f, b: %.4f, k: %.4f\n', theta_hat(1,end), theta_hat(2,end), theta_hat(3,end));

    % Plot estimated parameters
    figure('Name', ['Estimated Parameters - ' input_case.name], 'Position', [100, 100, 1280, 860]);
    normalization_text = ternary(use_normalization, 'Normalized', 'Unnormalized');
    sgtitle(sprintf('Input: %s   |   Gamma = %.3f   |   %s', input_case.name, gamma, normalization_text), 'FontWeight', 'bold');

    subplot(3,1,1);
    plot(t, theta_hat(1,:), 'b', 'LineWidth', 1.5);
    ylabel('$$\hat{m}(t)$$ [kg]', 'Interpreter', 'latex'); 
    grid on;
    title('Εκτίμηση μάζας');

    subplot(3,1,2);
    plot(t, theta_hat(2,:), 'r', 'LineWidth', 1.5);
    ylabel('$$\hat{b}(t)$$ [Ns/m]', 'Interpreter', 'latex'); 
    grid on;
    title('Εκτίμηση απόσβεσης');

    subplot(3,1,3);
    plot(t, theta_hat(3,:), 'k', 'LineWidth', 1.5);
    ylabel('$$\hat{k}(t)$$ [N/m]', 'Interpreter', 'latex'); 
    xlabel('t [sec]');
    grid on;
    title('Εκτίμηση ελαστικότητας');

    % Save figure
    if ~exist('output', 'dir')
        mkdir('output');
    end
    saveas(gcf, sprintf('output/Prob1a_estimation_%s_gamma%.3f_%s_%ds.png', input_case.name, gamma, normalization_text, T_total));
end

function out = ternary(cond, val_true, val_false)
    if cond
        out = val_true;
    else
        out = val_false;
    end
end