#!/usr/bin/env julia
#
# reach_scram_pj.jl — nonlinear reach on scram, prompt-jump model.
#
# Scram obligation: from any operating-envelope state, drive n down to
# n <= 1e-4 within T_max = 60 s.  Constant control u = -8*beta (rods
# slammed in).  Q_sg = 3% P0 (decay-heat-level sink, placeholder).
#
# 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

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

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

        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)
    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_subcrit = r.n_final[2] <= 1e-4 ? "✓ subcritical (n ≤ 1e-4)" : "× still above 1e-4"
        println("  T = $(T_probe) s: $(r.n_sets) sets, $(round(r.elapsed; digits=1))s wall — n ∈ [$(round(r.n_final[1]; sigdigits=3)), $(round(r.n_final[2]; sigdigits=3))]  $ok_subcrit")
    else
        println("  T = $(T_probe) s: FAILED")
    end
end

mat_out = joinpath(@__DIR__, "..", "..", "reachability", "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]
    end
end
matwrite(mat_out, saved)
println("\nSaved scram envelope summaries to $mat_out")