clear all; close all;

% System parameters
tau_values = [0.2, 1];  % Time delay values
dt = 0.01;              % Time step
t_max = 100;            % Maximum simulation time
num_steps = round(t_max / dt) + 1;  % Number of steps including the initial point

% Time vector
t = 0:dt:t_max;

% Preallocate arrays to store results for both tau values
x_all = zeros(3, num_steps, length(tau_values));
xh_all = zeros(3, num_steps, length(tau_values));
u_all = zeros(2, num_steps, length(tau_values));
fa_all = zeros(2, num_steps, length(tau_values));
fa_hat_all = zeros(2, num_steps, length(tau_values));
Wc_hist_all = zeros(6, num_steps, length(tau_values));
integral_all = zeros(num_steps, length(tau_values));

% Loop over different tau values
for tau_idx = 1:length(tau_values)
    tau = tau_values(tau_idx);
    
    % Initial conditions
    x = zeros(3, num_steps);
    X0=[1; 0; 0]
    x(:, 1) = X0; % Initial states
    xh = zeros(3, num_steps);
    xh(:, 1) = [0; 0.3; 0.4]; % Initial observer states
    u = zeros(2, num_steps); 
    fa_hat = zeros(2, num_steps); % Initial fault estimation
    fa = zeros(2, num_steps); % Fault
    
    % Observer and learning parameters
    L1 = 2 * eye(3); % Observer gain
    Lr = 0.01; % Learning rate
    L2 = 8; % Gain in fault estimation
    ro = 3; % For fault in HJB
    Q = eye(3);
    R = eye(2);
    Qd = 10 * eye(3);
    % Initial weights for the critic structure
    Wc = [3; 0.2; 5.4; 3.4; 4.7; 1.3];
    Wc_hist = zeros(length(Wc), num_steps);
    Wc_hist(:, 1) = Wc;
    
    % Initialize integral array
    integral_values = zeros(1, num_steps);
    
    % Simulation loop
    for k = 2:num_steps
        % Define the thresholds based on num_steps
        quarter_steps = num_steps / 4;
        three_quarter_steps = 3 * num_steps / 4;

        % Occurring fault logic
        if k <= 30/dt
            fa(:, k) = [0; 0];
        elseif k > 30/dt && k <= 70/dt
            fa(:, k) = [1; 0.5];
        else
            fa(:, k) = [1; 0.5];
        end

        if t(k) <= tau
            x_tau = X0;
            xh_tau = [0; 0.3; 0.4];
        else
            delay_idx = round((t(k) - tau) / dt) + 1;
            x_tau = x(:, delay_idx);
            xh_tau = xh(:, delay_idx);
        end

        % System dynamics
        x_dot = [-x(1, k-1) + sin(x_tau(2)) + x(3, k-1);
                 -x(2, k-1) + cos(x_tau(1)) + u(1, k-1);
                 -x(3, k-1) + x_tau(1) + x(2, k-1) + u(2, k-1)];

        ghx = [0  0; 1  0; 0  1];
        x(:, k) = x(:, k-1) + dt * (x_dot - ghx * fa(:, k-1));
        
        % Observer dynamics
        xh_dot = [-xh(1, k-1) + sin(xh_tau(2)) + xh(3, k-1);
                  -xh(2, k-1) + cos(xh_tau(1));
                  -xh(3, k-1) + xh_tau(1) + xh(2, k-1)];
        
        xh(:, k) = xh(:, k-1) + dt * (xh_dot + ghx * (u(:, k-1) - fa_hat(:, k-1)) + L1 * (x(:, k-1) - xh(:, k-1)));

        % Fault estimation
        ex(:, k) = x(:, k) - xh(:, k);
        fa_hat(:, k) = fa_hat(:, k-1) + dt * (-L2 * ghx' * ex(:, k));

        % Calculate the integral within the window [t-tau, t]
        if k * dt <= tau
            integral_window = xh(:, 1:k);
        else
            start_idx = max(1, k - round(tau / dt));
            integral_window = xh(:, start_idx:k);
        end
        integral_value = dt * sum(sum(integral_window .* (Qd * integral_window), 1));
        integral_values(k) = integral_value;

        % Critic update
        x_critic = [xh(1, k); xh(2, k); xh(3, k)];
        phi = polynomial_features(x_critic);
        grad_phi = polynomial_features_gradient(x_critic);
        dot_xh = (xh(:, k) - xh(:, k-1)) / dt;

        ec(k) = ro * fa_hat(:, k)' * fa_hat(:, k) + x(:, k)' * Q * x(:, k) + u(:, k)' * R * u(:, k) + (Wc' * grad_phi * dot_xh) +integral_value;
        Wc = Wc + dt * (-Lr * ec(k) * grad_phi * dot_xh);
        Wc_hist(:, k) = Wc;

        % Control input update
        u(:, k) = -0.5 * (R^-1) * (ghx' * grad_phi' * Wc)+fa_hat(:,k);
         u1(:, k)=-0.5 * (R^-1) * (ghx' * grad_phi' * Wc);
    end
    
    % Store results for this tau value
    x_all(:,:,tau_idx) = x;
    xh_all(:,:,tau_idx) = xh;
    u_all(:,:,tau_idx) = u1;
    fa_all(:,:,tau_idx) = fa;
    fa_hat_all(:,:,tau_idx) = fa_hat;
    Wc_hist_all(:,:,tau_idx) = Wc_hist;
    integral_all(:, tau_idx) = integral_values;
end

% Assuming t, x_all, and xh_all are already defined in the workspace

% Plotting results
figure;
for i = 1:3
    subplot(3, 1, i);
    
    % Plot actual state with disturbance d = 0.2
    plot(t, squeeze(x_all(i, :, 1)), 'b-', 'LineWidth', 2); hold on;
    
    % Plot estimated state with disturbance d = 0.2
    plot(t, squeeze(xh_all(i, :, 1)), 'b--', 'LineWidth', 2);
    
    % Plot actual state with disturbance d = 1
    plot(t, squeeze(x_all(i, :, 2)), 'r-', 'LineWidth', 2);
    
    % Plot estimated state with disturbance d = 1
    plot(t, squeeze(xh_all(i, :, 2)), 'r--', 'LineWidth', 2);
    
    % Set labels and title
    xlabel('Time t');
    ylabel(['State x_' num2str(i)]);
    title(['Response of {x}_' num2str(i) ' and \hat{x}_' num2str(i)]);
    
    % Set legend with LaTeX interpreter
    legend(['$x_' num2str(i) '$ (d=0.2)'], ['$\hat{x}_' num2str(i) '$ (d=0.2)'], ...
           ['$x_' num2str(i) '$ (d=1)'], ['$\hat{x}_' num2str(i) '$ (d=1)'], ...
           'Interpreter', 'latex');
    
    % Display grid
    grid on;
end



figure;
for i = 1:6
    subplot(3, 2, i);
    plot(t, Wc_hist_all(i, :, 1), 'b-', 'LineWidth', 2); hold on;
    plot(t, Wc_hist_all(i, :, 2), 'r-', 'LineWidth', 2);
    xlabel('Time t');
    ylabel(['Weight Wc_' num2str(i)]);
    title(['Evolution of Critic Weight Wc_' num2str(i)]);
    legend('d=0.2', 'd=1');
    grid on;
end

% New figure for control inputs
figure;
for i = 1:2
    subplot(2, 1, i);
    plot(t, u_all(i, :, 1), 'b-', 'LineWidth', 2); hold on;
    plot(t, u_all(i, :, 2), 'r-', 'LineWidth', 2);
    xlabel('Time t');
    ylabel(['Control Input $u_' num2str(i) '$'], 'Interpreter', 'latex');
    title(['Control Input $u_' num2str(i) '$ over Time'], 'Interpreter', 'latex');
    legend('d=0.2', 'd=1');
    grid on;
end

% New figure for integral values
% New figure for integral values
figure;
plot(t, integral_all(:, 1), 'b-', 'LineWidth', 2); hold on;
plot(t, integral_all(:, 2), 'r-', 'LineWidth', 2);
xlabel('Time t');
ylabel('Integral Value');
title('Integral of $z^T Q z$ over $[t-d, t]$', 'Interpreter', 'latex');
legend('d=0.2', 'd=1');
grid on;



% New figure for fault fa and its estimation fa_hat
figure;
for i = 1:2
    subplot(2, 1, i);
    plot(t, fa_all(i, :, 1), 'b-', 'LineWidth', 2); hold on;
    plot(t, fa_hat_all(i, :, 1), 'b--', 'LineWidth', 2);
    plot(t, fa_all(i, :, 2), 'r-', 'LineWidth', 2);
    plot(t, fa_hat_all(i, :, 2), 'r--', 'LineWidth', 2);
    xlabel('Time t');
    ylabel(['Fault $fa_' num2str(i) '$'], 'Interpreter', 'latex');
    title(['Fault $fa_' num2str(i) '$ and Estimated Fault $\hat{fa}_' num2str(i) '$'], 'Interpreter', 'latex');
    legend(['$fa_' num2str(i) ' \, (d=0.2)$'], ['$\hat{fa}_' num2str(i) ' \, (d=0.2)$'], ...
           ['$fa_' num2str(i) ' \, (d=1)$'], ['$\hat{fa}_' num2str(i) ' \, (d=1)$'], ...
           'Interpreter', 'latex');
    grid on;
end



t = 0:dt:t_max;
% Helper functions
function phi = polynomial_features(x)
    N = x(1); T = x(2); I = x(3);
    phi = [N^2; T^2; I^2; N*T; N*I; T*I];
end

function grad_phi = polynomial_features_gradient(x)
    N = x(1); T = x(2); I = x(3);
    grad_phi = [
        2*N,  0,  0;
        0,  2*T,  0;
        0,  0,  2*I;
        T,  N,  0;
        I,  0,  N;
        0,  I,  T
    ];
end
