#!/usr/bin/env julia
#
# reach_heatup_pj_tight.jl — heatup PJ reach with a tighter X_entry.
#
# The default X_entry (from predicates.json) has T_c ∈ [281, 295] — 14 K
# wide.  The baseline PJ reach at T=300s produced T_c envelope
# [272.4, 295.0], violating the low-T_avg trip at 280.
#
# Hypothesis: entry-box width is contributing to tube growth.  Try
# T_c ∈ [285, 291] (6 K) and T_f matched, see if the lower envelope
# rises above 280.

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

using LinearAlgebra
using ReachabilityAnalysis, LazySets
using MAT

# Same constants as reach_heatup_pj.jl.
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 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 T_STANDBY    = T_C0 - 33.333333
const RAMP_RATE_CS = 28.0 / 3600
const KP_HEATUP    = 1e-4

@taylorize function rhs_heatup_pj_tight!(dx, x, p, t)
    rho = KP_HEATUP * (T_STANDBY + RAMP_RATE_CS * x[10] - x[8])
    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])) / (M_SG * C_C)
    dx[10] = one(x[1])
    return nothing
end

# Tighter X_entry on T_c and T_f specifically.
n_lo, n_hi         = 1.0e-3, 2.0e-3    # narrower n
T_f_lo, T_f_hi     = 285.0, 291.0      # was 275–295
T_c_lo, T_c_hi     = 285.0, 291.0      # was 281–295
T_cold_lo, T_cold_hi = 278.0, 285.0    # was 270–281 (shifted up)

n_mid = 0.5 * (n_lo + n_hi)
C_mid = [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_mid

x_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]
x_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=== Heatup PJ reach with TIGHTENED X_entry ===")
println("  T_c ∈ [$(T_c_lo), $(T_c_hi)] (width 6 K, was 14 K)")
println("  T_f ∈ [$(T_f_lo), $(T_f_hi)]")
println("  n-implied ∈ [$(n_lo), $(n_hi)]")

for T_probe in (60.0, 300.0)
    println("\n--- Probe T = $T_probe s ---")
    sys  = BlackBoxContinuousSystem(rhs_heatup_pj_tight!, 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)
        flow_hr = overapproximate(flow, Hyperrectangle)
        Tc_lo_env = minimum(low(set(R), 8)  for R in flow_hr)
        Tc_hi_env = maximum(high(set(R), 8) for R in flow_hr)
        Tf_lo_env = minimum(low(set(R), 7)  for R in flow_hr)
        Tf_hi_env = maximum(high(set(R), 7) for R in flow_hr)
        println("  $(length(flow_hr)) sets in $(round(elapsed; digits=1))s")
        println("  T_c envelope: [$(round(Tc_lo_env; digits=2)), $(round(Tc_hi_env; digits=2))] °C")
        println("  T_f envelope: [$(round(Tf_lo_env; digits=2)), $(round(Tf_hi_env; digits=2))] °C")
        println("  Low-T_avg trip (T_c ≥ 280): $(Tc_lo_env >= 280 ? "✅ DISCHARGED" : "× still loose")")
    catch err
        println("  FAILED: ", first(sprint(showerror, err), 200))
    end
end