#!/usr/bin/env julia
#
# validate_pj.jl — quantify the prompt-jump approximation error.
#
# Two parallel sims of the same heatup scenario:
#   (i)  full 10-state PKE+T/H (pke_th_rhs!)
#   (ii) reduced 9-state prompt-jump model (pke_th_rhs_pj!)
#
# After the prompt transient (~few hundred microseconds) the two
# trajectories should agree closely on T_c, T_f, T_cold, and on the
# reconstructed n vs the dynamic n.  Quantify the error explicitly.

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

using Printf
using LinearAlgebra
using OrdinaryDiffEq
using Plots

include(joinpath(@__DIR__, "..", "src", "pke_params.jl"))
include(joinpath(@__DIR__, "..", "src", "pke_th_rhs.jl"))
include(joinpath(@__DIR__, "..", "src", "pke_th_rhs_pj.jl"))
include(joinpath(@__DIR__, "..", "controllers", "controllers.jl"))

plant = pke_params()
T_standby = plant.T_c0 - 33.333333

# Heatup scenario — same as sim_heatup.jl.
ref_heat = (; T_start=T_standby, T_target=plant.T_c0, ramp_rate=28/3600)
Q_sg = t -> 0.0

# Initial conditions — both at n=1e-3, same temperatures.
n0 = 1e-3
C0 = (plant.beta_i ./ (plant.lambda_i .* plant.Lambda)) .* n0
x0_full = [n0; C0; T_standby; T_standby; T_standby]
x0_pj   = [C0; T_standby; T_standby; T_standby]

tspan = (0.0, 3000.0)

# --- Full-state sim ---
function rhs_full!(dx, x, p, t)
    u = ctrl_heatup(t, x, plant, ref_heat)
    pke_th_rhs!(dx, x, t, plant, Q_sg, u)
end
prob_full = ODEProblem(rhs_full!, x0_full, tspan)
sol_full  = solve(prob_full, Rodas5(); reltol=1e-8, abstol=1e-10)

# --- Prompt-jump sim ---
# ctrl_heatup expects 10-state x with T_f at index 8, T_c at 9.
# We pass an adapter that maps 9-state to the controller's expected layout.
function ctrl_heatup_pj(t, x_pj, plant_arg, ref_arg)
    # Construct a fake 10-state with n placeholder; controller doesn't read n.
    x10 = [0.0; x_pj[1:6]; x_pj[7]; x_pj[8]; x_pj[9]]
    return ctrl_heatup(t, x10, plant_arg, ref_arg)
end
function rhs_pj!(dx, x, p, t)
    u = ctrl_heatup_pj(t, x, plant, ref_heat)
    pke_th_rhs_pj!(dx, x, t, plant, Q_sg, u)
end
prob_pj = ODEProblem(rhs_pj!, x0_pj, tspan)
sol_pj  = solve(prob_pj, Rodas5(); reltol=1e-8, abstol=1e-10)

# --- Compare at sampled times ---
function get_n_full(sol, t) sol(t)[1] end
function get_T_c_full(sol, t) sol(t)[9] end
function get_T_f_full(sol, t) sol(t)[8] end
function get_T_cold_full(sol, t) sol(t)[10] end

function get_T_c_pj(sol, t) sol(t)[8] end
function get_T_f_pj(sol, t) sol(t)[7] end
function get_T_cold_pj(sol, t) sol(t)[9] end
function get_n_pj(sol, t)
    x = sol(t)
    u = ctrl_heatup_pj(t, x, plant, ref_heat)
    pj_reconstruct_n(x, plant, u)
end

println("\n=== PJ vs full-state, heatup scenario ===")
println("  t [s]    n_full    n_pj      |Δn|/n_full   T_c err   T_f err   T_cold err")
for t_check in (1.0, 5.0, 10.0, 60.0, 300.0, 1200.0, 3000.0)
    n_f = get_n_full(sol_full, t_check)
    n_p = get_n_pj(sol_pj, t_check)
    Tc_err = abs(get_T_c_full(sol_full, t_check) - get_T_c_pj(sol_pj, t_check))
    Tf_err = abs(get_T_f_full(sol_full, t_check) - get_T_f_pj(sol_pj, t_check))
    Tcold_err = abs(get_T_cold_full(sol_full, t_check) - get_T_cold_pj(sol_pj, t_check))
    @printf "  %6.1f   %.3e %.3e   %8.2e     %.3e   %.3e   %.3e\n" t_check n_f n_p abs(n_f-n_p)/abs(n_f) Tc_err Tf_err Tcold_err
end

# --- Plot trajectory overlay ---
figdir = joinpath(@__DIR__, "..", "..", "docs", "figures")
isdir(figdir) || mkpath(figdir)

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

t_grid = range(0, 3000, length=600)
n_full_arr   = [get_n_full(sol_full, t)    for t in t_grid]
n_pj_arr     = [get_n_pj(sol_pj, t)        for t in t_grid]
Tc_full_arr  = [get_T_c_full(sol_full, t)  for t in t_grid]
Tc_pj_arr    = [get_T_c_pj(sol_pj, t)      for t in t_grid]

p_n = plot(t_grid ./ 60, n_full_arr,  lw=2, color=:blue, label="full-state",
           xlabel="t [min]", ylabel="n", title="Power: full vs PJ")
plot!(p_n, t_grid ./ 60, n_pj_arr, lw=1.5, ls=:dash, color=:red, label="prompt-jump")

p_Tc = plot(t_grid ./ 60, CtoF.(Tc_full_arr),  lw=2, color=:blue, label="full-state",
            xlabel="t [min]", ylabel="T_avg [F]", title="T_avg: full vs PJ")
plot!(p_Tc, t_grid ./ 60, CtoF.(Tc_pj_arr), lw=1.5, ls=:dash, color=:red, label="prompt-jump")

fig = plot(p_n, p_Tc, layout=(1,2), size=(1200, 450),
           plot_title="PJ vs full-state, heatup scenario (3000 s)")
savefig(fig, joinpath(figdir, "validate_pj_heatup.png"))
println("\nFigure: $figdir/validate_pj_heatup.png")