#!/usr/bin/env julia
#
# barrier_sos_2d.jl — SOS polynomial barrier on a 2-state projection.
#
# Proof of concept that SumOfSquares.jl + CSDP can fit a polynomial
# barrier certificate on a reduced version of the operation-mode
# closed-loop.  If this works, scaling to full 10-state is a matter
# of increasing degree and throughput.
#
# Reduced dynamics: project the LQR closed-loop onto (dT_c, dn), the
# primary safety direction and the dominant unregulated direction.
# A_red, B_w_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.
# Disturbance: Q_sg deviation |dw| ≤ 0.15·P0.
#
# Barrier specification (Prajna-Jadbabaie):
#     B(x) ≤ 0 on X_entry
#     B(x) ≥ 0 on X_unsafe (= complement of safety)
#     ∂B/∂x · f(x) ≤ 0 on {B(x) = 0}  (for all w in W)
# Using SOS multipliers σ_i(x), w-dependence via lossless-disturbance bound.

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"))

plant = pke_params()
x_op  = pke_initial_conditions(plant)

# Full linearization.
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

println("\n=== SOS barrier attempt — 2-state (n, T_c) projection ===")
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()

# --- SOS formulation ---
# dx = [dn; dTc] = [x[1]; x[2]] in polynomial variables.
@polyvar x1 x2

# Dynamics with worst-case constant w:
w_bar = 0.15 * plant.P0
# Split disturbance into its mid + extreme, handle as bounded constant.
# For the Lie derivative check we use the WORST-CASE w that maximizes
# the outward velocity.  Since B_w_red is a known 2-vector and ∂B/∂x
# is polynomial in x, the max-over-w is achieved at w ∈ {-w_bar, +w_bar}.
# Defer that max — check both worst cases separately.

f_nom = A_cl_red * [x1; x2]            # 2-vector of polynomials in x1, x2

# Safety set as intersection of halfspaces g_i ≥ 0:
#   g1 = 5 - x2        (dT_c ≤ 5)
#   g2 = x2 + 5        (dT_c ≥ -5)
#   g3 = 0.15 - x1     (dn ≤ 0.15)
#   g4 = x1 + 0.15     (dn ≥ -0.15)
# Unsafe set = complement; for SOS we use the Putinar formulation where
# B ≥ 0 on unsafe.  With multiple unsafe regions (each =complement of
# one safety halfspace) we'd need one constraint per unsafe region.
# Simpler: pick one unsafe halfspace to focus on — say n >= 1.15
# (high-flux trip).  g_u1 = x1 - 0.15.

# Entry set:
#   g_e1 = 0.1 - x2; g_e2 = x2 + 0.1; g_e3 = 0.01 - x1; g_e4 = x1 + 0.01.

g_s1 = 5.0 - x2
g_s2 = x2 + 5.0
g_s3 = 0.15 - x1
g_s4 = x1 + 0.15
g_u_high = x1 - 0.15    # unsafe when n > 1.15 (dn > 0.15)
g_u_low  = -0.15 - x1   # unsafe when n < 0.85 (dn < -0.15)
g_e1 = 0.1 - x2
g_e2 = x2 + 0.1
g_e3 = 0.01 - x1
g_e4 = x1 + 0.01

# --- Build the SOS program ---
solver = optimizer_with_attributes(CSDP.Optimizer, "printlevel" => 0)
model = SOSModel(solver)

# Barrier polynomial, degree 4.
monos_B = monomials([x1, x2], 0:4)
@variable(model, B_poly, Poly(monos_B))

# SOS multipliers for each set constraint, degree 2.
monos_σ = monomials([x1, x2], 0:2)

# (1) B ≤ 0 on X_entry: -B - Σᵢ σ_eᵢ · g_eᵢ is SOS.
@variable(model, σ_e1, SOSPoly(monos_σ))
@variable(model, σ_e2, SOSPoly(monos_σ))
@variable(model, σ_e3, SOSPoly(monos_σ))
@variable(model, σ_e4, SOSPoly(monos_σ))
@constraint(model, -B_poly - σ_e1*g_e1 - σ_e2*g_e2 - σ_e3*g_e3 - σ_e4*g_e4 in SOSCone())

# (2) B ≥ 0 on X_unsafe (using the "high" unsafe region).  Include safety
# constraints so we stay inside the relevant half:
#   B - σ_u_high · g_u_high - σ_u_s2 · g_s2 - σ_u_s3 · (-1) is SOS (dummy)
# Actually: unsafe-high = {x1 ≥ 0.15} alone (unconstrained in x2).
# Simplest form:
@variable(model, σ_u, SOSPoly(monos_σ))
@constraint(model, B_poly - σ_u * g_u_high in SOSCone())

# (3) Lie derivative:  ∇B · f ≤ 0 EVERYWHERE (not just on B=0 boundary).
# Stronger than needed, but keeps the SDP convex.  The bilinear
# Putinar form -(∇B·f) - σ_b·B ≥ SOS requires iterative BMI methods;
# we skip that for this first attempt and use the stronger "global
# decrease" condition.  If the Hurwitz system admits a quadratic B
# this should still be solvable.
dB_dx = [differentiate(B_poly, x1), differentiate(B_poly, x2)]
# B_w_red is [0, 0] in this projection (Q_sg doesn't directly couple
# into n or T_c in the linearization), so the disturbance term drops
# out and the Lie-derivative condition simplifies.
f_tot = A_cl_red * [x1; x2]
lie = dB_dx[1] * f_tot[1] + dB_dx[2] * f_tot[2]
@constraint(model, -lie in SOSCone())

# Feasibility problem — no objective needed.  Any B that satisfies the
# three SOS constraints is a valid barrier.

println("  Solving SOS program (CSDP)…")
optimize!(model)
status = termination_status(model)
println("  Status: $status")
if status == MOI.OPTIMAL
    println("  ✅ SOS barrier found.")
    println("  B(x) = ", round(value(B_poly); digits=4))
elseif status == MOI.INFEASIBLE
    println("  ❌ SOS program infeasible — no degree-4 polynomial B exists")
    println("     with the given sets and dynamics.  Try higher degree,")
    println("     larger X_unsafe margin, or different formulation.")
else
    println("  ⚠ Solver stopped with: $status")
end