#!/usr/bin/env julia
#
# barrier_sos_2d_shutdown.jl — SOS polynomial barrier on the hot-standby
# (q_shutdown) closed-loop, on a 2-state thermal projection.
#
# Companion to barrier_sos_2d.jl (operation-mode LQR projection).
#
# Hot-standby controller: u = -5*beta (constant rod insertion). With
# Q_sg = 0 (no SG load) and small n, the closed loop has a thermal
# equilibrium where rod-induced negative reactivity balances temperature
# feedback and decay-heat balances are negligible. We:
#   1. Find that equilibrium by long-horizon simulation.
#   2. Linearize there.
#   3. Reduce to (T_c, T_cold) — the slow safety-relevant thermal modes.
#      n is decoupled at low power and not the safety driver in this mode.
#   4. Build a degree-4 SOS barrier on the reduced closed-loop.
#
# Safety set (deviation from equilibrium):
#   |dT_c|    <= 10 K
#   |dT_cold| <= 15 K
# Entry set (X_entry from predicates.json::q_shutdown, recentered on x_eq):
#   |dT_c|    <= 5 K
#   |dT_cold| <= 5 K
#
# X_unsafe-high focus: dT_c >= 10 (over-warming → leaving hot-standby on
# the wrong side; would imply the shutdown controller cannot hold the
# plant cool enough).

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

using Printf
using LinearAlgebra
using OrdinaryDiffEq
using DynamicPolynomials
using SumOfSquares
using CSDP
using JSON

include(joinpath(@__DIR__, "..", "..", "src", "pke_params.jl"))
include(joinpath(@__DIR__, "..", "..", "src", "pke_states.jl"))
include(joinpath(@__DIR__, "..", "..", "src", "pke_th_rhs.jl"))
include(joinpath(@__DIR__, "..", "..", "src", "pke_linearize.jl"))
include(joinpath(@__DIR__, "..", "..", "src", "pke_solver.jl"))
include(joinpath(@__DIR__, "..", "..", "src", "sos_barrier.jl"))
include(joinpath(@__DIR__, "..", "..", "controllers", "controllers.jl"))

plant = pke_params()

pred_path = joinpath(@__DIR__, "..", "..", "..", "reachability", "predicates.json")
predicates = JSON.parsefile(pred_path)
states = pke_states(plant; predicates=predicates)
println("T_standby = $(round(states.T_standby; digits=2)) °C")

# --- (1) Simulate to quasi-equilibrium ---
println("\n=== (1) Finding shutdown quasi-equilibrium ===")
Q_shut = t -> 0.0
t_sim, X_sim, U_sim = pke_solver(plant, Q_shut, ctrl_shutdown, nothing,
                                  (0.0, 50000.0); x0=states.shutdown, verbose=false)
x_eq = X_sim[end, :]
@printf "  After %.0f s: n=%.3e  T_f=%.2f  T_c=%.2f  T_cold=%.2f\n" t_sim[end] x_eq[1] x_eq[8] x_eq[9] x_eq[10]
@printf "  Final-step rate ‖dx/dt‖ ≈ %.3e (should be ~0)\n" norm(
    pke_th_rhs(x_eq, t_sim[end], plant, Q_shut, U_sim[end]))

# --- (2) Linearize at x_eq with u_star = U_SHUTDOWN, Q_star = 0 ---
println("\n=== (2) Linearizing at x_eq ===")
u_eq = -5.0 * plant.beta
A_full, B_full, B_w_full, _, _, _ = pke_linearize(plant;
    x_star=x_eq, u_star=u_eq, Q_star=0.0)

# u is constant → closed-loop A_cl = A_full (no state feedback term).
A_cl = A_full

# --- (3) Reduce to (T_c, T_cold) = state indices 9, 10 ---
reduce_idx = [9, 10]
A_red = A_cl[reduce_idx, reduce_idx]
cross = A_cl[reduce_idx, setdiff(1:10, reduce_idx)]
println("  A_red =")
show(stdout, "text/plain", A_red); println()
println("  eigenvalues: ", round.(eigvals(A_red); sigdigits=4))
println("  ‖dropped-coupling‖ = $(round(norm(cross); sigdigits=3))")

# --- (4) SOS feasibility ---
println("\n=== (4) SOS barrier (degree-4, ε-slack capped at 1.0) ===")
@polyvar x1 x2   # x1 = δT_c, x2 = δT_cold

entry_halfspaces = [
    5.0 - x1,
    x1 + 5.0,
    5.0 - x2,
    x2 + 5.0,
]

# Unsafe-high focus: dT_c ≥ +10. Asymmetric — over-warming is the harder
# case for a shutdown controller (rods already maxed in negative; can't
# add more negative reactivity to compensate).
unsafe_halfspaces = [x1 - 10.0]

result = solve_sos_barrier_2d(A_red, (x1, x2),
                              entry_halfspaces, unsafe_halfspaces;
                              barrier_degree=4, multiplier_degree=2,
                              eps_cap=1.0)

println("  Status: $(result.status)")
if result.status == MOI.OPTIMAL && result.ε > 1e-8
    println(@sprintf "  ✅ ε* = %.4e — real certificate." result.ε)
    println("  B(x) = $(result.B)")
elseif result.status == MOI.OPTIMAL
    println("  ⚠ ε ≈ 0 — solver returned trivial B ≡ 0. No real barrier")
    println("     at degree 4 for these sets.")
else
    println("  ❌ $(result.status). Try higher degree or 3-D extension")
    println("     (include T_f, since dropped-coupling is non-trivial).")
end

println("\n=== Equilibrium summary ===")
@printf "  x_eq = (n=%.3e, T_f=%.3f, T_c=%.3f, T_cold=%.3f)\n" x_eq[1] x_eq[8] x_eq[9] x_eq[10]
@printf "  u_eq = -5*beta = %.4f\n" u_eq
@printf "  rho_eq = u_eq + alpha_f*(T_f-T_f0) + alpha_c*(T_c-T_c0) = %.4f\n" (u_eq + plant.alpha_f*(x_eq[8] - plant.T_f0) + plant.alpha_c*(x_eq[9] - plant.T_c0))