using ReachabilityAnalysis
using Plots

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

Using ReachabilityAnalysis.jl with set-based methods (not parameter sweeps!)

Key insight: We'll treat Q1 as an uncertain INPUT, not a state variable.
This avoids the @taylorize issues we had before.
"""

println("=== Two-Loop Reactor: Formal Reachability Analysis ===\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)

println("=== Parameters ===")
println("C_0 = $C_0 %-sec/°F")
println("τ_0 = $τ_0 sec")
println("μ = $μ (fixed)")
println("P_r = $P_r")
println("Q2 = $Q2 (SG2 - KNOWN)")

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

println("\nComputed parameters:")
println("C_r = $C_r")
println("C_sg = $C_sg")
println("W = $W")

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

println("\n=== Initial Conditions ===")
println("T_hot_0 = $T_hot_0 °F")
println("T_cold1_0 = $T_cold1_0 °F")
println("T_cold2_0 = $T_cold2_0 °F")

# Q1 uncertainty range
const Q1_min = 45.0
const Q1_max = 55.0
const Q1_center = (Q1_min + Q1_max) / 2

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

# Define the system using @taylorize
# We'll approximate Q1 uncertainty by adding it as a 4th state with Q̇1 = 0
@taylorize function two_loop_with_Q1!(du, u, p, t)
    T_hot, T_cold1, T_cold2, Q1 = u

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

    return du
end

# Initial set: Small box around initial point for temps, larger for Q1
# Using Hyperrectangle for initial set
X0_low = [T_hot_0 - 0.1, T_cold1_0 - 0.1, T_cold2_0 - 0.1, Q1_min]
X0_high = [T_hot_0 + 0.1, T_cold1_0 + 0.1, T_cold2_0 + 0.1, Q1_max]
X0 = Hyperrectangle(low=X0_low, high=X0_high)

println("\n=== Initial Set ===")
println("X0 = Hyperrectangle")
println("  T_hot ∈ [$(T_hot_0 - 0.1), $(T_hot_0 + 0.1)]")
println("  T_cold1 ∈ [$(T_cold1_0 - 0.1), $(T_cold1_0 + 0.1)]")
println("  T_cold2 ∈ [$(T_cold2_0 - 0.1), $(T_cold2_0 + 0.1)]")
println("  Q1 ∈ [$Q1_min, $Q1_max]")

# Create the initial value problem
prob = @ivp(x' = two_loop_with_Q1!(x), dim: 4, x(0) ∈ X0)

println("\n=== Solving Reachability Problem ===")
println("Using TMJets algorithm (Taylor models)...")

# Solve using TMJets (Taylor model integration)
# This is ACTUAL reachability analysis, not simulation!
sol = solve(prob, T=30.0, alg=TMJets(abstol=1e-10, orderT=7, orderQ=1))

println("✓ Reachability computation complete!")
println("Number of reach sets: $(length(sol))")

# Extract flowpipe projections
println("\n=== Creating Plots ===")

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

# Plot 2: T_cold1 vs time (affected by Q1 uncertainty)
p2 = plot(sol, vars=(0, 2),
         xlabel="Time (s)", ylabel="T_cold1 (°F)",
         title="Cold Leg 1 (SG1 - Uncertain)",
         lw=0, alpha=0.5, color=:blue,
         lab="Reachable set")

# Plot 3: T_cold2 vs time (not affected by Q1)
p3 = plot(sol, vars=(0, 3),
         xlabel="Time (s)", ylabel="T_cold2 (°F)",
         title="Cold Leg 2 (SG2 - Known)",
         lw=0, alpha=0.5, color=:green,
         lab="Reachable set")

# Plot 4: Phase portrait T_hot vs T_cold1
p4 = plot(sol, vars=(2, 1),
         xlabel="T_cold1 (°F)", ylabel="T_hot (°F)",
         title="Phase Portrait: T_hot vs T_cold1",
         lw=0, alpha=0.5, color=:purple,
         lab="Reachable set")

plot_combined = plot(p1, p2, p3, p4, layout=(2, 2), size=(1400, 1000),
                    plot_title="Formal Reachability Analysis: SG1 Uncertainty Q1 ∈ [$Q1_min, $Q1_max]")
savefig(plot_combined, "two_loop_reachability.png")
println("Saved: two_loop_reachability.png")

# Compute T_avg for each reach set and create phase portrait
println("\n=== Computing T_avg Reachability ===")

# For T_avg phase portrait, we need to project to (T_avg, T_hot) space
# T_avg = μ*T_hot + (1-μ)*(T_cold1 + T_cold2)/2
# This requires linear transformation of the flowpipe

# Let's create a detailed phase portrait plot
p_phase = plot(xlabel="T_avg (°F)", ylabel="T_hot (°F)",
              title="Phase Portrait: T_hot vs T_avg\n(Formal Reachability: Q1 ∈ [$Q1_min, $Q1_max])",
              size=(1000, 900), legend=:topright)

# We'll sample the boundary of each reach set
println("Extracting reach set boundaries for T_avg computation...")
for (i, reach_set) in enumerate(sol)
    if i % 10 == 0
        # Sample this reach set
        T_hot_vals = Float64[]
        T_avg_vals = Float64[]

        # Get the set at this time step
        set = reach_set.X

        # Sample vertices and some interior points
        for _ in 1:50
            # Sample a random point from the set
            pt = sample(set)
            T_h = pt[1]
            T_c1 = pt[2]
            T_c2 = pt[3]

            T_avg = μ * T_h + (1 - μ) * (T_c1 + T_c2) / 2

            push!(T_hot_vals, T_h)
            push!(T_avg_vals, T_avg)
        end

        # Plot the cloud of points
        scatter!(p_phase, T_avg_vals, T_hot_vals,
                markersize=2, alpha=0.3, color=:purple, label="")
    end
end

savefig(p_phase, "phase_portrait_reachability.png")
println("Saved: phase_portrait_reachability.png")

println("\n✓ Complete!")
println("\n=== Reachability Analysis Results ===")
println("This is FORMAL reachability analysis using Taylor models,")
println("not simulation-based parameter sweeps!")
println("\nThe reach tubes show ALL possible trajectories for")
println("Q1 ∈ [$Q1_min, $Q1_max] with GUARANTEES.")