#!/usr/bin/env julia
#
# reach_scram_pj.jl — nonlinear reach on scram, prompt-jump model.
#
# Scram obligation: from any operating-envelope state, drive total
# reactivity below the shutdown-margin threshold (rho <= -0.01, i.e.
# 1% dk/k subcritical) within T_max = 60 s.  Constant control
# u = -8*beta (rods slammed in).  Q_sg = 3% P0 (decay-heat-level sink,
# placeholder).
#
# X_exit halfspace (from reachability/predicates.json::shutdown_margin):
#   alpha_f * T_f + alpha_c * T_c <= 0.00402
# discharged when sup over reach set of LHS <= 0.00402.
#
# 9-state PJ model (10D with augmented time).

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

using LinearAlgebra
using ReachabilityAnalysis, LazySets
using JSON
using MAT

# 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          # rod worth applied at scram
const Q_SG_DECAY = 0.03 * P0       # constant decay-heat-level sink

# 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.00402

# Taylorized scram RHS, PJ form.
@taylorize function rhs_scram_pj_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

# X_entry — small box around operating point: scram could fire from anywhere
# in operation, but for demo we take a tight envelope and propagate.
n_op  = 1.0
C_op  = [BETA_1/(LAM_1*LAMBDA), BETA_2/(LAM_2*LAMBDA),
         BETA_3/(LAM_3*LAMBDA), BETA_4/(LAM_4*LAMBDA),
         BETA_5/(LAM_5*LAMBDA), BETA_6/(LAM_6*LAMBDA)] .* n_op

x_lo = [C_op[1] * 0.99, C_op[2] * 0.99, C_op[3] * 0.99,
        C_op[4] * 0.99, C_op[5] * 0.99, C_op[6] * 0.99,
        T_F0 - 1.0, T_C0 - 1.0, T_COLD0 - 1.0, 0.0]
x_hi = [C_op[1] * 1.01, C_op[2] * 1.01, C_op[3] * 1.01,
        C_op[4] * 1.01, C_op[5] * 1.01, C_op[6] * 1.01,
        T_F0 + 1.0, T_C0 + 1.0, T_COLD0 + 1.0, 0.0]
X0 = Hyperrectangle(low=x_lo, high=x_hi)

println("\n=== Nonlinear scram reach, prompt-jump model ===")
println("  X_entry: small box around operating point (n ≈ 1.0)")
println("  Constant u = -8*beta = $(round(U_SCRAM; digits=4))")
println("  Q_sg = 3% P0 (decay-heat sink)")
println("  T_max = 60 s")
println("  X_exit: alpha_f*T_f + alpha_c*T_c <= $(round(SDM_RHS; sigdigits=4))  (rho <= -$(RHO_SDM))")

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_pj_taylor!, 10)
    prob = InitialValueProblem(sys, X0)
    try
        alg = TMJets(orderT=4, orderQ=2, abstol=1e-9, maxsteps=100000)
        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")

        flow_hr = overapproximate(flow, Hyperrectangle)

        # --- Per-step envelopes for plotting tubes ---
        n_steps = length(flow_hr)
        t_arr     = zeros(n_steps)
        n_lo_ts   = zeros(n_steps); n_hi_ts   = zeros(n_steps)
        rho_lo_ts = zeros(n_steps); rho_hi_ts = zeros(n_steps)
        Tc_lo_ts  = zeros(n_steps); Tc_hi_ts  = zeros(n_steps)
        Tf_lo_ts  = zeros(n_steps); Tf_hi_ts  = zeros(n_steps)
        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

        # Reconstruct n at last time step from C and T_c.
        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_final_lo = LAMBDA * sumLC_lo / denom_hi
        n_final_hi = LAMBDA * sumLC_hi / denom_lo
        Tc_final = (low(last, 8), high(last, 8))
        Tf_final = (low(last, 7), high(last, 7))
        Tcold_final = (low(last, 9), high(last, 9))
        # shutdown_margin halfspace LHS: alpha_f*T_f + alpha_c*T_c.
        # 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_lhs_lo = ALPHA_F*high(last,7) + ALPHA_C*high(last,8)
        rho_max = U_SCRAM + ALPHA_F*(low(last,7) - T_F0) +
                            ALPHA_C*(low(last,8) - T_C0)
        rho_min = U_SCRAM + ALPHA_F*(high(last,7) - T_F0) +
                            ALPHA_C*(high(last,8) - T_C0)
        sdm_ok = sdm_lhs_hi <= SDM_RHS

        println("  n at T_probe (reconstructed): [$(round(n_final_lo; sigdigits=4)), $(round(n_final_hi; sigdigits=4))]")
        println("  T_c at T_probe: [$(round(Tc_final[1]; digits=2)), $(round(Tc_final[2]; digits=2))] °C")
        println("  T_f at T_probe: [$(round(Tf_final[1]; digits=2)), $(round(Tf_final[2]; digits=2))] °C")
        println("  rho at T_probe: [$(round(rho_min; sigdigits=4)), $(round(rho_max; sigdigits=4))]  (shutdown margin = $(round(-rho_max; 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_sets=n_sets, elapsed=elapsed,
                            n_final=(n_final_lo, n_final_hi),
                            Tc=Tc_final, Tf=Tf_final, Tcold=Tcold_final,
                            sdm_lhs=(sdm_lhs_lo, sdm_lhs_hi),
                            rho=(rho_min, rho_max),
                            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
        msg = sprint(showerror, err)
        println("  FAILED: ", first(msg, 300))
        results[T_probe] = (status="FAILED", err=first(msg, 300))
        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"
        println("  T = $(T_probe) s: $(r.n_sets) sets, $(round(r.elapsed; digits=1))s wall — rho ∈ [$(round(r.rho[1]; sigdigits=3)), $(round(r.rho[2]; sigdigits=3))]  $ok_str")
    else
        println("  T = $(T_probe) s: FAILED")
    end
end

mat_out = joinpath(@__DIR__, "..", "..", "..", "results", "reach_scram_pj_result.mat")
saved = Dict{String, Any}("probe_horizons" => collect((10.0, 30.0, 60.0)))
for T_probe in (10.0, 30.0, 60.0)
    haskey(results, T_probe) || continue
    r = results[T_probe]
    if r.status == "OK"
        saved["T_$(Int(T_probe))_n_lo"]    = r.n_final[1]
        saved["T_$(Int(T_probe))_n_hi"]    = r.n_final[2]
        saved["T_$(Int(T_probe))_Tc_lo"]   = r.Tc[1]
        saved["T_$(Int(T_probe))_Tc_hi"]   = r.Tc[2]
        saved["T_$(Int(T_probe))_Tf_lo"]   = r.Tf[1]
        saved["T_$(Int(T_probe))_Tf_hi"]   = r.Tf[2]
        saved["T_$(Int(T_probe))_Tcold_lo"] = r.Tcold[1]
        saved["T_$(Int(T_probe))_Tcold_hi"] = r.Tcold[2]
        saved["T_$(Int(T_probe))_sdm_lhs_hi"] = r.sdm_lhs[2]
        saved["T_$(Int(T_probe))_rho_max"]    = r.rho[2]
        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
saved["sdm_rhs"] = SDM_RHS
saved["rho_sdm"] = RHO_SDM
matwrite(mat_out, saved)
println("\nSaved scram envelope summaries to $mat_out")