#!/usr/bin/env julia
#
# reach_heatup_pj_tight_full.jl — tight-entry heatup PJ reach,
# per-timestep envelopes saved for plotting (T_c, T_h, T_cold, rho, n).
#
# Identical dynamics to reach_heatup_pj_tight.jl but keeps the full
# tube data so we can overlay T_c and T_h (and infer power) on one plot.

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

using LinearAlgebra
using ReachabilityAnalysis, LazySets
using MAT

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_tf!(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

# Tight X_entry.
n_lo, n_hi         = 1.0e-3, 2.0e-3
T_f_lo, T_f_hi     = 285.0, 291.0
T_c_lo, T_c_hi     = 285.0, 291.0
T_cold_lo, T_cold_hi = 278.0, 285.0

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=== Tight-entry heatup PJ reach — saving full per-step envelopes ===")

T_probe = 300.0
sys  = BlackBoxContinuousSystem(rhs_heatup_pj_tf!, 10)
prob = InitialValueProblem(sys, X0)
alg  = TMJets(orderT=4, orderQ=2, abstol=1e-9, maxsteps=100000)
t_start = time()
sol = solve(prob; T=T_probe, alg=alg)
println("  Wall time: $(round(time() - t_start; digits=1)) s")
flow_hr = overapproximate(flowpipe(sol), Hyperrectangle)
n_steps = length(flow_hr)
println("  $n_steps reach-sets")

t_arr     = 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)
Tco_lo_ts = zeros(n_steps); Tco_hi_ts = 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)

for (k, R) in enumerate(flow_hr)
    s = set(R)
    t_arr[k]    = high(s, 10)
    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)
    Tco_lo_ts[k] = low(s, 9); Tco_hi_ts[k] = high(s, 9)
    sumLC_lo = 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 = 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)
    # Ramp-ref bounds at this reach-set's t.
    t_hi_here = high(s, 10)
    t_lo_here = low(s, 10)
    # T_ref at these times (monotone increasing):
    Tref_lo = T_STANDBY + RAMP_RATE_CS * t_lo_here
    Tref_hi = T_STANDBY + RAMP_RATE_CS * t_hi_here
    rho_lo_here = KP_HEATUP * (Tref_lo - high(s, 8))
    rho_hi_here = KP_HEATUP * (Tref_hi - low(s, 8))
    rho_lo_ts[k] = rho_lo_here
    rho_hi_ts[k] = rho_hi_here
    denom_lo = BETA - rho_hi_here
    denom_hi = BETA - rho_lo_here
    if denom_lo > 0
        n_lo_ts[k] = LAMBDA * sumLC_lo / denom_hi
        n_hi_ts[k] = LAMBDA * sumLC_hi / denom_lo
    end
end

println("  T_c envelope:   [$(round(minimum(Tc_lo_ts); digits=2)), $(round(maximum(Tc_hi_ts); digits=2))] °C")
println("  T_hot envelope: [$(round(minimum(2 .* Tc_lo_ts .- Tco_hi_ts); digits=2)), $(round(maximum(2 .* Tc_hi_ts .- Tco_lo_ts); digits=2))] °C")
println("  T_cold env:     [$(round(minimum(Tco_lo_ts); digits=2)), $(round(maximum(Tco_hi_ts); digits=2))] °C")
println("  rho envelope:   [$(round(minimum(rho_lo_ts); sigdigits=4)), $(round(maximum(rho_hi_ts); sigdigits=4))]")
println("  rho / beta:     [$(round(minimum(rho_lo_ts)/BETA; digits=3)), $(round(maximum(rho_hi_ts)/BETA; digits=3))]")

mat_out = joinpath(@__DIR__, "..", "..", "reachability", "reach_heatup_pj_tight_full.mat")
matwrite(mat_out, Dict(
    "t_arr"     => t_arr,
    "Tc_lo_ts"  => Tc_lo_ts,  "Tc_hi_ts"  => Tc_hi_ts,
    "Tf_lo_ts"  => Tf_lo_ts,  "Tf_hi_ts"  => Tf_hi_ts,
    "Tco_lo_ts" => Tco_lo_ts, "Tco_hi_ts" => Tco_hi_ts,
    "n_lo_ts"   => n_lo_ts,   "n_hi_ts"   => n_hi_ts,
    "rho_lo_ts" => rho_lo_ts, "rho_hi_ts" => rho_hi_ts,
    "T_probe"   => T_probe,
    "beta"      => BETA,
    "Kp"        => KP_HEATUP,
    "T_c0"      => T_C0,
    "T_cold0"   => T_COLD0,
    "T_standby" => T_STANDBY,
))
println("\nSaved full per-step envelopes to $mat_out")