%% main_mode_sweep.m — demo each DRC continuous-mode controller
%
%  One scenario per mode, from an initial condition plausible for that
%  mode.  Figures saved to ../docs/figures/mode_sweep_*.png so the runs
%  are reviewable without the MATLAB IDE.
%
%  Modes, in the order DRC visits them starting from shutdown:
%    1. shutdown   — deep subcritical, cold IC, Q_sg = 0
%    2. heatup     — same cold IC, ramped T_avg reference
%    3. operation  — operating IC, 100% -> 80% Q_sg step
%    4. scram      — operating IC, rods slammed in, Q_sg decays to 0
%
%  Each run produces plot_pke_results 4-panel plus saves a lightweight
%  summary figure.

clear; clc; close all;
addpath('controllers');
addpath(fullfile('..', 'reachability'));   % for load_predicates

plant = pke_params();
pred  = load_predicates(plant);            % T_standby and FRET-predicate concretizations
figdir = fullfile('..', 'docs', 'figures');
if ~exist(figdir, 'dir'), mkdir(figdir); end

CtoF = @(T) T*9/5 + 32;

%% ===== IC helpers =====
% Operating IC: full-power steady state.
x0_op = pke_initial_conditions(plant);

% Hot-standby shutdown IC: coolant flat at T_standby = T_c0 - 60 F ~ 275 C,
% reactor deep subcritical with only a trace neutron population.  T_standby
% comes from reachability/predicates.json (single source of truth).  Well
% inside the model's +/-50 C trust region around the operating point, and
% above coolant saturation at reduced operating pressure.
T_standby = pred.constants.T_standby;
n_shut  = 1e-6;
C_shut  = (plant.beta_i ./ (plant.lambda_i * plant.Lambda)) * n_shut;
x0_shut = [n_shut; C_shut; T_standby; T_standby; T_standby];

% Heatup IC: reactor already taken critical at 0.1% power at hot-standby
% temperature.  Mirrors real plant practice: achieve criticality, then
% heat up.  Same T_standby as the shutdown IC — heatup begins from where
% shutdown left off.
n_heat  = 1e-3;
C_heat  = (plant.beta_i ./ (plant.lambda_i * plant.Lambda)) * n_heat;
x0_heat = [n_heat; C_heat; T_standby; T_standby; T_standby];

%% ===== Mode 1: SHUTDOWN =====
fprintf('\n===== Mode 1: ctrl_shutdown =====\n');
Q_sg_shut = @(t) 0;                       % no SG demand
tspan_shut = [0, 600];
[t1, X1, U1] = pke_solver(plant, Q_sg_shut, @ctrl_shutdown, [], tspan_shut, x0_shut);

%% ===== Mode 2: HEATUP =====
fprintf('\n===== Mode 2: ctrl_heatup =====\n');
ref_heatup = struct();
ref_heatup.T_start    = T_standby;         % ~275 C (hot standby, = T_c0 - 60 F)
ref_heatup.T_target   = plant.T_c0;        % 308.35 C (operating setpoint)
ref_heatup.ramp_rate  = 28 / 3600;         % 28 C/hr, tech-spec limit
Q_sg_heat = @(t) 0;                        % no SG demand during heatup
tspan_heat = [0, 5400];                    % ~90 min (33 C span at 28 C/hr = 71 min + settling)
[t2, X2, U2] = pke_solver(plant, Q_sg_heat, @ctrl_heatup, ref_heatup, tspan_heat, x0_heat);

%% ===== Mode 3a: OPERATION (plain P) =====
fprintf('\n===== Mode 3a: ctrl_operation (P) =====\n');
Q_sg_op = @(t) plant.P0 * (1.0 - 0.2 * (t >= 30));   % 100% -> 80% step
ref_op  = struct('T_avg', plant.T_c0);
tspan_op = [0, 600];
[t3, X3, U3] = pke_solver(plant, Q_sg_op, @ctrl_operation, ref_op, tspan_op, x0_op);

%% ===== Mode 3b: OPERATION (LQR) =====
fprintf('\n===== Mode 3b: ctrl_operation_lqr =====\n');
[t3b, X3b, U3b] = pke_solver(plant, Q_sg_op, @ctrl_operation_lqr, [], tspan_op, x0_op);

%% ===== Mode 4: SCRAM =====
fprintf('\n===== Mode 4: ctrl_scram =====\n');
% After a scram signal, the turbine trips and SG isolation occurs fast.
% Q_sg drops from P0 to a decay-heat-level sink (3% of P0) in ~10 s,
% then holds.  Not a true decay-heat model — good enough to show the
% post-scram thermal trajectory without coolant going unphysically cold.
Q_sg_scr = @(t) plant.P0 * max(0.03, 1 - max(0, t - 10) / 10);
tspan_scr = [0, 600];
[t4, X4, U4] = pke_solver(plant, Q_sg_scr, @ctrl_scram, [], tspan_scr, x0_op);

%% ===== Per-mode 4-panel plots =====
plot_pke_results(t1, X1, U1, plant, Q_sg_shut, 'ctrl\_shutdown (cold IC)');
exportgraphics(gcf, fullfile(figdir, 'mode_sweep_1_shutdown.png'), 'Resolution', 150);

plot_pke_results(t2, X2, U2, plant, Q_sg_heat, 'ctrl\_heatup (ramp T\_avg)');
exportgraphics(gcf, fullfile(figdir, 'mode_sweep_2_heatup.png'), 'Resolution', 150);

plot_pke_results(t3, X3, U3, plant, Q_sg_op, 'ctrl\_operation (P on T\_avg)');
exportgraphics(gcf, fullfile(figdir, 'mode_sweep_3_operation.png'), 'Resolution', 150);

plot_pke_results(t3b, X3b, U3b, plant, Q_sg_op, 'ctrl\_operation\_lqr');
exportgraphics(gcf, fullfile(figdir, 'mode_sweep_3b_operation_lqr.png'), 'Resolution', 150);

plot_pke_results(t4, X4, U4, plant, Q_sg_scr, 'ctrl\_scram');
exportgraphics(gcf, fullfile(figdir, 'mode_sweep_4_scram.png'), 'Resolution', 150);

%% ===== Heatup ramp-tracking figure =====
% Overlay the T_ref signal on T_avg so the lag is visible.
T_ref_trace = min(ref_heatup.T_start + ref_heatup.ramp_rate .* t2, ref_heatup.T_target);
figure('Position', [100 80 900 400], 'Name', 'Heatup ramp tracking');
plot(t2/60, CtoF(T_ref_trace), 'k--', 'LineWidth', 1.2); hold on;
plot(t2/60, CtoF(X2(:,9)),    'r-',  'LineWidth', 1.5);
xlabel('Time [min]'); ylabel('T_{avg} [F]');
legend('T_{ref}(t) (28 C/hr ramp)', 'T_{avg} (plant)', 'Location', 'southeast');
title('ctrl\_heatup: reference tracking');
grid on;
exportgraphics(gcf, fullfile(figdir, 'mode_sweep_heatup_tracking.png'), 'Resolution', 150);

%% ===== P vs LQR head-to-head on T_avg =====
figure('Position', [100 80 900 400], 'Name', 'Operation: P vs LQR');
plot(t3,  CtoF(X3(:,9)),  'b-',  'LineWidth', 1.5); hold on;
plot(t3b, CtoF(X3b(:,9)), 'r-',  'LineWidth', 1.5);
yline(CtoF(plant.T_c0), 'k--');
xlabel('Time [s]'); ylabel('T_{avg} [F]');
legend('ctrl\_operation (P)', 'ctrl\_operation\_lqr', 'setpoint', 'Location', 'best');
title('Operation mode: P vs LQR under 100% \rightarrow 80% Q_{sg} step');
grid on;
exportgraphics(gcf, fullfile(figdir, 'mode_sweep_op_P_vs_LQR.png'), 'Resolution', 150);

%% ===== Summary figure: normalized power across all modes =====
figure('Position', [100 80 1000 450], 'Name', 'Normalized power, all modes');
subplot(2,2,1); semilogy(t1, max(X1(:,1), 1e-20)); grid on;
    title('Shutdown');  xlabel('t [s]'); ylabel('n (log)');
subplot(2,2,2); plot(t2/60, X2(:,1)); grid on;
    title('Heatup');    xlabel('t [min]'); ylabel('n');
subplot(2,2,3); plot(t3, X3(:,1)); grid on;
    title('Operation'); xlabel('t [s]'); ylabel('n');
subplot(2,2,4); semilogy(t4, max(X4(:,1), 1e-20)); grid on;
    title('Scram');     xlabel('t [s]'); ylabel('n (log)');
sgtitle('Normalized reactor power n(t) across DRC modes');
exportgraphics(gcf, fullfile(figdir, 'mode_sweep_power_overview.png'), 'Resolution', 150);

fprintf('\n=== Figures written to %s ===\n', figdir);