%% test_linearize.m — sanity-check the Jacobians against the nonlinear sim
%
%  Simulate two parallel trajectories under a small Q_sg step (5% down):
%    1. Full nonlinear plant with u = 0
%    2. Linear dx/dt = A*dx + B_w*dw propagated by ode45
%  Compare deviations from x_op.  For a small-enough disturbance, the
%  trajectories should overlay to within a few percent over a few minutes.

clear; clc; close all;
addpath('controllers');

plant = pke_params();
x_op  = pke_initial_conditions(plant);

[A, B, B_w, xs, us, Qs] = pke_linearize(plant, x_op, 0, plant.P0);

%% ===== Eigenvalue/conditioning check =====
eigs_A = eig(A);
fprintf('\n=== Linearization at operating point ===\n');
fprintf('  ||A||      = %.3e\n', norm(A));
fprintf('  ||B||      = %.3e\n', norm(B));
fprintf('  ||B_w||    = %.3e\n', norm(B_w));
fprintf('  stable?    = %s (max Re(eig) = %.3e)\n', ...
        iif(all(real(eigs_A) < 1e-8), 'YES', 'NO'), max(real(eigs_A)));
fprintf('  fastest pole ~ %.3e /s\n', max(abs(real(eigs_A))));
fprintf('  slowest pole ~ %.3e /s\n', min(abs(real(eigs_A(real(eigs_A) < -1e-10)))));

%% ===== Linear vs nonlinear trajectory under small Q_sg step =====
% 5% down step at t = 30s.
dQ   = -0.05 * plant.P0;
Q_sg = @(t) plant.P0 + dQ * (t >= 30);

tspan = [0, 300];

% Nonlinear
opts = odeset('RelTol', 1e-8, 'AbsTol', 1e-10);
rhs_nl = @(t, x) pke_th_rhs(t, x, plant, Q_sg, 0);
[t_nl, X_nl] = ode15s(rhs_nl, tspan, x_op, opts);

% Linear (deviation)
dQ_of_t = @(t) dQ * (t >= 30);
rhs_lin = @(t, dx) A*dx + B_w*dQ_of_t(t);
[t_lin, DX_lin] = ode15s(rhs_lin, tspan, zeros(size(x_op)), opts);

% Reconstruct absolute linear state for plotting
X_lin = DX_lin + x_op.';

%% ===== Plot comparison — focus on thermal states =====
CtoF = @(T) T*9/5 + 32;
figdir = fullfile('..', 'docs', 'figures');
if ~exist(figdir, 'dir'), mkdir(figdir); end

figure('Position', [100 80 1100 700], 'Name', 'Linear vs nonlinear sanity');

subplot(2,2,1);
plot(t_nl,  X_nl(:,1),  'b-', 'LineWidth', 1.5); hold on;
plot(t_lin, X_lin(:,1), 'r--','LineWidth', 1.5);
xlabel('t [s]'); ylabel('n'); grid on;
title('Normalized power'); legend('nonlinear', 'linear', 'Location', 'best');

subplot(2,2,2);
plot(t_nl,  CtoF(X_nl(:,8)),  'b-', 'LineWidth', 1.5); hold on;
plot(t_lin, CtoF(X_lin(:,8)), 'r--','LineWidth', 1.5);
xlabel('t [s]'); ylabel('T_f [F]'); grid on;
title('Fuel temperature');

subplot(2,2,3);
plot(t_nl,  CtoF(X_nl(:,9)),  'b-', 'LineWidth', 1.5); hold on;
plot(t_lin, CtoF(X_lin(:,9)), 'r--','LineWidth', 1.5);
xlabel('t [s]'); ylabel('T_c [F]'); grid on;
title('Avg coolant');

subplot(2,2,4);
plot(t_nl,  CtoF(X_nl(:,10)), 'b-', 'LineWidth', 1.5); hold on;
plot(t_lin, CtoF(X_lin(:,10)),'r--','LineWidth', 1.5);
xlabel('t [s]'); ylabel('T_{cold} [F]'); grid on;
title('Cold-leg coolant');

sgtitle('Linear-model sanity check: 5% Q_{sg} down-step, u = 0');
exportgraphics(gcf, fullfile(figdir, 'linearize_sanity.png'), 'Resolution', 150);

%% ===== Quantitative error at a few times =====
fprintf('\n=== Max |linear - nonlinear| at sampled times ===\n');
fprintf('  (normalized by operating-point magnitude where nonzero)\n');
for ts = [60, 120, 300]
    xi_nl  = interp1(t_nl,  X_nl,  ts);
    xi_lin = interp1(t_lin, X_lin, ts);
    rel_err = abs(xi_nl - xi_lin) ./ max(abs(x_op.'), 1e-6);
    fprintf('  t = %3d s:  max rel err = %.2e  (%s)\n', ts, max(rel_err), ...
            state_name(find(rel_err == max(rel_err), 1)));
end

%% ===== Save linearization =====
save(fullfile('..', 'reachability', 'linearization_at_op.mat'), ...
     'A', 'B', 'B_w', 'xs', 'us', 'Qs', '-v7');
fprintf('\nSaved A/B/B_w to ../reachability/linearization_at_op.mat\n');

%% ===== helpers =====
function s = state_name(k)
    names = {'n','C1','C2','C3','C4','C5','C6','T_f','T_c','T_cold'};
    s = names{k};
end

function y = iif(cond, a, b)
    if cond, y = a; else, y = b; end
end