#!/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

# 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")
        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
        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")

        results[T_probe] = (status="OK", n=(n_lo, n_hi), elapsed=elapsed)
    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"
        @printf "  T = %4.0f s: n ∈ [%.3e, %.3e]\n" T_probe r.n[1] r.n[2]
    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"])
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]
    end
end
matwrite(mat_out, saved)
println("\nSaved fat-entry scram reach to $mat_out")