#!/usr/bin/env julia
#
# reach_heatup_pj.jl — nonlinear reach on heatup, prompt-jump model.
#
# Reduced from 10-state to 9-state (n is algebraic). Removes the Λ⁻¹
# stiffness that capped the full-state reach at ~10 s.
#
# State (10D with augmented time):
#   x[1..6] = C_1..C_6   (delayed-neutron precursors)
#   x[7]    = T_f
#   x[8]    = T_c
#   x[9]    = T_cold
#   x[10]   = t (augmented time, dt/dt = 1)
#
# n is algebraic:  n = Λ·Σ λ_i C_i / (β - ρ),  ρ = K_p·(T_ref - T_c).
#
# Configuration-driven: pass a TOML config path as the first CLI arg,
# or omit for the baseline config.
#
#   julia --project=. scripts/reach/reach_heatup_pj.jl                  # baseline
#   julia --project=. scripts/reach/reach_heatup_pj.jl configs/heatup/tight.toml
#
# Configs live in code/configs/heatup/*.toml. See baseline.toml for
# the full schema.

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

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

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

# --- Taylorized heatup PJ RHS ---
@taylorize function rhs_heatup_pj_taylor!(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

# --- Config loader ---
function load_config(config_path)
    if isfile(config_path)
        return TOML.parsefile(config_path)
    else
        error("Config file not found: $config_path")
    end
end

function build_entry_box(config)
    if get(config, "use_predicates_entry", false)
        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"]
    else
        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"]
    end

    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]

    return Hyperrectangle(low=x_lo, high=x_hi),
           (n_lo=n_lo, n_hi=n_hi,
            T_f_lo=T_f_lo, T_f_hi=T_f_hi,
            T_c_lo=T_c_lo, T_c_hi=T_c_hi,
            T_cold_lo=T_cold_lo, T_cold_hi=T_cold_hi)
end

# --- Per-step envelope extraction ---
function extract_envelopes(flow_hr)
    n_steps = length(flow_hr)
    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)
        t_hi_here = high(s, 10)
        t_lo_here = low(s, 10)
        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
    return (; t_arr,
              Tc_lo_ts, Tc_hi_ts, Tf_lo_ts, Tf_hi_ts,
              Tco_lo_ts, Tco_hi_ts, n_lo_ts, n_hi_ts,
              rho_lo_ts, rho_hi_ts)
end

# --- Main ---
function main()
    default_config = joinpath(@__DIR__, "..", "..", "configs", "heatup", "baseline.toml")
    config_path = length(ARGS) > 0 ? ARGS[1] : default_config
    # Allow a config path relative to repo root or code/.
    if !isfile(config_path)
        alt = joinpath(@__DIR__, "..", "..", config_path)
        isfile(alt) && (config_path = alt)
    end
    config = load_config(config_path)

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

    X0, entry_info = build_entry_box(config)
    println("  X_entry: n ∈ [$(entry_info.n_lo), $(entry_info.n_hi)], " *
            "T_c ∈ [$(entry_info.T_c_lo), $(entry_info.T_c_hi)] °C")

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

    results = Dict{Float64, Any}()

    for T_probe in probes
        println("\n--- Probe T = $T_probe s ($(round(T_probe/60; digits=1)) min) ---")
        sys  = BlackBoxContinuousSystem(rhs_heatup_pj_taylor!, 10)
        prob = InitialValueProblem(sys, X0)
        try
            alg = TMJets(
                orderT   = tmjets_cfg["orderT"],
                orderQ   = tmjets_cfg["orderQ"],
                abstol   = tmjets_cfg["abstol"],
                maxsteps = tmjets_cfg["maxsteps"],
            )
            t_start = time()
            sol = solve(prob; T=Float64(T_probe), alg=alg)
            elapsed = time() - t_start
            flow = flowpipe(sol)
            n_sets = length(flow)
            println("  TMJets: $n_sets reach-sets, wall $(round(elapsed; digits=1)) s")

            flow_hr = overapproximate(flow, Hyperrectangle)
            env = extract_envelopes(flow_hr)

            println("  n envelope:   [$(round(minimum(env.n_lo_ts); sigdigits=4)), $(round(maximum(env.n_hi_ts); sigdigits=4))]")
            println("  T_c envelope: [$(round(minimum(env.Tc_lo_ts); digits=2)), $(round(maximum(env.Tc_hi_ts); digits=2))] °C")
            println("  T_f envelope: [$(round(minimum(env.Tf_lo_ts); digits=2)), $(round(maximum(env.Tf_hi_ts); digits=2))] °C")
            println("  T_cold env:   [$(round(minimum(env.Tco_lo_ts); digits=2)), $(round(maximum(env.Tco_hi_ts); digits=2))] °C")
            println("  rho env:      [$(round(minimum(env.rho_lo_ts); sigdigits=4)), $(round(maximum(env.rho_hi_ts); sigdigits=4))]")

            results[T_probe] = (status="OK", n_sets=n_sets, elapsed=elapsed, env=env)
        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 probes
        haskey(results, T_probe) || continue
        r = results[T_probe]
        if r.status == "OK"
            println("  T = $(T_probe) s: OK, $(r.n_sets) sets, $(round(r.elapsed; digits=1))s wall")
        else
            println("  T = $(T_probe) s: FAILED")
        end
    end

    # --- Save ---
    if get(config, "output", Dict()) |> (o -> get(o, "save_per_step", false))
        result_file = config["output"]["result_file"]
        mat_out = joinpath(@__DIR__, "..", "..", "..", "results", result_file)
        saved = Dict{String, Any}(
            "config_name"   => config["name"],
            "probe_horizons" => collect(probes),
            "beta"          => BETA,
            "Kp"            => KP_HEATUP,
            "T_c0"          => T_C0,
            "T_cold0"       => T_COLD0,
            "T_standby"     => T_STANDBY,
        )
        for T_probe in probes
            haskey(results, T_probe) || continue
            r = results[T_probe]
            r.status == "OK" || continue
            pre = "T_$(Int(T_probe))_"
            env = r.env
            saved[pre * "t_arr"]     = env.t_arr
            saved[pre * "Tc_lo_ts"]  = env.Tc_lo_ts;  saved[pre * "Tc_hi_ts"]  = env.Tc_hi_ts
            saved[pre * "Tf_lo_ts"]  = env.Tf_lo_ts;  saved[pre * "Tf_hi_ts"]  = env.Tf_hi_ts
            saved[pre * "Tco_lo_ts"] = env.Tco_lo_ts; saved[pre * "Tco_hi_ts"] = env.Tco_hi_ts
            saved[pre * "n_lo_ts"]   = env.n_lo_ts;   saved[pre * "n_hi_ts"]   = env.n_hi_ts
            saved[pre * "rho_lo_ts"] = env.rho_lo_ts; saved[pre * "rho_hi_ts"] = env.rho_hi_ts
        end
        matwrite(mat_out, saved)
        println("\nSaved per-step envelopes to $mat_out")
    end
end

main()