using ReachabilityAnalysis
using Plots

"""
Two-Loop Reactor: Formal Reachability Analysis with SG1 Efficiency Uncertainty

Attempt #2: Using LGG09 algorithm which doesn't require @taylorize
This uses zonotope overapproximation instead of Taylor models.
"""

println("=== Two-Loop Reactor: Formal Reachability (LGG09) ===\n")

# Constants from Python script
const C_0 = 33.33       # Base Heat Capacity [%-sec/°F]
const τ_0 = 0.75 * C_0  # Base Time Constant [sec]
const μ = 0.6           # Fixed reactor water mass fraction
const P_r = 100.0       # Reactor power
const Q2 = 50.0         # SG2 heat removal (KNOWN)

# System parameters
C_r = μ * C_0
C_sg = (1 - μ) * C_0 / 2
W = C_0 / (2 * τ_0)

println("=== Parameters ===")
println("C_0 = $C_0, τ_0 = $τ_0, μ = $μ")
println("C_r = $C_r, C_sg = $C_sg, W = $W")
println("P_r = $P_r, Q2 = $Q2")

# Initial conditions (non-equilibrium)
const T_hot_0 = 455.0
const T_cold1_0 = 450.0
const T_cold2_0 = 450.0

# Q1 uncertainty range
const Q1_min = 45.0
const Q1_max = 55.0

println("\n=== Q1 Uncertainty ===")
println("Q1 ∈ [$Q1_min, $Q1_max]")

# Define system WITHOUT @taylorize (for LGG09 algorithm)
# Treat Q1 as 4th state with Q̇1 = 0
function two_loop_with_Q1!(dx, x, p, t)
    T_hot = x[1]
    T_cold1 = x[2]
    T_cold2 = x[3]
    Q1 = x[4]

    # Energy balances
    dx[1] = (P_r - W * (T_hot - T_cold1) - W * (T_hot - T_cold2)) / C_r
    dx[2] = (W * (T_hot - T_cold1) - Q1) / C_sg
    dx[3] = (W * (T_hot - T_cold2) - Q2) / C_sg
    dx[4] = 0.0  # Q1 is constant uncertain parameter

    return dx
end

# Initial set
X0 = Hyperrectangle(
    low=[T_hot_0 - 0.1, T_cold1_0 - 0.1, T_cold2_0 - 0.1, Q1_min],
    high=[T_hot_0 + 0.1, T_cold1_0 + 0.1, T_cold2_0 + 0.1, Q1_max]
)

println("\n=== Initial Set ===")
println("X0 = Hyperrectangle(4D)")

# Create IVP
prob = @ivp(x' = two_loop_with_Q1!(x), dim: 4, x(0) ∈ X0)

println("\n=== Solving with LGG09 (Zonotope) ===")
println("This uses zonotope overapproximation...")

# Try LGG09 algorithm (doesn't require @taylorize)
try
    sol = solve(prob, T=30.0, alg=LGG09(δ=0.1))

    println("✓ Success! Computed $(length(sol)) reach sets")

    # Plot results
    p1 = plot(sol, vars=(0, 1),
             xlabel="Time (s)", ylabel="T_hot (°F)",
             title="Hot Leg Temperature",
             lw=0.5, alpha=0.6, color=:red, lab="Reach sets")

    p2 = plot(sol, vars=(0, 2),
             xlabel="Time (s)", ylabel="T_cold1 (°F)",
             title="Cold Leg 1 (SG1 - Uncertain)",
             lw=0.5, alpha=0.6, color=:blue, lab="Reach sets")

    p3 = plot(sol, vars=(0, 3),
             xlabel="Time (s)", ylabel="T_cold2 (°F)",
             title="Cold Leg 2 (SG2 - Known)",
             lw=0.5, alpha=0.6, color=:green, lab="Reach sets")

    p4 = plot(sol, vars=(2, 1),
             xlabel="T_cold1 (°F)", ylabel="T_hot (°F)",
             title="Phase Portrait",
             lw=0.5, alpha=0.6, color=:purple, lab="Reach sets")

    plot_combined = plot(p1, p2, p3, p4, layout=(2, 2), size=(1400, 1000),
                        plot_title="Formal Reachability (LGG09): Q1 ∈ [$Q1_min, $Q1_max]")
    savefig(plot_combined, "two_loop_reachability_lgg09.png")
    println("\nSaved: two_loop_reachability_lgg09.png")

catch e
    println("❌ LGG09 failed: $e")
    println("\nTrying GLGM06 instead...")

    try
        sol = solve(prob, T=30.0, alg=GLGM06(δ=0.1))

        println("✓ Success with GLGM06! Computed $(length(sol)) reach sets")

        # Plot results
        p1 = plot(sol, vars=(0, 1),
                 xlabel="Time (s)", ylabel="T_hot (°F)",
                 title="Hot Leg Temperature",
                 lw=0.5, alpha=0.6, color=:red, lab="Reach sets")

        p2 = plot(sol, vars=(0, 2),
                 xlabel="Time (s)", ylabel="T_cold1 (°F)",
                 title="Cold Leg 1 (SG1 - Uncertain)",
                 lw=0.5, alpha=0.6, color=:blue, lab="Reach sets")

        p3 = plot(sol, vars=(0, 3),
                 xlabel="Time (s)", ylabel="T_cold2 (°F)",
                 title="Cold Leg 2 (SG2 - Known)",
                 lw=0.5, alpha=0.6, color=:green, lab="Reach sets")

        p4 = plot(sol, vars=(2, 1),
                 xlabel="T_cold1 (°F)", ylabel="T_hot (°F)",
                 title="Phase Portrait",
                 lw=0.5, alpha=0.6, color=:purple, lab="Reach sets")

        plot_combined = plot(p1, p2, p3, p4, layout=(2, 2), size=(1400, 1000),
                            plot_title="Formal Reachability (GLGM06): Q1 ∈ [$Q1_min, $Q1_max]")
        savefig(plot_combined, "two_loop_reachability_glgm06.png")
        println("\nSaved: two_loop_reachability_glgm06.png")

    catch e2
        println("❌ GLGM06 also failed: $e2")
        println("\n=== Both algorithms failed ===")
        println("The system may be too nonlinear for these approximation methods.")
        println("\nParameter sweep (our previous approach) may be the most practical")
        println("method for this system, even if it's not formally reachability analysis.")
    end
end

println("\n✓ Complete!")