#!/usr/bin/env julia
#
# reach_scram_pj_fat.jl — scram reach from the union of all mode-entry
# reach envelopes + the LOCA scenario envelope.
#
# Morning-review point 2: the real scram X_entry is the set of all
# states the plant could plausibly be in across every mode + accident
# scenarios.  Computes the bounding-box hull of:
#
#   1. Hot-standby IC box (narrow, from mode_boundaries.q_shutdown.X_entry_polytope)
#   2. Heatup reach envelope (from results/reach_heatup_pj_tight.mat)
#   3. Operation reach envelope (from results/reach_operation_result.mat)
#   4. LOCA reach final-state envelope (from results/reach_loca_operation.mat)
#
# Then runs scram PJ on the fat X_entry and reports per-halfspace result.

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

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

# Plant constants.
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

# X_exit threshold: shutdown_margin halfspace, mirrors predicates.json.
const RHO_SDM = 0.01                                            # 1% dk/k
const SDM_RHS = -RHO_SDM - U_SCRAM + ALPHA_F*T_F0 + ALPHA_C*T_C0   # ≈ 0.00297

# Taylorized scram RHS (same as reach_scram_pj.jl).
@taylorize function rhs_scram_fat_taylor!(dx, x, p, t)
    rho = U_SCRAM + ALPHA_F * (x[7] - T_F0) + ALPHA_C * (x[8] - T_C0)
    sum_lam_C = LAM_1*x[1] + LAM_2*x[2] + LAM_3*x[3] +
                LAM_4*x[4] + LAM_5*x[5] + LAM_6*x[6]
    denom = BETA - rho
    n = LAMBDA * sum_lam_C / denom
    inv_factor = sum_lam_C / denom
    dx[1] = BETA_1 * inv_factor - LAM_1 * x[1]
    dx[2] = BETA_2 * inv_factor - LAM_2 * x[2]
    dx[3] = BETA_3 * inv_factor - LAM_3 * x[3]
    dx[4] = BETA_4 * inv_factor - LAM_4 * x[4]
    dx[5] = BETA_5 * inv_factor - LAM_5 * x[5]
    dx[6] = BETA_6 * inv_factor - LAM_6 * x[6]
    dx[7] = (P0 * n - HA * (x[7] - x[8])) / (M_F * C_F)
    dx[8] = (HA * (x[7] - x[8]) - 2 * W_M * C_C * (x[8] - x[9])) / (M_C * C_C)
    dx[9] = (2 * W_M * C_C * (x[8] - x[9]) - Q_SG_DECAY) / (M_SG * C_C)
    dx[10] = one(x[1])
    return nothing
end

# --- Build fat X_entry from union of reach envelopes ---
results_dir = joinpath(@__DIR__, "..", "..", "..", "results")

# 1. Hot-standby box from predicates.json.
pred_raw = JSON.parsefile(joinpath(@__DIR__, "..", "..", "..", "reachability", "predicates.json"))
sh_entry = pred_raw["mode_boundaries"]["q_shutdown"]["X_entry_polytope"]
hs_n_lo, hs_n_hi = sh_entry["n_range"]
hs_Tf_lo, hs_Tf_hi = sh_entry["T_f_range_C"]
hs_Tc_lo, hs_Tc_hi = sh_entry["T_c_range_C"]
hs_Tcold_lo, hs_Tcold_hi = sh_entry["T_cold_range_C"]

# Precursor bounds at hot-standby: C_i = β_i/(λ_i·Λ) · n, with n in hs_n_range.
function C_range_for_n(n_lo, n_hi)
    C_lo = [BETA_1/(LAM_1*LAMBDA)*n_lo, BETA_2/(LAM_2*LAMBDA)*n_lo,
            BETA_3/(LAM_3*LAMBDA)*n_lo, BETA_4/(LAM_4*LAMBDA)*n_lo,
            BETA_5/(LAM_5*LAMBDA)*n_lo, BETA_6/(LAM_6*LAMBDA)*n_lo]
    C_hi = [BETA_1/(LAM_1*LAMBDA)*n_hi, BETA_2/(LAM_2*LAMBDA)*n_hi,
            BETA_3/(LAM_3*LAMBDA)*n_hi, BETA_4/(LAM_4*LAMBDA)*n_hi,
            BETA_5/(LAM_5*LAMBDA)*n_hi, BETA_6/(LAM_6*LAMBDA)*n_hi]
    return C_lo, C_hi
end

shutdown_C_lo, shutdown_C_hi = C_range_for_n(hs_n_lo, hs_n_hi)
shutdown_lo = [shutdown_C_lo; hs_Tf_lo; hs_Tc_lo; hs_Tcold_lo]   # 9 entries
shutdown_hi = [shutdown_C_hi; hs_Tf_hi; hs_Tc_hi; hs_Tcold_hi]

# 2. Heatup tight envelope (read from results/reach_heatup_pj_tight.mat).
heatup_path = joinpath(results_dir, "reach_heatup_pj_tight.mat")
heatup_lo = fill(NaN, 9); heatup_hi = fill(NaN, 9)
if isfile(heatup_path)
    h = matread(heatup_path)
    # Use the longest-horizon probe's envelope (T=300 or whatever's there).
    Tprobes = sort(collect([parse(Int, replace(split(k, "_")[2], ".0" => "")) for k in keys(h) if startswith(k, "T_") && endswith(k, "_Tc_lo_ts")]))
    if !isempty(Tprobes)
        T = Tprobes[end]
        pre = "T_$(T)_"
        heatup_lo = [minimum(vec(h[pre*"Tf_lo_ts"])) - 5, # pad slightly
                     fill(0.0, 6)...]   # 6 Cs from full-state operating range
        # Rebuild more carefully:
        heatup_C_lo, heatup_C_hi = C_range_for_n(1e-3, 2e-3)
        heatup_lo = [heatup_C_lo; minimum(vec(h[pre*"Tf_lo_ts"]));
                     minimum(vec(h[pre*"Tc_lo_ts"])); minimum(vec(h[pre*"Tco_lo_ts"]))]
        heatup_hi = [heatup_C_hi; maximum(vec(h[pre*"Tf_hi_ts"]));
                     maximum(vec(h[pre*"Tc_hi_ts"])); maximum(vec(h[pre*"Tco_hi_ts"]))]
    end
else
    println("Warning: $heatup_path not found; using hot-standby as fallback.")
    heatup_lo = shutdown_lo; heatup_hi = shutdown_hi
end

# 3. Operation reach envelope from results/reach_operation_result.mat.
op_path = joinpath(results_dir, "reach_operation_result.mat")
op_lo = fill(NaN, 9); op_hi = fill(NaN, 9)
if isfile(op_path)
    o = matread(op_path)
    X_lo_op = o["X_lo"]; X_hi_op = o["X_hi"]  # 10 × M
    # Drop row 1 (n) for the 9-state PJ scram model.
    op_lo = [minimum(X_lo_op[i, :]) for i in 2:10]
    op_hi = [maximum(X_hi_op[i, :]) for i in 2:10]
end

# 4. LOCA operation reach final envelope.
loca_path = joinpath(results_dir, "reach_loca_operation.mat")
loca_lo = fill(NaN, 9); loca_hi = fill(NaN, 9)
if isfile(loca_path)
    l = matread(loca_path)
    X_env_lo = vec(l["X_envelope_lo"])  # 10 entries
    X_env_hi = vec(l["X_envelope_hi"])
    loca_lo = X_env_lo[2:10]   # drop n
    loca_hi = X_env_hi[2:10]
end

# Union = element-wise min/max across all sources.
sources = [
    ("shutdown", shutdown_lo, shutdown_hi),
    ("heatup", heatup_lo, heatup_hi),
    ("operation", op_lo, op_hi),
    ("loca", loca_lo, loca_hi),
]

fat_lo = fill(+Inf, 9); fat_hi = fill(-Inf, 9)
for (name, lo, hi) in sources
    any(isnan, lo) || any(isnan, hi) && continue
    for i in 1:9
        fat_lo[i] = min(fat_lo[i], lo[i])
        fat_hi[i] = max(fat_hi[i], hi[i])
    end
end

# LOCA values are absurd on precursors (linear reach numeric blowup) — clamp
# C_i bounds to something physically plausible.  Cap at 1000× the
# operating-point C values (generous).
C_clamp_hi = [BETA_i/(LAM_i*LAMBDA) * 1.2e0 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])]
C_clamp_lo = [BETA_i/(LAM_i*LAMBDA) * 1e-7 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_lo[i] = max(fat_lo[i], C_clamp_lo[i])
    fat_hi[i] = min(fat_hi[i], C_clamp_hi[i])
end
# Temperatures: clamp to model's trust region ~[200, 400] °C.
for i in 7:9
    fat_lo[i] = max(fat_lo[i], 200.0)
    fat_hi[i] = min(fat_hi[i], 400.0)
end

println("\n=== Fat X_entry(scram) from union of reach envelopes ===")
state_names_9 = ["C1","C2","C3","C4","C5","C6","T_f","T_c","T_cold"]
for i in 1:9
    println("  $(state_names_9[i]):  [$(round(fat_lo[i]; sigdigits=4)), $(round(fat_hi[i]; sigdigits=4))]")
end

# Add augmented time state.
x_lo_full = [fat_lo; 0.0]
x_hi_full = [fat_hi; 0.0]
X0 = Hyperrectangle(low=x_lo_full, high=x_hi_full)

# --- Scram reach ---
results = Dict{Float64, Any}()
for T_probe in (10.0, 30.0, 60.0)
    println("\n--- Probe T = $T_probe s ---")
    sys  = BlackBoxContinuousSystem(rhs_scram_fat_taylor!, 10)
    prob = InitialValueProblem(sys, X0)
    try
        alg = TMJets(orderT=4, orderQ=2, abstol=1e-9, maxsteps=100000)
        t0 = time()
        sol = solve(prob; T=T_probe, alg=alg)
        elapsed = time() - t0
        flow = flowpipe(sol)
        flow_hr = overapproximate(flow, Hyperrectangle)
        n_sets = length(flow_hr)
        println("  TMJets: $n_sets reach-sets in $(round(elapsed; digits=1)) s")

        # --- Per-step envelopes for plotting tubes ---
        t_arr     = zeros(n_sets)
        n_lo_ts   = zeros(n_sets); n_hi_ts   = zeros(n_sets)
        rho_lo_ts = zeros(n_sets); rho_hi_ts = zeros(n_sets)
        Tc_lo_ts  = zeros(n_sets); Tc_hi_ts  = zeros(n_sets)
        Tf_lo_ts  = zeros(n_sets); Tf_hi_ts  = zeros(n_sets)
        for (k, R) in enumerate(flow_hr)
            s = set(R)
            t_arr[k] = high(s, 10)
            sumLC_lo_k = LAM_1*low(s,1) + LAM_2*low(s,2) + LAM_3*low(s,3) +
                         LAM_4*low(s,4) + LAM_5*low(s,5) + LAM_6*low(s,6)
            sumLC_hi_k = LAM_1*high(s,1) + LAM_2*high(s,2) + LAM_3*high(s,3) +
                         LAM_4*high(s,4) + LAM_5*high(s,5) + LAM_6*high(s,6)
            rho_lo_k = U_SCRAM + ALPHA_F*(high(s,7) - T_F0) + ALPHA_C*(high(s,8) - T_C0)
            rho_hi_k = U_SCRAM + ALPHA_F*(low(s,7)  - T_F0) + ALPHA_C*(low(s,8)  - T_C0)
            denom_lo_k = BETA - rho_hi_k
            denom_hi_k = BETA - rho_lo_k
            n_lo_ts[k]   = denom_lo_k > 0 ? LAMBDA * sumLC_lo_k / denom_hi_k : NaN
            n_hi_ts[k]   = denom_lo_k > 0 ? LAMBDA * sumLC_hi_k / denom_lo_k : NaN
            rho_lo_ts[k] = rho_lo_k
            rho_hi_ts[k] = rho_hi_k
            Tc_lo_ts[k]  = low(s, 8); Tc_hi_ts[k] = high(s, 8)
            Tf_lo_ts[k]  = low(s, 7); Tf_hi_ts[k] = high(s, 7)
        end

        last = set(flow_hr[end])
        sumLC_lo = LAM_1*low(last,1) + LAM_2*low(last,2) + LAM_3*low(last,3) +
                   LAM_4*low(last,4) + LAM_5*low(last,5) + LAM_6*low(last,6)
        sumLC_hi = LAM_1*high(last,1) + LAM_2*high(last,2) + LAM_3*high(last,3) +
                   LAM_4*high(last,4) + LAM_5*high(last,5) + LAM_6*high(last,6)
        rho_lo = U_SCRAM + ALPHA_F*(low(last,7) - T_F0) + ALPHA_C*(high(last,8) - T_C0)
        rho_hi = U_SCRAM + ALPHA_F*(high(last,7) - T_F0) + ALPHA_C*(low(last,8) - T_C0)
        denom_lo = BETA - rho_hi
        denom_hi = BETA - rho_lo
        n_lo = denom_lo > 0 ? LAMBDA * sumLC_lo / denom_hi : NaN
        n_hi = denom_lo > 0 ? LAMBDA * sumLC_hi / denom_lo : NaN

        # shutdown_margin discharge: alpha_f*T_f + alpha_c*T_c ≤ SDM_RHS.
        # Coefficients negative → sup over the box at low(T_f), low(T_c).
        sdm_lhs_hi = ALPHA_F*low(last,7) + ALPHA_C*low(last,8)
        sdm_ok = sdm_lhs_hi <= SDM_RHS

        println("  n envelope:    [$(round(n_lo; sigdigits=4)), $(round(n_hi; sigdigits=4))]")
        println("  T_c envelope:  [$(round(low(last,8); digits=2)), $(round(high(last,8); digits=2))] °C")
        println("  T_f envelope:  [$(round(low(last,7); digits=2)), $(round(high(last,7); digits=2))] °C")
        println("  rho envelope:  [$(round(rho_lo; sigdigits=4)), $(round(rho_hi; sigdigits=4))]  (shutdown margin = $(round(-rho_hi; sigdigits=4)) dk/k)")
        println("  shutdown_margin LHS sup: $(round(sdm_lhs_hi; sigdigits=4))  vs RHS $(round(SDM_RHS; sigdigits=4))  → $(sdm_ok ? "✓ DISCHARGED" : "× not yet")")

        results[T_probe] = (status="OK", n=(n_lo, n_hi), elapsed=elapsed,
                            rho=(rho_lo, rho_hi), sdm_lhs_hi=sdm_lhs_hi, sdm_ok=sdm_ok,
                            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
        println("  FAILED: ", first(sprint(showerror, err), 300))
        results[T_probe] = (status="FAILED",)
        break
    end
end

println("\n=== Summary ===")
for T_probe in (10.0, 30.0, 60.0)
    haskey(results, T_probe) || continue
    r = results[T_probe]
    if r.status == "OK"
        ok_str = r.sdm_ok ? "✓ shutdown_margin DISCHARGED" : "× shutdown_margin not yet"
        @printf "  T = %4.0f s: n ∈ [%.3e, %.3e]  ρ ∈ [%.4f, %.4f]  %s\n" T_probe r.n[1] r.n[2] r.rho[1] r.rho[2] ok_str
    else
        println("  T = $T_probe s: FAILED")
    end
end

mat_out = joinpath(results_dir, "reach_scram_pj_fat.mat")
saved = Dict{String,Any}("fat_lo" => fat_lo, "fat_hi" => fat_hi,
                         "sources" => ["shutdown", "heatup_tight", "operation", "loca_operation"],
                         "sdm_rhs" => SDM_RHS, "rho_sdm" => RHO_SDM)
for (T_probe, r) in results
    if r.status == "OK"
        saved["T_$(Int(T_probe))_n_lo"] = r.n[1]
        saved["T_$(Int(T_probe))_n_hi"] = r.n[2]
        saved["T_$(Int(T_probe))_rho_lo"] = r.rho[1]
        saved["T_$(Int(T_probe))_rho_hi"] = r.rho[2]
        saved["T_$(Int(T_probe))_sdm_lhs_hi"] = r.sdm_lhs_hi
        saved["T_$(Int(T_probe))_sdm_ok"] = r.sdm_ok ? 1.0 : 0.0
        # Per-step time series for tube plotting.
        saved["T_$(Int(T_probe))_t_arr"]     = r.t_arr
        saved["T_$(Int(T_probe))_n_lo_ts"]   = r.n_lo_ts
        saved["T_$(Int(T_probe))_n_hi_ts"]   = r.n_hi_ts
        saved["T_$(Int(T_probe))_rho_lo_ts"] = r.rho_lo_ts
        saved["T_$(Int(T_probe))_rho_hi_ts"] = r.rho_hi_ts
        saved["T_$(Int(T_probe))_Tc_lo_ts"]  = r.Tc_lo_ts
        saved["T_$(Int(T_probe))_Tc_hi_ts"]  = r.Tc_hi_ts
        saved["T_$(Int(T_probe))_Tf_lo_ts"]  = r.Tf_lo_ts
        saved["T_$(Int(T_probe))_Tf_hi_ts"]  = r.Tf_hi_ts
    end
end
matwrite(mat_out, saved)
println("\nSaved fat-entry scram reach to $mat_out")