#!/usr/bin/env julia
#
# reach_heatup_pj_sd.jl — heatup PJ reach WITH bounded Q_sg steam dump.
#
# Adds an augmented state x[11] = Q_sg as a bounded parameter (dx[11] = 0).
# The reach propagates the entry-box range of Q_sg forward; at each reach
# set, the Q_sg extent is the disturbance envelope the controller has to
# reject.  Physical story: operator controls secondary steam dump; actual
# value unknown but bounded in [0, 0.05·P_0].
#
# Time is carried as x[12] = t (dt/dt = 1).
#
# Diagnostic script — uses the tight-entry + steam-dump config.

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

using LinearAlgebra
using ReachabilityAnalysis, LazySets
using JSON
using MAT
using TOML

# Plant constants.
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

# 12-state RHS: [C_1..C_6, T_f, T_c, T_cold, Q_sg, t]
# dx[10] = 0   (Q_sg is a bounded parameter)
# dx[11] = 1   (time)
# Hmm — indexing conflict with original version.  Reorganize:
#   x[1..6]  = C_1..C_6
#   x[7]     = T_f
#   x[8]     = T_c
#   x[9]     = T_cold
#   x[10]    = Q_sg (constant-parameter, dx[10]=0)
#   x[11]    = t (dx[11]=1)
@taylorize function rhs_heatup_sd_taylor!(dx, x, p, t)
    rho = KP_HEATUP * (T_STANDBY + RAMP_RATE_CS * x[11] - 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]) - x[10]) / (M_SG * C_C)  # Q_sg now enters!

    dx[10] = zero(x[1])   # Q_sg constant
    dx[11] = one(x[1])    # time

    return nothing
end

function main()
    config_path = length(ARGS) > 0 ? ARGS[1] :
        joinpath(@__DIR__, "..", "..", "configs", "heatup", "with_steam_dump.toml")
    if !isfile(config_path)
        alt = joinpath(@__DIR__, "..", "..", config_path)
        isfile(alt) && (config_path = alt)
    end
    config = TOML.parsefile(config_path)

    println("\n=== Heatup PJ reach with steam dump — config: $(config["name"]) ===")
    println("  $(get(config, "description", ""))")

    e = config["entry"]
    n_lo, n_hi           = e["n_range"]
    T_f_lo, T_f_hi       = e["T_f_range_C"]
    T_c_lo, T_c_hi       = e["T_c_range_C"]
    T_cold_lo, T_cold_hi = e["T_cold_range_C"]
    sd = e["steam_dump"]
    Q_lo = sd["Q_lo_fraction_P0"] * P0
    Q_hi = sd["Q_hi_fraction_P0"] * P0

    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, Q_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, Q_hi, 0.0]

    X0 = Hyperrectangle(low=x_lo, high=x_hi)
    println("  n ∈ [$n_lo, $n_hi], T_c ∈ [$T_c_lo, $T_c_hi] °C")
    println("  Q_sg (steam dump) ∈ [$(Q_lo/1e6), $(Q_hi/1e6)] MW")

    t_cfg = config["tmjets"]
    probes = config["probes"]["horizons_seconds"]

    results = Dict{Float64, Any}()
    for T_probe in probes
        println("\n--- Probe T = $T_probe s ---")
        sys  = BlackBoxContinuousSystem(rhs_heatup_sd_taylor!, 11)
        prob = InitialValueProblem(sys, X0)
        try
            alg = TMJets(orderT=t_cfg["orderT"], orderQ=t_cfg["orderQ"],
                         abstol=t_cfg["abstol"], maxsteps=t_cfg["maxsteps"])
            t0 = time()
            sol = solve(prob; T=Float64(T_probe), alg=alg)
            elapsed = time() - t0
            flow_hr = overapproximate(flowpipe(sol), Hyperrectangle)
            n_sets = length(flow_hr)
            println("  TMJets: $n_sets reach-sets in $(round(elapsed; digits=1)) s")

            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)
            Tco_lo_env = minimum(low(set(R), 9) for R in flow_hr)
            Tco_hi_env = maximum(high(set(R), 9) 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("  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("  T_cold env:     [$(round(Tco_lo_env; digits=2)), $(round(Tco_hi_env; digits=2))] °C")

            # Compare low-T_avg trip (280 °C).
            low_trip_ok = Tc_lo_env >= 280.0
            println("  Low-T_avg trip (T_c ≥ 280): $(low_trip_ok ? "✅" : "× loose")")

            results[T_probe] = (status="OK", Tc=(Tc_lo_env, Tc_hi_env),
                                Tf=(Tf_lo_env, Tf_hi_env),
                                Tcold=(Tco_lo_env, Tco_hi_env),
                                low_trip_ok=low_trip_ok)
        catch err
            println("  FAILED: ", first(sprint(showerror, err), 300))
            results[T_probe] = (status="FAILED",)
            break
        end
    end

    mat_out = joinpath(@__DIR__, "..", "..", "..", "results", config["output"]["result_file"])
    saved = Dict{String, Any}("config_name" => config["name"],
                              "Q_lo" => Q_lo, "Q_hi" => Q_hi)
    for (T_probe, r) in results
        if r.status == "OK"
            pre = "T_$(Int(T_probe))_"
            saved[pre * "Tc_lo"] = r.Tc[1]; saved[pre * "Tc_hi"] = r.Tc[2]
            saved[pre * "Tf_lo"] = r.Tf[1]; saved[pre * "Tf_hi"] = r.Tf[2]
            saved[pre * "Tcold_lo"] = r.Tcold[1]; saved[pre * "Tcold_hi"] = r.Tcold[2]
        end
    end
    matwrite(mat_out, saved)
    println("\nSaved to $mat_out")
end

main()
