#!/usr/bin/env julia
#
# barrier_sos_2d.jl — SOS polynomial barrier on a 2-state projection
# of the operation-mode LQR closed loop.
#
# Proof of concept that SumOfSquares.jl + CSDP can fit a polynomial
# barrier certificate on a reduced model.  If this works, scaling to
# full 10-state is a matter of increasing degree and throughput.
#
# Reduced dynamics: project the LQR closed-loop onto (dn, dT_c), the
# dominant unregulated direction and the primary safety direction.
# A_red, B_red are the 2x2 / 2x1 submatrices corresponding to these
# components (ignoring cross-coupling into the 8 other states, which is
# a modeling simplification but keeps the SOS tractable).
#
# Safety: |dT_c| ≤ 5 K AND |dn| ≤ 0.15 (i.e. 0.85 ≤ n ≤ 1.15).
# Entry:  |dT_c| ≤ 0.1 AND |dn| ≤ 0.01.
# Unsafe focus: dn ≥ +0.15 (high-flux trip; the harder direction
# under LQR because positive-n excursions trip n_high before T_c trips).

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

using LinearAlgebra
using MatrixEquations
using DynamicPolynomials
using SumOfSquares
using CSDP

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

plant = pke_params()

# Full linearization at full-power steady state.
A_full, B_full, B_w_full, _, _, _ = pke_linearize(plant)

# Reduced 2x2: rows/cols (1, 9) — n and T_c.
reduce_idx = [1, 9]
A_red   = A_full[reduce_idx, reduce_idx]
B_red   = B_full[reduce_idx]
B_w_red = B_w_full[reduce_idx]

# LQR on the reduced system.  Light weighting on n, heavy on T_c.
Q_lqr = Diagonal([1.0, 1e2])
R_lqr = 1e6 * ones(1, 1)
X_ric, _, _ = arec(A_red, reshape(B_red, :, 1), R_lqr, Matrix(Q_lqr))
K_red = (R_lqr \ reshape(B_red, 1, :)) * X_ric
A_cl_red = A_red - reshape(B_red, :, 1) * K_red

# Cross-coupling check from dropped states.
cross = A_full[reduce_idx, setdiff(1:10, reduce_idx)]

println("\n=== SOS barrier — 2-state (dn, dT_c) projection of operation LQR ===")
println("  A_cl_red =")
show(stdout, "text/plain", A_cl_red); println()
println("  B_w_red = $B_w_red")
println("  eigenvalues: ", round.(eigvals(A_cl_red); sigdigits=4))
println("  ‖dropped-coupling‖ = $(round(norm(cross); sigdigits=3))")
println()

# --- SOS sets ---
@polyvar x1 x2   # x1 = dn, x2 = dT_c

entry_halfspaces = [
    0.01 - x1,    # dn ≤ 0.01
    x1 + 0.01,    # dn ≥ -0.01
    0.1  - x2,    # dT_c ≤ 0.1
    x2  + 0.1,    # dT_c ≥ -0.1
]

# Unsafe focus: dn ≥ +0.15 (high-flux trip). Asymmetric — n_high trips
# at 1.15 (dn = +0.15), n_low at 0.15 (dn = -0.85), so the +0.15
# direction is the binding one for LQR which has tightly bounded n.
unsafe_halfspaces = [x1 - 0.15]

# --- Solve ---
println("  Solving SOS feasibility (degree-4 B, ε-slack capped at 1.0)...")
result = solve_sos_barrier_2d(A_cl_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("  ✅ ε* = $(round(result.ε; digits=4)) — real certificate.")
    println("  B(x) = $(result.B)")
elseif result.status == MOI.OPTIMAL
    println("  ⚠ ε ≈ 0 — solver returned trivial B ≡ 0. No real barrier")
    println("     at degree 4 with these sets.")
else
    println("  ❌ $(result.status). Try higher degree or relax sets.")
end