#!/usr/bin/env julia
#
# plot_reach_tubes.jl — multi-panel reach-tube visualization.
#
# Produces a four-panel figure from a reach-result .mat file:
#   (1) Temperature tube overlay: T_c, T_hot, T_cold envelopes together.
#       The gap between T_hot and T_cold is a direct proxy for core
#       thermal power (P = W·c_c·ΔT).
#   (2) ΔT_core = T_hot - T_cold envelope (the power proxy, explicit).
#   (3) Reactivity ρ envelope, normalized by β (in dollars).
#   (4) Normalized power n envelope.
#
# Two input formats supported:
#   operation: reach_operation_result.mat (linear reach, has R_lo/R_hi
#              matrices, time vector T, nominal X_nom).
#   heatup_pj: reach_heatup_pj_tight_full.mat (per-timestep envelopes
#              Tc_lo_ts/Tc_hi_ts/... already extracted; rho from PJ
#              controller).
#
# Usage:
#   julia --project=. scripts/plot_reach_tubes.jl operation
#   julia --project=. scripts/plot_reach_tubes.jl heatup_pj

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

using MAT
using Plots
gr()

include(joinpath(@__DIR__, "..", "..", "src", "pke_params.jl"))
const PLANT = pke_params()

function plot_tubes_operation()
    mat_path = joinpath(@__DIR__, "..", "..", "..", "results", "reach_operation_result.mat")
    d = matread(mat_path)

    T_vec = vec(d["T"])                        # time grid
    X_lo  = d["X_lo"]                          # 10 × M lower bounds
    X_hi  = d["X_hi"]                          # 10 × M upper bounds
    X_nom = d["X_nom"]                         # 10 × M nominal

    # States at indices: n=1, T_f=8, T_c=9, T_cold=10.
    n_lo = X_lo[1, :];   n_hi = X_hi[1, :]
    Tf_lo = X_lo[8, :];  Tf_hi = X_hi[8, :]
    Tc_lo = X_lo[9, :];  Tc_hi = X_hi[9, :]
    Tco_lo = X_lo[10, :]; Tco_hi = X_hi[10, :]

    # T_hot = 2*T_c - T_cold; envelope via worst-case signed combination.
    Th_lo = 2 .* Tc_lo .- Tco_hi
    Th_hi = 2 .* Tc_hi .- Tco_lo

    # ΔT_core = T_hot - T_cold = 2*(T_c - T_cold).
    dT_lo = 2 .* (Tc_lo .- Tco_hi)
    dT_hi = 2 .* (Tc_hi .- Tco_lo)

    # ρ under LQR: ρ_total = u + α_f·dT_f + α_c·dT_c where u = -K(x-x_op).
    # For a quick envelope, compute ρ at the nominal trajectory and
    # inflate by bounds of the feedback contributions from the tube.
    # Simpler: use total reactivity = u + α_f*(T_f-T_f0) + α_c*(T_c-T_c0).
    # u under LQR is small; we approximate ρ envelope by α feedback
    # alone (the u contribution is ≤ few cents).
    rho_nom = PLANT.alpha_f .* (X_nom[8, :] .- PLANT.T_f0) .+
              PLANT.alpha_c .* (X_nom[9, :] .- PLANT.T_c0)
    # Envelope via worst-case T_f, T_c.
    rho_lo = PLANT.alpha_f .* (Tf_hi .- PLANT.T_f0) .+
             PLANT.alpha_c .* (Tc_hi .- PLANT.T_c0)   # both α < 0, max T → min ρ
    rho_hi = PLANT.alpha_f .* (Tf_lo .- PLANT.T_f0) .+
             PLANT.alpha_c .* (Tc_lo .- PLANT.T_c0)

    title_stem = "Operation-mode LQR reach tubes"
    _plot_common(T_vec, Tc_lo, Tc_hi, Th_lo, Th_hi, Tco_lo, Tco_hi,
                 dT_lo, dT_hi, rho_lo, rho_hi, n_lo, n_hi, title_stem,
                 "reach_operation_tubes.png")
end

function plot_tubes_heatup_pj()
    mat_path = joinpath(@__DIR__, "..", "..", "..", "reachability",
                       "reach_heatup_pj_tight_full.mat")
    d = matread(mat_path)

    t_arr   = vec(d["t_arr"])
    Tc_lo   = vec(d["Tc_lo_ts"]);  Tc_hi   = vec(d["Tc_hi_ts"])
    Tf_lo   = vec(d["Tf_lo_ts"]);  Tf_hi   = vec(d["Tf_hi_ts"])
    Tco_lo  = vec(d["Tco_lo_ts"]); Tco_hi  = vec(d["Tco_hi_ts"])
    n_lo    = vec(d["n_lo_ts"]);   n_hi    = vec(d["n_hi_ts"])
    rho_lo  = vec(d["rho_lo_ts"]); rho_hi  = vec(d["rho_hi_ts"])

    Th_lo = 2 .* Tc_lo .- Tco_hi
    Th_hi = 2 .* Tc_hi .- Tco_lo
    dT_lo = 2 .* (Tc_lo .- Tco_hi)
    dT_hi = 2 .* (Tc_hi .- Tco_lo)

    title_stem = "Heatup PJ (tight entry) reach tubes"
    _plot_common(t_arr, Tc_lo, Tc_hi, Th_lo, Th_hi, Tco_lo, Tco_hi,
                 dT_lo, dT_hi, rho_lo, rho_hi, n_lo, n_hi, title_stem,
                 "reach_heatup_pj_tubes.png")
end

function _plot_common(t, Tc_lo, Tc_hi, Th_lo, Th_hi, Tco_lo, Tco_hi,
                     dT_lo, dT_hi, rho_lo, rho_hi, n_lo, n_hi,
                     title_stem, outname)
    CtoF(T) = T*9/5 + 32

    # Panel 1: T_c / T_hot / T_cold overlaid.
    p1 = plot(xlabel="Time [s]", ylabel="T [°C]",
              title="T tubes", legend=:right)
    plot!(p1, t, Tc_hi, fillrange=Tc_lo, fillalpha=0.30, color=:red,
          linealpha=0, label="T_c tube")
    plot!(p1, t, Th_hi, fillrange=Th_lo, fillalpha=0.25, color=:orange,
          linealpha=0, label="T_hot tube")
    plot!(p1, t, Tco_hi, fillrange=Tco_lo, fillalpha=0.25, color=:blue,
          linealpha=0, label="T_cold tube")

    # Panel 2: ΔT_core = T_hot - T_cold (power proxy at constant flow).
    P_lo_MW = PLANT.W * PLANT.c_c .* dT_lo ./ 1e6
    P_hi_MW = PLANT.W * PLANT.c_c .* dT_hi ./ 1e6
    p2 = plot(xlabel="Time [s]", ylabel="ΔT_core = T_hot - T_cold [K]",
              title="Core ΔT envelope (power proxy)", legend=:right)
    plot!(p2, t, dT_hi, fillrange=dT_lo, fillalpha=0.30, color=:purple,
          linealpha=0, label="ΔT_core [K]")
    p2b = twinx(p2)
    plot!(p2b, t, P_hi_MW, fillrange=P_lo_MW, fillalpha=0.0, color=:gray,
          linealpha=0.5, linestyle=:dot, label="P=W·c_c·ΔT [MW]",
          ylabel="Primary thermal power [MWth]")

    # Panel 3: ρ envelope in dollars.
    rho_lo_d = rho_lo ./ PLANT.beta
    rho_hi_d = rho_hi ./ PLANT.beta
    p3 = plot(xlabel="Time [s]", ylabel="ρ [\$]",
              title="Reactivity envelope (1 \$ = β = prompt-critical)",
              legend=:right)
    plot!(p3, t, rho_hi_d, fillrange=rho_lo_d, fillalpha=0.3,
          color=:darkgreen, linealpha=0, label="ρ / β")
    hline!(p3, [1.0, -1.0], ls=:dash, color=:red,
           label="prompt ±1 \$")
    hline!(p3, [0.0], ls=:dot, color=:black, label="critical")

    # Panel 4: n envelope (log scale if needed).
    p4 = plot(xlabel="Time [s]", ylabel="n (normalized power)",
              title="n envelope", legend=:right)
    plot!(p4, t, n_hi, fillrange=n_lo, fillalpha=0.3, color=:black,
          linealpha=0, label="n tube")

    fig = plot(p1, p2, p3, p4, layout=(2, 2), size=(1300, 800),
               plot_title=title_stem)
    figdir = joinpath(@__DIR__, "..", "..", "..", "docs", "figures")
    isdir(figdir) || mkpath(figdir)
    outpath = joinpath(figdir, outname)
    savefig(fig, outpath)
    println("Saved $outpath")
end

# CLI dispatch.
which_plot = length(ARGS) > 0 ? ARGS[1] : "both"
if which_plot in ("operation", "both")
    plot_tubes_operation()
end
if which_plot in ("heatup_pj", "both")
    mat_path = joinpath(@__DIR__, "..", "..", "..", "reachability",
                       "reach_heatup_pj_tight_full.mat")
    if isfile(mat_path)
        plot_tubes_heatup_pj()
    else
        println("Skipping heatup_pj plot — $mat_path not found.")
        println("(Run scripts/reach_heatup_pj_tight_full.jl first.)")
    end
end