PWR-HYBRID-3/code/scripts/plot/plot_reach_tubes.jl

#!/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(probe_seconds::Int=8000)
    # Reads the per-probe save format from reach_heatup_pj.jl with the
    # latest tight.toml config.  Format: keys prefixed with "T_{probe}_"
    # for each requested probe horizon (e.g. T_8000_Tc_lo_ts).
    mat_path = joinpath(@__DIR__, "..", "..", "..", "results",
                       "reach_heatup_pj_tight.mat")
    d = matread(mat_path)

    pre = "T_$(probe_seconds)_"
    if !haskey(d, pre * "t_arr")
        # Fall back to legacy bare-key format (older mat files).
        pre = ""
        haskey(d, "t_arr") || error("Neither '$pre' nor bare keys found in $mat_path. Available probes: " *
            join([replace(k, r"_t_arr$" => "") for k in keys(d) if endswith(k, "_t_arr")], ", "))
    end

    t_arr   = vec(d[pre * "t_arr"])
    Tc_lo   = vec(d[pre * "Tc_lo_ts"]);  Tc_hi   = vec(d[pre * "Tc_hi_ts"])
    Tf_lo   = vec(d[pre * "Tf_lo_ts"]);  Tf_hi   = vec(d[pre * "Tf_hi_ts"])
    Tco_lo  = vec(d[pre * "Tco_lo_ts"]); Tco_hi  = vec(d[pre * "Tco_hi_ts"])
    n_lo    = vec(d[pre * "n_lo_ts"]);   n_hi    = vec(d[pre * "n_hi_ts"])
    rho_lo  = vec(d[pre * "rho_lo_ts"]); rho_hi  = vec(d[pre * "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, T=$probe_seconds s) 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

function plot_tubes_scram_full(probe_seconds::Int=60)
    # 3-panel scram tube plot using FULL-PKE reach (no prompt-jump
    # algebraic substitution).  T_c on top, T_f in the middle, ρ on the
    # bottom.  Per Dane's 2026-04-30 evening request: "I'd love tubes for
    # t_c t_f and rho."
    mat_path = joinpath(@__DIR__, "..", "..", "..", "results",
                       "reach_scram_full_fat.mat")
    d = matread(mat_path)
    pre = "T_$(probe_seconds)_"
    haskey(d, pre * "t_arr") || error("No per-step data for probe $probe_seconds in $mat_path")

    t_arr  = vec(d[pre * "t_arr"])
    Tc_lo  = vec(d[pre * "Tc_lo_ts"]);  Tc_hi  = vec(d[pre * "Tc_hi_ts"])
    Tf_lo  = vec(d[pre * "Tf_lo_ts"]);  Tf_hi  = vec(d[pre * "Tf_hi_ts"])
    rho_lo = vec(d[pre * "rho_lo_ts"]); rho_hi = vec(d[pre * "rho_hi_ts"])

    # Convert ρ to dollars.
    rho_lo_d = rho_lo ./ PLANT.beta
    rho_hi_d = rho_hi ./ PLANT.beta
    rho_sdm_dollars = 0.01 / PLANT.beta

    # Panel 1: T_c envelope
    p1 = plot(xlabel="Time after rod insertion [s]", ylabel="T_c [°C]",
              title="Average coolant temperature envelope",
              legend=:right, dpi=180)
    plot!(p1, t_arr, Tc_hi, fillrange=Tc_lo, fillalpha=0.35,
          color=:red, linealpha=0, label="T_c tube")

    # Panel 2: T_f envelope
    p2 = plot(xlabel="Time after rod insertion [s]", ylabel="T_f [°C]",
              title="Fuel temperature envelope",
              legend=:right, dpi=180)
    plot!(p2, t_arr, Tf_hi, fillrange=Tf_lo, fillalpha=0.35,
          color=:orange, linealpha=0, label="T_f tube")

    # Panel 3: ρ envelope in dollars
    p3 = plot(xlabel="Time after rod insertion [s]",
              ylabel="ρ [\$]   (1 \$ = β = prompt-critical)",
              title="Reactivity envelope",
              legend=:right, dpi=180)
    plot!(p3, t_arr, rho_hi_d, fillrange=rho_lo_d, fillalpha=0.35,
          color=:darkgreen, linealpha=0, label="ρ tube")
    hline!(p3, [0.0], ls=:dot, color=:black, label="critical")
    hline!(p3, [-rho_sdm_dollars], ls=:dash, color=:red, lw=2,
           label="ρ_SDM = -0.01 (\$ = $(round(-rho_sdm_dollars; digits=2)))")

    fig = plot(p1, p2, p3, layout=(3, 1), size=(900, 900),
               plot_title="Full-PKE scram reach tubes (T=$probe_seconds s, fat entry)")
    figdir = joinpath(@__DIR__, "..", "..", "..", "docs", "figures")
    isdir(figdir) || mkpath(figdir)
    outpath = joinpath(figdir, "reach_scram_full_tubes.png")
    savefig(fig, outpath)
    println("Saved $outpath")
end

function plot_tubes_scram_pj(probe_seconds::Int=60; variant::Symbol=:fat)
    # 2-panel scram tube plot: rho(t) on top, n(t) below.
    # variant=:tight reads reach_scram_pj_result.mat (small box around
    # the operating point).  variant=:fat reads reach_scram_pj_fat.mat
    # (union over shutdown + heatup-tight + operation + LOCA envelopes).
    mat_filename = variant == :fat ? "reach_scram_pj_fat.mat" :
                                      "reach_scram_pj_result.mat"
    mat_path = joinpath(@__DIR__, "..", "..", "..", "results", mat_filename)
    d = matread(mat_path)
    pre = "T_$(probe_seconds)_"
    haskey(d, pre * "t_arr") || error("No per-step data for probe $probe_seconds in $mat_path")

    t_arr  = vec(d[pre * "t_arr"])
    rho_lo = vec(d[pre * "rho_lo_ts"])
    rho_hi = vec(d[pre * "rho_hi_ts"])
    n_lo   = vec(d[pre * "n_lo_ts"])
    n_hi   = vec(d[pre * "n_hi_ts"])

    # Convert rho to dollars (rho / beta).
    rho_lo_d = rho_lo ./ PLANT.beta
    rho_hi_d = rho_hi ./ PLANT.beta

    # Panel 1: rho envelope in dollars.
    p1 = plot(xlabel="Time after rod insertion [s]",
              ylabel="ρ [\$]   (1 \$ = β = prompt-critical)",
              title="Reactivity envelope",
              legend=:right, dpi=180)
    plot!(p1, t_arr, rho_hi_d, fillrange=rho_lo_d, fillalpha=0.35,
          color=:darkgreen, linealpha=0, label="ρ tube")
    hline!(p1, [0.0], ls=:dot, color=:black, label="critical")
    # Shutdown margin threshold rho_SDM = 0.01 dk/k → in dollars: 0.01/beta
    rho_sdm_dollars = 0.01 / PLANT.beta
    hline!(p1, [-rho_sdm_dollars], ls=:dash, color=:red, lw=2,
           label="ρ_SDM = -0.01 (\$ = $(round(-rho_sdm_dollars; digits=2)))")

    # Panel 2: n envelope (log scale).
    # Filter NaN / non-positive (PJ reconstruction can produce slack-negatives).
    n_lo_pos = max.(n_lo, 1e-9)
    n_hi_pos = max.(n_hi, 1e-9)
    p2 = plot(xlabel="Time after rod insertion [s]",
              ylabel="n (normalized power)",
              title="Power envelope (log scale)",
              legend=:right, yaxis=:log, dpi=180)
    plot!(p2, t_arr, n_hi_pos, fillrange=n_lo_pos, fillalpha=0.35,
          color=:black, linealpha=0, label="n tube")
    hline!(p2, [1e-4], ls=:dash, color=:orange, lw=2, label="n = 10⁻⁴ (p_above_crit floor)")

    title_suffix = variant == :fat ? "fat entry, T=$probe_seconds s" :
                                       "T=$probe_seconds s"
    fig = plot(p1, p2, layout=(2, 1), size=(900, 700),
               plot_title="Scram reach tubes (PJ-reduced PKE, $title_suffix)")
    figdir = joinpath(@__DIR__, "..", "..", "..", "docs", "figures")
    isdir(figdir) || mkpath(figdir)
    out_filename = variant == :fat ? "reach_scram_pj_fat_tubes.png" :
                                      "reach_scram_pj_tubes.png"
    outpath = joinpath(figdir, out_filename)
    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 ("scram_full", "both")
    mat_path = joinpath(@__DIR__, "..", "..", "..", "results",
                       "reach_scram_full_fat.mat")
    probe_s = length(ARGS) >= 2 ? parse(Int, ARGS[2]) : 60
    if isfile(mat_path)
        plot_tubes_scram_full(probe_s)
    else
        println("Skipping scram_full plot — $mat_path not found.")
    end
end
if which_plot in ("scram_pj", "both")
    mat_path = joinpath(@__DIR__, "..", "..", "..", "results",
                       "reach_scram_pj_result.mat")
    probe_s = length(ARGS) >= 2 ? parse(Int, ARGS[2]) : 60
    if isfile(mat_path)
        plot_tubes_scram_pj(probe_s)
    else
        println("Skipping scram_pj plot — $mat_path not found.")
        println("(Run scripts/reach/reach_scram_pj.jl first.)")
    end
end
if which_plot in ("heatup_pj", "both")
    mat_path = joinpath(@__DIR__, "..", "..", "..", "results",
                       "reach_heatup_pj_tight.mat")
    # Optional second CLI arg: probe horizon in seconds (default 8000).
    probe = length(ARGS) >= 2 ? parse(Int, ARGS[2]) : 8000
    if isfile(mat_path)
        plot_tubes_heatup_pj(probe)
    else
        println("Skipping heatup_pj plot — $mat_path not found.")
        println("(Run scripts/reach/reach_heatup_pj.jl configs/heatup/tight.toml first.)")
    end
end