#!/usr/bin/env julia
#
# reach_scram_full_fat.jl — FULL nonlinear PKE scram reach (no PJ).
#
# Companion to reach_scram_pj_fat.jl which uses the prompt-jump-reduced
# 9-state model.  This version keeps n as a state and uses the full
# 10-state PKE — captures the prompt jump dynamics directly (rather
# than via the algebraic n substitution).
#
# Dane wanted the "real" prompt jump visible in the tubes, even at the
# cost of stiffness-driven slowness in TMJets.  The penalty: each
# integration step is bounded by Λ ≈ 10⁻⁴ s.  60 s plant horizon means
# very many steps and a real risk of over-approximation slack growth.
#
# State (11D with augmented time):
#   x[1]    = n
#   x[2..7] = C_1..C_6
#   x[8]    = T_f
#   x[9]    = T_c
#   x[10]   = T_cold
#   x[11]   = t (augmented time)
#
# Constant control u = -8β (rod-fall scram).
# Q_sg = 3% P0 (decay-heat-level sink).
#
# Saves per-step envelopes for T_c, T_f, ρ, n at every probe horizon.

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

using LinearAlgebra
using ReachabilityAnalysis, LazySets
using JSON
using MAT
using Printf

# Plant constants — must match pke_params.
const LAMBDA  = 1e-4
const BETA_1, BETA_2, BETA_3, BETA_4, BETA_5, BETA_6 =
    0.000215, 0.001424, 0.001274, 0.002568, 0.000748, 0.000273
const BETA    = BETA_1 + BETA_2 + BETA_3 + BETA_4 + BETA_5 + BETA_6
const LAM_1, LAM_2, LAM_3, LAM_4, LAM_5, LAM_6 =
    0.0124, 0.0305, 0.111, 0.301, 1.14, 3.01

const P0   = 1e9
const M_F, C_F, M_C, C_C, HA, W_M, M_SG =
    50000.0, 300.0, 20000.0, 5450.0, 5e7, 5000.0, 30000.0
const ALPHA_F, ALPHA_C = -2.5e-5, -1e-4

const T_COLD0 = 290.0
const DT_CORE = P0 / (W_M * C_C)
const T_HOT0  = T_COLD0 + DT_CORE
const T_C0    = (T_HOT0 + T_COLD0) / 2
const T_F0    = T_C0 + P0 / HA

const U_SCRAM    = -8 * BETA
const Q_SG_DECAY = 0.03 * P0

const RHO_SDM = 0.01
const SDM_RHS = -RHO_SDM - U_SCRAM + ALPHA_F*T_F0 + ALPHA_C*T_C0

# Full-PKE Taylorized scram RHS (11 states with augmented t).
@taylorize function rhs_scram_full_taylor!(dx, x, p, t)
    rho = U_SCRAM + ALPHA_F * (x[8] - T_F0) + ALPHA_C * (x[9] - T_C0)
    sum_lam_C = LAM_1*x[2] + LAM_2*x[3] + LAM_3*x[4] +
                LAM_4*x[5] + LAM_5*x[6] + LAM_6*x[7]

    # n equation — full PKE (the source of stiffness via 1/Λ ≈ 10^4)
    dx[1] = ((rho - BETA) / LAMBDA) * x[1] + sum_lam_C

    # Precursor balance equations
    dx[2] = (BETA_1 / LAMBDA) * x[1] - LAM_1 * x[2]
    dx[3] = (BETA_2 / LAMBDA) * x[1] - LAM_2 * x[3]
    dx[4] = (BETA_3 / LAMBDA) * x[1] - LAM_3 * x[4]
    dx[5] = (BETA_4 / LAMBDA) * x[1] - LAM_4 * x[5]
    dx[6] = (BETA_5 / LAMBDA) * x[1] - LAM_5 * x[6]
    dx[7] = (BETA_6 / LAMBDA) * x[1] - LAM_6 * x[7]

    # Thermal-hydraulics
    dx[8]  = (P0 * x[1] - HA * (x[8] - x[9])) / (M_F * C_F)
    dx[9]  = (HA * (x[8] - x[9]) - 2 * W_M * C_C * (x[9] - x[10])) / (M_C * C_C)
    dx[10] = (2 * W_M * C_C * (x[9] - x[10]) - Q_SG_DECAY) / (M_SG * C_C)

    dx[11] = one(x[1])
    return nothing
end

# --- Build fat X_entry from union of reach envelopes (11D with n) ---
results_dir = joinpath(@__DIR__, "..", "..", "..", "results")
pred_raw = JSON.parsefile(joinpath(@__DIR__, "..", "..", "..", "reachability", "predicates.json"))

# Helper: compute equilibrium C_i for a given n.
C_eq(n) = [BETA_1/(LAM_1*LAMBDA)*n, BETA_2/(LAM_2*LAMBDA)*n,
           BETA_3/(LAM_3*LAMBDA)*n, BETA_4/(LAM_4*LAMBDA)*n,
           BETA_5/(LAM_5*LAMBDA)*n, BETA_6/(LAM_6*LAMBDA)*n]

# 1. Hot-standby (n ≈ 1e-7..1e-4, T near 275 °C).
sh = pred_raw["mode_boundaries"]["q_shutdown"]["X_entry_polytope"]
hs_n_lo, hs_n_hi   = sh["n_range"]
shutdown_lo = [hs_n_lo; C_eq(hs_n_lo); sh["T_f_range_C"][1];
               sh["T_c_range_C"][1]; sh["T_cold_range_C"][1]]
shutdown_hi = [hs_n_hi; C_eq(hs_n_hi); sh["T_f_range_C"][2];
               sh["T_c_range_C"][2]; sh["T_cold_range_C"][2]]

# 2. Operation (n near 1.0, full-power thermals).
op_path = joinpath(results_dir, "reach_operation_result.mat")
operation_lo = nothing; operation_hi = nothing
if isfile(op_path)
    o = matread(op_path)
    X_lo_op = o["X_lo"]; X_hi_op = o["X_hi"]   # 10 × M (n at row 1)
    operation_lo = [minimum(X_lo_op[i, :]) for i in 1:10]
    operation_hi = [maximum(X_hi_op[i, :]) for i in 1:10]
end

# X_entry construction note:
#   We tried widening to a single hyperrectangle covering shutdown
#   THROUGH high-flux pre-trip (n=1.15, T_c=320, T_f=348).  TMJets
#   refused — single hyperrectangle that contains both (n=1.15, T_c=320)
#   AND (n=1e-7, T_c=270) also contains physically-impossible
#   combinations like (n=1.15, T_c=270).  Over-approximation slack at
#   t=0 was so wide TMJets failed at the first step.
#
#   Right fix: SET SPLITTING — split entry into multiple narrower boxes,
#   run reach on each separately, take union of tubes.  Multi-hour
#   build; flagged for follow-up.
#
#   For now: union shutdown + "operation transient" box (which itself
#   spans n=0.85..1.15, the load-follow + high-flux range).  This gives
#   the widest entry that TMJets can actually solve.  ~7 orders of
#   magnitude on n; ~50 °C on temperatures.

# Operation-transient box (covers steady + load-follow + pre-trip,
# physically consistent — high n correlates with high T_f).
trans_n_lo, trans_n_hi   = 0.85, 1.15
trans_Tf_lo, trans_Tf_hi = 320.0, 332.0   # ~T_c0 + P/hA at the relevant n
trans_Tc_lo, trans_Tc_hi = 305.0, 311.0   # tight around T_c0
trans_Tco_lo, trans_Tco_hi = 285.0, 295.0
trans_lo = [trans_n_lo; C_eq(trans_n_lo); trans_Tf_lo; trans_Tc_lo; trans_Tco_lo]
trans_hi = [trans_n_hi; C_eq(trans_n_hi); trans_Tf_hi; trans_Tc_hi; trans_Tco_hi]

# Union: shutdown + transients box.
fat_lo = copy(shutdown_lo); fat_hi = copy(shutdown_hi)
for src_pair in [(trans_lo, trans_hi)]
    lo, hi = src_pair
    for i in 1:10
        fat_lo[i] = min(fat_lo[i], lo[i])
        fat_hi[i] = max(fat_hi[i], hi[i])
    end
end

# Defensive clamping (n must be > 0; precursors must be in physical range).
fat_lo[1] = max(fat_lo[1], 1e-9)
C_max = [BETA_i/(LAM_i*LAMBDA) * 1.3 for (BETA_i, LAM_i) in
         zip([BETA_1,BETA_2,BETA_3,BETA_4,BETA_5,BETA_6],
             [LAM_1,LAM_2,LAM_3,LAM_4,LAM_5,LAM_6])]
for i in 1:6
    fat_hi[1+i] = min(fat_hi[1+i], C_max[i])
    fat_lo[1+i] = max(fat_lo[1+i], 0.0)
end

# Build the 11D X0 (add augmented time = 0).
x_lo = vcat(fat_lo, 0.0)
x_hi = vcat(fat_hi, 0.0)
X0 = Hyperrectangle(low=x_lo, high=x_hi)

println("\n=== Fat-entry FULL-PKE scram reach (no PJ) ===")
println("  X_entry (n):       [", round(fat_lo[1]; sigdigits=3), ", ",
        round(fat_hi[1]; sigdigits=3), "]")
println("  X_entry (T_f, °C): [", round(fat_lo[8]; digits=2), ", ",
        round(fat_hi[8]; digits=2), "]")
println("  X_entry (T_c, °C): [", round(fat_lo[9]; digits=2), ", ",
        round(fat_hi[9]; digits=2), "]")
println("  Constant u = -8β = ", round(U_SCRAM; digits=4))
println("  Probe horizons: 10 s only (full-PKE stiffness wall + slack growth)")

results = Dict{Float64, Any}()
for T_probe in (10.0,)
    println("\n--- Probe T = $T_probe s ---")
    sys  = BlackBoxContinuousSystem(rhs_scram_full_taylor!, 11)
    prob = InitialValueProblem(sys, X0)
    try
        alg = TMJets(orderT=4, orderQ=2, abstol=1e-9, maxsteps=2_000_000)
        t_start = time()
        sol = solve(prob; T=T_probe, alg=alg)
        elapsed = time() - t_start
        flow = flowpipe(sol)
        n_sets = length(flow)
        println("  TMJets: $n_sets reach-sets in $(round(elapsed; digits=1)) s wall")

        # Per-step overapproximation can produce NaN radii deep into the
        # run when slack growth blows out (full-PKE scram is stiff and
        # this is expected after a few seconds of plant time).  Wrap
        # each step in try/catch so we save the longest contiguous
        # NaN-free prefix.
        n_steps_full = length(flow)
        t_arr_buf     = Float64[]
        n_lo_buf      = Float64[]; n_hi_buf  = Float64[]
        rho_lo_buf    = Float64[]; rho_hi_buf = Float64[]
        Tc_lo_buf     = Float64[]; Tc_hi_buf = Float64[]
        Tf_lo_buf     = Float64[]; Tf_hi_buf = Float64[]
        for R in flow
            try
                s = set(overapproximate(R, Hyperrectangle))
                # Skip if any extracted radius is NaN.
                tval = high(s, 11)
                nlo  = low(s, 1);  nhi  = high(s, 1)
                Tflo = low(s, 8);  Tfhi = high(s, 8)
                Tclo = low(s, 9);  Tchi = high(s, 9)
                if any(isnan, (tval, nlo, nhi, Tflo, Tfhi, Tclo, Tchi))
                    break
                end
                push!(t_arr_buf,  tval)
                push!(n_lo_buf,   nlo);  push!(n_hi_buf,   nhi)
                push!(Tf_lo_buf,  Tflo); push!(Tf_hi_buf,  Tfhi)
                push!(Tc_lo_buf,  Tclo); push!(Tc_hi_buf,  Tchi)
                push!(rho_lo_buf, U_SCRAM + ALPHA_F*(Tfhi - T_F0) + ALPHA_C*(Tchi - T_C0))
                push!(rho_hi_buf, U_SCRAM + ALPHA_F*(Tflo - T_F0) + ALPHA_C*(Tclo - T_C0))
            catch
                break
            end
        end
        n_steps = length(t_arr_buf)
        t_arr     = t_arr_buf
        n_lo_ts   = n_lo_buf;   n_hi_ts   = n_hi_buf
        rho_lo_ts = rho_lo_buf; rho_hi_ts = rho_hi_buf
        Tc_lo_ts  = Tc_lo_buf;  Tc_hi_ts  = Tc_hi_buf
        Tf_lo_ts  = Tf_lo_buf;  Tf_hi_ts  = Tf_hi_buf

        # Build a "last" hyperrectangle from the final NaN-free reach set.
        if n_steps == 0
            error("All reach-sets produced NaN under overapproximate.")
        end
        # Use last buffered values for the "at T_probe" report.
        last_n_lo = n_lo_ts[end];  last_n_hi = n_hi_ts[end]
        last_Tf_lo = Tf_lo_ts[end]; last_Tf_hi = Tf_hi_ts[end]
        last_Tc_lo = Tc_lo_ts[end]; last_Tc_hi = Tc_hi_ts[end]
        last_rho_lo = rho_lo_ts[end]; last_rho_hi = rho_hi_ts[end]
        println("  Extracted $n_steps NaN-free reach-sets out of $(n_steps_full) total")

        @printf "  n at last good step:   [%.4e, %.4e]\n" last_n_lo last_n_hi
        @printf "  T_f at last good step: [%.2f, %.2f] °C\n" last_Tf_lo last_Tf_hi
        @printf "  T_c at last good step: [%.2f, %.2f] °C\n" last_Tc_lo last_Tc_hi
        @printf "  ρ at last good step:   [%.5f, %.5f]\n" last_rho_lo last_rho_hi
        @printf "  last good time: %.4f s\n" t_arr[end]

        results[T_probe] = (status="OK", n_sets=n_sets, elapsed=elapsed,
                            t_arr=t_arr,
                            n_lo_ts=n_lo_ts, n_hi_ts=n_hi_ts,
                            rho_lo_ts=rho_lo_ts, rho_hi_ts=rho_hi_ts,
                            Tc_lo_ts=Tc_lo_ts, Tc_hi_ts=Tc_hi_ts,
                            Tf_lo_ts=Tf_lo_ts, Tf_hi_ts=Tf_hi_ts)
    catch err
        msg = sprint(showerror, err)
        println("  FAILED: ", first(msg, 300))
        results[T_probe] = (status="FAILED", err=first(msg, 300))
        break
    end
end

# Save .mat — same key naming convention as the PJ scram script.
mat_out = joinpath(results_dir, "reach_scram_full_fat.mat")
saved = Dict{String, Any}()
saved["fat_lo"]   = fat_lo
saved["fat_hi"]   = fat_hi
saved["sources"]  = ["shutdown", "operation", "loca"]
saved["sdm_rhs"]  = SDM_RHS
saved["rho_sdm"]  = RHO_SDM
for T_probe in (10.0, 30.0, 60.0)
    haskey(results, T_probe) || continue
    r = results[T_probe]
    r.status == "OK" || continue
    pre = "T_$(Int(T_probe))_"
    saved[pre * "t_arr"]     = r.t_arr
    saved[pre * "n_lo_ts"]   = r.n_lo_ts
    saved[pre * "n_hi_ts"]   = r.n_hi_ts
    saved[pre * "rho_lo_ts"] = r.rho_lo_ts
    saved[pre * "rho_hi_ts"] = r.rho_hi_ts
    saved[pre * "Tc_lo_ts"]  = r.Tc_lo_ts
    saved[pre * "Tc_hi_ts"]  = r.Tc_hi_ts
    saved[pre * "Tf_lo_ts"]  = r.Tf_lo_ts
    saved[pre * "Tf_hi_ts"]  = r.Tf_hi_ts
end
matwrite(mat_out, saved)
println("\nSaved fat-entry full-PKE scram reach to $mat_out")