#!/usr/bin/env julia
#
# reach_heatup_nonlinear.jl — nonlinear reach on heatup mode via TMJets.
#
# STATUS: partial success. 10 s horizon works; longer horizons fail.
#
# Full 10-state nonlinear reach of the PKE + thermal-hydraulics
# closed-loop is limited to ~10 s horizons by prompt-neutron stiffness
# (Lambda = 1e-4 s).  TMJets' adaptive step size shrinks to ~1 ms to
# resolve prompt dynamics, then hits numerical precision floor and
# produces NaN around t ~ 10 s.  10,583 steps for 10 s of sim time.
#
# For heatup's 5-hour horizon (18 million prompt timescales), we need
# singular-perturbation reduction of the neutronics: treat n as a
# quasi-static algebraic function of (T, C, rho) rather than a dynamic
# state.  The reduced-order PKE replaces the stiff dn/dt equation with
# an algebraic relation, removing the fast timescale from the reach
# problem.
#
# The 10 s result is a validation of the Taylor-model approach, not
# a usable safety proof.  The machinery works; the model needs
# reduction before the horizon is meaningful.
#
# Saturation disabled (ctrl_heatup_unsat structure), all constants
# inlined so @taylorize can generate a specialized Taylor-model RHS.
# Time carried as augmented state x[11].

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

using LinearAlgebra
using ReachabilityAnalysis, LazySets
using JSON

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

const P0   = 1e9
const M_F  = 50000.0
const C_F  = 300.0
const M_C  = 20000.0
const C_C  = 5450.0
const HA   = 5e7
const W_M  = 5000.0
const M_SG = 30000.0

const ALPHA_F = -2.5e-5
const ALPHA_C = -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

# --- Controller constants ---
const T_STANDBY    = T_C0 - 33.333333
const RAMP_RATE_CS = 28.0 / 3600
const KP_HEATUP    = 1e-4

# --- Taylorized RHS with augmented time (x[11] = t, dx[11] = 1) ---
# TMJets handles augmented time better than accessing the solver's t arg
# directly inside the Taylor-model rewriting.
@taylorize function rhs_heatup_taylor!(dx, x, p, t)
    # Inlined cancellation: rho_total = Kp * (T_ref - T_c) after cancellation.
    # T_ref uses x[11] as the time-state (no min; valid while T_ref < T_c0).
    rho = KP_HEATUP * (T_STANDBY + RAMP_RATE_CS * x[11] - x[9])

    # Neutronics (explicit — no loops, no function calls)
    dx[1] = (rho - BETA) / LAMBDA * x[1] +
            LAM_1*x[2] + LAM_2*x[3] + LAM_3*x[4] +
            LAM_4*x[5] + LAM_5*x[6] + LAM_6*x[7]
    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])) / (M_SG * C_C)

    # Time (augmented state for reach; dt/dt = 1)
    dx[11] = one(x[1])
    return nothing
end

# --- Build X_entry from predicates.json ---
pred_path = joinpath(@__DIR__, "..", "..", "reachability", "predicates.json")
pred_raw  = JSON.parsefile(pred_path)
entry     = pred_raw["mode_boundaries"]["q_heatup"]["X_entry_polytope"]

n_lo, n_hi           = entry["n_range"]
T_f_lo, T_f_hi       = entry["T_f_range_C"]
T_c_lo, T_c_hi       = entry["T_c_range_C"]
T_cold_lo, T_cold_hi = entry["T_cold_range_C"]

n_mid = 0.5 * (n_lo + n_hi)
C_mid_1 = (BETA_1 / (LAM_1 * LAMBDA)) * n_mid
C_mid_2 = (BETA_2 / (LAM_2 * LAMBDA)) * n_mid
C_mid_3 = (BETA_3 / (LAM_3 * LAMBDA)) * n_mid
C_mid_4 = (BETA_4 / (LAM_4 * LAMBDA)) * n_mid
C_mid_5 = (BETA_5 / (LAM_5 * LAMBDA)) * n_mid
C_mid_6 = (BETA_6 / (LAM_6 * LAMBDA)) * n_mid

x_lo = [n_lo;
        C_mid_1 * (n_lo / n_mid); C_mid_2 * (n_lo / n_mid);
        C_mid_3 * (n_lo / n_mid); C_mid_4 * (n_lo / n_mid);
        C_mid_5 * (n_lo / n_mid); C_mid_6 * (n_lo / n_mid);
        T_f_lo; T_c_lo; T_cold_lo;
        0.0]   # augmented time

x_hi = [n_hi;
        C_mid_1 * (n_hi / n_mid); C_mid_2 * (n_hi / n_mid);
        C_mid_3 * (n_hi / n_mid); C_mid_4 * (n_hi / n_mid);
        C_mid_5 * (n_hi / n_mid); C_mid_6 * (n_hi / n_mid);
        T_f_hi; T_c_hi; T_cold_hi;
        0.0]

X0 = Hyperrectangle(low=x_lo, high=x_hi)

println("\n=== Nonlinear heatup reach (saturation disabled, @taylorize RHS) ===")
println("  X_entry: n ∈ [$(n_lo), $(n_hi)], T_c ∈ [$(T_c_lo), $(T_c_hi)] C")

for T_probe in (10.0, 60.0, 300.0)
    println("\n--- Probe T = $T_probe s ---")
    sys  = BlackBoxContinuousSystem(rhs_heatup_taylor!, 11)
    prob = InitialValueProblem(sys, X0)
    try
        alg = TMJets(orderT=4, orderQ=2, abstol=1e-10, maxsteps=50000)
        sol = solve(prob; T=T_probe, alg=alg)
        flow = flowpipe(sol)
        n_sets = length(flow)
        println("  TMJets: $n_sets reach-sets")

        # Overapproximate each Taylor-model reach-set as a hyperrectangle
        # so we can query envelopes.  This is standard LazySets usage.
        flow_hr = overapproximate(flow, Hyperrectangle)
        n_sets_hr = length(flow_hr)

        # Envelope per state over the whole flowpipe.
        n_lo_env  = minimum(low(set(R), 1)  for R in flow_hr)
        n_hi_env  = maximum(high(set(R), 1) for R in flow_hr)
        Tc_lo_env = minimum(low(set(R), 9)  for R in flow_hr)
        Tc_hi_env = maximum(high(set(R), 9) for R in flow_hr)
        Tf_lo_env = minimum(low(set(R), 8)  for R in flow_hr)
        Tf_hi_env = maximum(high(set(R), 8) for R in flow_hr)
        Tcold_lo  = minimum(low(set(R), 10) for R in flow_hr)
        Tcold_hi  = maximum(high(set(R), 10) for R in flow_hr)
        t_final   = maximum(high(set(R), 11) for R in flow_hr)

        println("  t reached:    $(round(t_final; digits=1)) s of $T_probe")
        println("  n envelope:   [$(round(n_lo_env; sigdigits=4)), $(round(n_hi_env; sigdigits=4))]")
        println("  T_f envelope: [$(round(Tf_lo_env; digits=2)), $(round(Tf_hi_env; digits=2))] C")
        println("  T_c envelope: [$(round(Tc_lo_env; digits=2)), $(round(Tc_hi_env; digits=2))] C")
        println("  T_cold env:   [$(round(Tcold_lo; digits=2)), $(round(Tcold_hi; digits=2))] C")
    catch err
        msg = sprint(showerror, err)
        println("  FAILED: ", first(msg, 300))
    end
end