#!/usr/bin/env julia
#
# barrier_polytopic.jl — polytopic barrier attempt for operation mode.
#
# Idea: instead of the loose quadratic-Lyapunov ellipsoid, use the
# polytope P = inv2_holds ∩ (precursor bounds) and verify
# forward-invariance face-by-face via LP.
#
# Nagumo's theorem for linear systems with bounded disturbance: a
# polytope P = {x : Ax ≤ b} is forward-invariant under dx/dt = A_cl x +
# B_w w iff for each face i of P (where a_i' x = b_i is active),
#   max over (x in P with a_i'x=b_i, w in W) of a_i'(A_cl x + B_w w) ≤ 0.
# This is one LP per face.
#
# The issue with inv2_holds alone: it's unbounded in the 6 precursor
# directions, so the LP is unbounded.  Fix: add precursor bounds
# inferred from the reach-tube envelope (which we already computed in
# reach_operation.jl).  The resulting augmented polytope is bounded
# and the LPs have finite solutions.
#
# Uses JuMP + HiGHS (open-source, ships with no license trouble).

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

using Printf
using LinearAlgebra
using MatrixEquations
using MAT
using JSON
using JuMP
using HiGHS

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", "load_predicates.jl"))

plant = pke_params()
x_op  = pke_initial_conditions(plant)
pred  = load_predicates(plant)

# --- Closed-loop A_cl (LQR) ---
A, B, B_w, _, _, _ = pke_linearize(plant)
Q_lqr = Diagonal([1.0, 1e-3, 1e-3, 1e-3, 1e-3, 1e-3, 1e-3, 1e-2, 1e2, 1.0])
R_lqr = 1e6 * ones(1, 1)
X_ric, _, _ = arec(A, reshape(B, :, 1), R_lqr, Matrix(Q_lqr))
K    = (R_lqr \ reshape(B, 1, :)) * X_ric
A_cl = A - reshape(B, :, 1) * K

# --- inv2_holds halfspaces (from JSON) ---
inv2 = pred.mode_invariants[:inv2_holds]
A_inv2 = inv2.A_poly        # 6 × 10
b_inv2 = inv2.b_poly        # 6

# --- Precursor bounds from reach-tube envelope ---
# Read reach_operation_result.mat; take min/max of X_lo, X_hi on precursors.
reach_mat_path = joinpath(@__DIR__, "..", "..", "..", "reachability",
                          "reach_operation_result.mat")
reach = matread(reach_mat_path)
X_lo = reach["X_lo"]; X_hi = reach["X_hi"]

# Build the augmented polytope: inv2_holds ∪ (C_i in tube bounds).
# Precursor halfspaces: C_i ≤ max(X_hi[i+1, :]), -C_i ≤ -min(X_lo[i+1, :])
# for i = 1..6 (rows 2..7 of state).
A_aug = copy(A_inv2)
b_aug = copy(b_inv2)
labels = copy(inv2.components)

for i in 1:6
    local idx = i + 1
    local C_min = minimum(X_lo[idx, :])
    local C_max = maximum(X_hi[idx, :])
    local row_hi = zeros(1, 10); row_hi[1, idx] = 1.0
    global A_aug = vcat(A_aug, row_hi); global b_aug = vcat(b_aug, C_max)
    push!(labels, "C$(i)_upper")
    local row_lo = zeros(1, 10); row_lo[1, idx] = -1.0
    global A_aug = vcat(A_aug, row_lo); global b_aug = vcat(b_aug, -C_min)
    push!(labels, "C$(i)_lower")
end

# (Skipping tube-based bounds on n, T_f, T_cold — those are REACH
#  envelopes, not forward-invariant limits.  We rely on the inv2_holds
#  halfspaces for n/T_f/T_cold and the precursor tube-bounds above to
#  keep the polytope bounded in all 10 dimensions.)

n_faces = size(A_aug, 1)
println("\n=== Polytopic barrier check ===")
println("  Polytope P = inv2_holds ∩ (precursor + n/T_f/T_cold tube bounds)")
println("  Total faces: $n_faces")
println("  Disturbance: Q_sg ∈ [0.85, 1.00]·P_0")

# --- Disturbance envelope ---
Q_nom = plant.P0
w_lo  = 0.85 * plant.P0 - Q_nom
w_hi  = 0.00 * plant.P0 + (plant.P0 - Q_nom)  # = 0
# Actually the disturbance is Q_sg deviation from nominal.
# Q_sg ∈ [0.85, 1.00]·P0 → w ∈ [-0.15·P0, 0].
w_lo = -0.15 * plant.P0
w_hi = 0.0

# --- Per-face LP check ---
# For face i: max over x, w of a_i' (A_cl x + B_w w) s.t. Ax ≤ b, a_i'x = b_i,
#             w ∈ [w_lo, w_hi].  All in DEVIATION coordinates (dx = x - x_op).
# Need to convert polytope: in absolute coords, P is {x : A_aug x ≤ b_aug}.
# Deviation polytope: P_dev = {dx : A_aug (dx + x_op) ≤ b_aug} = {dx : A_aug dx ≤ b_aug - A_aug x_op}.
b_dev = b_aug .- A_aug * x_op

results = []
for i in 1:n_faces
    a_i = A_aug[i, :]
    rhs_i = b_dev[i]

    # Worst-case disturbance contribution a_i' B_w w over w ∈ [w_lo, w_hi].
    # a_i' B_w is a scalar; max w is w_hi if that scalar > 0 else w_lo.
    aB = a_i' * B_w
    dist_worst = aB > 0 ? aB * w_hi : aB * w_lo

    # LP: maximize a_i' (A_cl dx) + dist_worst
    # subject to A_aug dx ≤ b_dev, a_i' dx = rhs_i.
    model = Model(HiGHS.Optimizer)
    set_silent(model)
    @variable(model, dx[1:10])
    @constraint(model, A_aug * dx .<= b_dev)
    @constraint(model, a_i' * dx == rhs_i)
    grad = (A_cl' * a_i)
    @objective(model, Max, grad' * dx + dist_worst)
    optimize!(model)

    status = termination_status(model)
    if status == MOI.OPTIMAL
        obj = objective_value(model)
        margin = -obj   # forward invariance requires obj ≤ 0
        badge = margin >= 0 ? "✅ forward-invariant" :
                "❌ can escape (a_i'·ẋ = $(round(obj; digits=4)))"
        @printf "  [%-28s]  max a_i'·ẋ = %+.4e  %s\n" labels[i] obj badge
        push!(results, (label=labels[i], obj=obj, status=status))
    else
        @printf "  [%-28s]  LP status: %s\n" labels[i] string(status)
        push!(results, (label=labels[i], obj=NaN, status=status))
    end
end

n_ok = count(r -> isfinite(r.obj) && r.obj <= 1e-10, results)
println("\n=== Summary ===")
println("  Faces forward-invariant:    $n_ok / $n_faces")
println("  Faces that can violate:     $(n_faces - n_ok)")

if n_ok == n_faces
    println("\n  ✅ Polytope P is forward-invariant under A_cl + bounded Q_sg.")
    println("     Polytopic barrier B(x) = max_i (a_i'(x - x_op) - b_dev_i) satisfies")
    println("     B(x₀) ≤ 0 on P, and dB/dt ≤ 0 on {B = 0}.")
else
    println("\n  ❌ Some faces can be crossed under A_cl; P as constructed is not a")
    println("     valid barrier.  Expected: safety halfspaces + tube-envelope bounds")
    println("     together form a set MUCH larger than what LQR can contract to, so")
    println("     the worst-case point on a face can have derivative outward.")
    println()
    println("     The right approach is the Blanchini pre-image algorithm:")
    println("       P₀ = safety polytope (inv2_holds + precursor bounds)")
    println("       P_{k+1} = P_k ∩ {x : A_cl x + B_w w ∈ P_k for all w ∈ W}")
    println("     iterating until fixed point or timeout.  The fixed point is the")
    println("     maximal robustly controllable invariant set inside the safety polytope.")
    println("     Each iteration adds faces; polytope grows combinatorially in size.")
    println("     Tools: Polyhedra.jl + CDDLib for the set operations, HiGHS for LPs.")
    println()
    println("     Rough effort estimate: 2-3 days focused to get a working implementation")
    println("     + a thesis-defensible result on operation mode.  Deferred for now.")
end