#!/usr/bin/env julia
#
# main_mode_sweep.jl — run every DRC-mode controller back-to-back.
#
# Julia equivalent of plant-model/main_mode_sweep.m.  Produces the
# same four-panel plots, saves to docs/figures/ with `mode_sweep_*.png`
# names (overwrites MATLAB outputs — the Julia versions take over).

using Pkg
Pkg.activate(joinpath(@__DIR__, ".."))

using Printf
using LinearAlgebra
using OrdinaryDiffEq
using Plots
using MatrixEquations
using JSON

include(joinpath(@__DIR__, "..", "src", "pke_params.jl"))
include(joinpath(@__DIR__, "..", "src", "pke_th_rhs.jl"))
include(joinpath(@__DIR__, "..", "src", "pke_linearize.jl"))
include(joinpath(@__DIR__, "..", "src", "pke_solver.jl"))
include(joinpath(@__DIR__, "..", "src", "plot_pke_results.jl"))
include(joinpath(@__DIR__, "..", "controllers", "controllers.jl"))

# --- Plant + predicates ---
plant = pke_params()

# Load T_standby from predicates.json for the hot-standby IC + heatup ref.
pred_path = joinpath(@__DIR__, "..", "..", "reachability", "predicates.json")
pred_raw  = JSON.parsefile(pred_path)
T_standby = plant.T_c0 + pred_raw["derived"]["T_standby_offset_C"]

# --- ICs ---
x_op = pke_initial_conditions(plant)

# Hot-standby IC: everything at T_standby, trace n.
n_shut = 1e-6
C_shut = (plant.beta_i ./ (plant.lambda_i .* plant.Lambda)) .* n_shut
x_shut = [n_shut; C_shut; T_standby; T_standby; T_standby]

# Heatup IC: already critical at 0.1% power, thermally matched.
n_heat = 1e-3
C_heat = (plant.beta_i ./ (plant.lambda_i .* plant.Lambda)) .* n_heat
x_heat = [n_heat; C_heat; T_standby; T_standby; T_standby]

# --- LQR gain (needed by ctrl_operation_lqr_factory) ---
A, B, B_w, _, _, _ = pke_linearize(plant)
Q_lqr = Diagonal([1.0, 1e-3, 1e-3, 1e-3, 1e-3, 1e-3, 1e-3, 1e-2, 1e2, 1.0])
R_lqr = 1e6 * ones(1, 1)
ctrl_op_lqr, K_lqr, _ = ctrl_operation_lqr_factory(
    plant, A, B; Q_lqr=Matrix(Q_lqr), R_lqr=R_lqr)

# --- Figure output directory ---
figdir = joinpath(@__DIR__, "..", "..", "docs", "figures")
isdir(figdir) || mkpath(figdir)

# --- Run each mode ---

println("\n===== Mode 1: ctrl_shutdown =====")
Q_shut = t -> 0.0
t1, X1, U1 = pke_solver(plant, Q_shut, ctrl_shutdown, nothing, (0.0, 600.0); x0=x_shut)

println("\n===== Mode 2: ctrl_heatup =====")
ref_heat = (; T_start=T_standby, T_target=plant.T_c0, ramp_rate=28/3600)
Q_heat = t -> 0.0
t2, X2, U2 = pke_solver(plant, Q_heat, ctrl_heatup, ref_heat, (0.0, 5400.0); x0=x_heat)

println("\n===== Mode 3a: ctrl_operation (P) =====")
Q_op   = t -> plant.P0 * (1.0 - 0.2 * (t >= 30))
ref_op = (; T_avg=plant.T_c0)
t3, X3, U3 = pke_solver(plant, Q_op, ctrl_operation, ref_op, (0.0, 600.0))

println("\n===== Mode 3b: ctrl_operation_lqr =====")
# LQR closure ignores ref; pass nothing.
t3b, X3b, U3b = pke_solver(plant, Q_op, ctrl_op_lqr, nothing, (0.0, 600.0))

println("\n===== Mode 4: ctrl_scram =====")
Q_scr = t -> plant.P0 * max(0.03, 1 - max(0, t - 10) / 10)
t4, X4, U4 = pke_solver(plant, Q_scr, ctrl_scram, nothing, (0.0, 600.0))

# --- Per-mode 4-panel plots ---
plot_pke_results(t1, X1, U1, plant, Q_shut;
    title="ctrl_shutdown (cold IC)",
    savepath=joinpath(figdir, "mode_sweep_1_shutdown.png"))
plot_pke_results(t2, X2, U2, plant, Q_heat;
    title="ctrl_heatup (ramp T_avg)",
    savepath=joinpath(figdir, "mode_sweep_2_heatup.png"))
plot_pke_results(t3, X3, U3, plant, Q_op;
    title="ctrl_operation (P on T_avg)",
    savepath=joinpath(figdir, "mode_sweep_3_operation.png"))
plot_pke_results(t3b, X3b, U3b, plant, Q_op;
    title="ctrl_operation_lqr",
    savepath=joinpath(figdir, "mode_sweep_3b_operation_lqr.png"))
plot_pke_results(t4, X4, U4, plant, Q_scr;
    title="ctrl_scram",
    savepath=joinpath(figdir, "mode_sweep_4_scram.png"))

# --- Heatup ramp-tracking figure ---
CtoF(T) = T*9/5 + 32
T_ref_trace = [min(T_standby + (28/3600)*ti, plant.T_c0) for ti in t2]
fig_heat = plot(t2 ./ 60, CtoF.(T_ref_trace), lw=1.3, ls=:dash, color=:black,
                label="T_ref (28 C/hr ramp)",
                xlabel="Time [min]", ylabel="T_avg [F]",
                title="ctrl_heatup: reference tracking", size=(1000, 400))
plot!(fig_heat, t2 ./ 60, CtoF.(X2[:, 9]), lw=2, color=:red,
      label="T_avg (plant)")
savefig(fig_heat, joinpath(figdir, "mode_sweep_heatup_tracking.png"))

# --- P vs LQR head-to-head ---
fig_ctrl = plot(t3, CtoF.(X3[:, 9]),  lw=2, color=:blue, label="ctrl_operation (P)",
                xlabel="Time [s]", ylabel="T_avg [F]",
                title="Operation mode: P vs LQR under 100% -> 80% Q_sg step",
                size=(1000, 400))
plot!(fig_ctrl, t3b, CtoF.(X3b[:, 9]), lw=2, color=:red, label="ctrl_operation_lqr")
hline!(fig_ctrl, [CtoF(plant.T_c0)], ls=:dash, color=:black, label="setpoint")
savefig(fig_ctrl, joinpath(figdir, "mode_sweep_op_P_vs_LQR.png"))

# --- Power overview ---
p_shut = plot(t1, max.(X1[:,1], 1e-20), yaxis=:log, title="Shutdown",
              xlabel="t [s]", ylabel="n (log)", legend=false, lw=1.5)
p_heat = plot(t2 ./ 60, X2[:,1], title="Heatup", xlabel="t [min]", ylabel="n",
              legend=false, lw=1.5)
p_op   = plot(t3, X3[:,1], title="Operation", xlabel="t [s]", ylabel="n",
              legend=false, lw=1.5)
p_scr  = plot(t4, max.(X4[:,1], 1e-20), yaxis=:log, title="Scram",
              xlabel="t [s]", ylabel="n (log)", legend=false, lw=1.5)
fig_ov = plot(p_shut, p_heat, p_op, p_scr, layout=(2,2), size=(1100, 650),
              plot_title="Normalized reactor power n(t) across DRC modes")
savefig(fig_ov, joinpath(figdir, "mode_sweep_power_overview.png"))

println("\n=== Figures written to $figdir ===")