#!/usr/bin/env julia
#
# reach_operation.jl — Julia port of reachability/reach_operation.m.
#
# Linear reach tube for operation-mode LQR closed-loop.  Per-halfspace
# margin check against inv2_holds (from predicates.json, single source
# of truth).  Figures saved to docs/figures/.
#
# Soundness caveat: this reach tube is a sound over-approximation of the
# LINEARIZED closed-loop system's reach set, not of the nonlinear plant.
# Linear-reach results here are an approximate-but-useful gut check.

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

using Printf
using LinearAlgebra
using MatrixEquations
using Plots
using JSON
using MAT

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"))
include(joinpath(@__DIR__, "..", "src", "reach_linear.jl"))

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

# --- Closed-loop linearization + LQR gain ---
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

eigs_cl = eigvals(A_cl)
println("\n=== Closed-loop spectrum (A - BK) ===")
@printf "  max Re(eig) = %.3e\n"  maximum(real.(eigs_cl))
@printf "  min Re(eig) = %.3e\n"  minimum(real.(eigs_cl))
@assert all(real.(eigs_cl) .< -1e-8) "A_cl not Hurwitz"

# --- Entry box ---
delta_entry = [0.01 * x_op[1];
               0.001 .* abs.(x_op[2:7]);
               0.1; 0.1; 0.1]

# --- Disturbance envelope ---
Q_nom = plant.P0
Q_min = 0.85 * plant.P0
Q_max = 1.00 * plant.P0
dQ_lo = Q_min - Q_nom
dQ_hi = Q_max - Q_nom

# --- Reach in deviation coordinates ---
tspan = (0.0, 600.0)
dt    = 0.5
T, R_lo, R_hi, Cnom = reach_linear(A_cl, B_w, zeros(10), delta_entry,
                                    dQ_lo, dQ_hi, tspan, dt)

# Translate back to absolute coordinates.
X_lo  = R_lo .+ x_op
X_hi  = R_hi .+ x_op
X_nom = Cnom .+ x_op

# --- Safety check: inv2_holds halfspace-by-halfspace ---
inv2 = pred.mode_invariants[:inv2_holds]
A_inv = inv2.A_poly
b_inv = inv2.b_poly
comps = inv2.components

println("\n=== Operation-mode reach vs inv2_holds safety limits ===")
for k in 1:size(A_inv, 1)
    a = A_inv[k, :]
    # Envelope max of a'*x: take Xhi where a>0, Xlo where a<0.
    env_for_max = (a .> 0) .* X_hi .+ (a .< 0) .* X_lo
    # Small correction: zero coeffs have no preference (both values equal).
    zero_mask = (a .== 0)
    env_for_max[zero_mask, :] .= X_hi[zero_mask, :]
    max_ax = maximum(a' * env_for_max)
    margin = b_inv[k] - max_ax
    status = margin < 0 ? "*** VIOLATED ***" : "OK"
    @printf "  [%-22s]  a'x <= %8.3f  |  max a'x = %8.3f  |  margin = %+8.3f  %s\n" comps[k] b_inv[k] max_ax margin status
end

# --- Per-state reach-set width diagnostic ---
state_names = ["n","C1","C2","C3","C4","C5","C6","T_f","T_c","T_cold"]
println("\n=== Reach-set width at t=0 vs t=T_final ===")
@printf "  %-7s  %-14s  %-14s  %-8s\n"  "state"  "init halfwidth"  "final halfwidth"  "ratio"
for i in 1:10
    hi = 0.5 * (R_hi[i, 1]   - R_lo[i, 1])
    hf = 0.5 * (R_hi[i, end] - R_lo[i, end])
    @printf "  %-7s  %.4e      %.4e      %.2f\n"  state_names[i]  hi  hf  hf/max(hi, eps())
end

# --- Plots: T_c reach tube, two views ---
CtoF(T) = T * 9/5 + 32
delta_safe_Tc = pred.constants.t_avg_in_range_halfwidth_C
figdir = joinpath(@__DIR__, "..", "..", "docs", "figures")
isdir(figdir) || mkpath(figdir)

p_safety = plot(T, CtoF.(X_nom[9, :]), lw=1.2, color=:red,
                label="nominal",
                xlabel="Time [s]", ylabel="T_avg [F]",
                title="Safety-band view")
plot!(p_safety, T, CtoF.(X_hi[9, :]), fillrange=CtoF.(X_lo[9, :]),
      fillalpha=0.3, color=:red, linealpha=0.0,
      label="reach tube")
hline!(p_safety, [CtoF(plant.T_c0 + delta_safe_Tc),
                  CtoF(plant.T_c0 - delta_safe_Tc)],
       ls=:dash, color=:black, label="safety +/- $(round(delta_safe_Tc; digits=2)) C")

dev_lo = X_lo[9, :] .- plant.T_c0
dev_hi = X_hi[9, :] .- plant.T_c0
max_dev = maximum(abs, [dev_lo; dev_hi])
p_zoom = plot(T, X_nom[9, :] .- plant.T_c0, lw=1.2, color=:red,
              label="nominal",
              xlabel="Time [s]", ylabel="T_avg - T_c0 [K]",
              title=@sprintf("Zoomed: max |dT_c| = %.3e K", max_dev))
plot!(p_zoom, T, dev_hi, fillrange=dev_lo, fillalpha=0.3,
      color=:red, linealpha=0.0, label="reach tube")
hline!(p_zoom, [0.0], ls=:dot, color=:black, label=nothing)

fig = plot(p_safety, p_zoom, layout=(1, 2), size=(1400, 500),
           plot_title=@sprintf("Operation-mode reach tube, LQR, Q_sg in [%.0f%%, %.0f%%] P0",
                               100*Q_min/Q_nom, 100*Q_max/Q_nom))
savefig(fig, joinpath(figdir, "reach_operation_Tc.png"))

# n tube
fig_n = plot(T, X_nom[1, :], lw=1.2, color=:blue, label="nominal",
             xlabel="Time [s]", ylabel="n",
             title="Operation mode reach tube on normalized power")
plot!(fig_n, T, X_hi[1, :], fillrange=X_lo[1, :], fillalpha=0.3,
      color=:blue, linealpha=0.0, label="reach tube")
savefig(fig_n, joinpath(figdir, "reach_operation_n.png"))

# --- Save result ---
matfile = joinpath(@__DIR__, "..", "..", "reachability", "reach_operation_result.mat")
matwrite(matfile, Dict("T" => collect(T), "R_lo" => R_lo, "R_hi" => R_hi,
                       "center" => Cnom, "X_lo" => X_lo, "X_hi" => X_hi,
                       "X_nom" => X_nom, "K" => Matrix(K), "A_cl" => A_cl,
                       "delta_entry" => delta_entry,
                       "Q_min" => Q_min, "Q_max" => Q_max,
                       "delta_safe_Tc" => delta_safe_Tc))
println("\nSaved reach result to $matfile")