using ReachabilityAnalysis
using Plots

"""
Two-Loop Reactor: Temperature Uncertainty Propagation

Simple, clean approach:
- Uncertainty ONLY in initial temperatures: T_hot, T_cold1, T_cold2
- Everything else is FIXED: μ = 0.6, Q1 = 50, Q2 = 50, P_r = 100
- T_avg is derived from the uncertain temperatures

This shows how initial temperature measurement uncertainty propagates.
"""

println("=== Two-Loop Reactor: Initial Temperature Uncertainty ===\n")

# ALL parameters are FIXED (no uncertainty)
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       # FIXED reactor power
const Q1 = 50.0         # FIXED SG1 heat removal
const Q2 = 50.0         # FIXED SG2 heat removal

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

println("=== Fixed Parameters ===")
println("C_0 = $C_0 %-sec/°F")
println("τ_0 = $τ_0 sec")
println("μ = $μ (FIXED)")
println("P_r = $P_r (FIXED)")
println("Q1 = $Q1 (FIXED)")
println("Q2 = $Q2 (FIXED)")
println("\nComputed:")
println("C_r = $C_r")
println("C_sg = $C_sg")
println("W = $W")

# Initial temperature uncertainty
# Realistic: ±2°F measurement uncertainty
const T_hot_nominal = 455.0
const T_cold_nominal = 450.0
const δT = 2.0  # ±2°F uncertainty

println("\n=== Initial Temperature Uncertainty ===")
println("T_hot ∈ [$(T_hot_nominal - δT), $(T_hot_nominal + δT)] °F")
println("T_cold1 ∈ [$(T_cold_nominal - δT), $(T_cold_nominal + δT)] °F")
println("T_cold2 ∈ [$(T_cold_nominal - δT), $(T_cold_nominal + δT)] °F")
println("\nPhysical meaning: ±$δT°F measurement uncertainty")

# Define system (no @taylorize, just plain function)
function two_loop!(dx, x, p, t)
    T_hot = x[1]
    T_cold1 = x[2]
    T_cold2 = x[3]

    # Energy balances with FIXED parameters
    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

    return dx
end

# Try using @taylorize for TMJets
@taylorize function two_loop_taylor!(dx, x, p, t)
    T_hot = x[1]
    T_cold1 = x[2]
    T_cold2 = x[3]

    # Energy balances with FIXED parameters
    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

    return dx
end

# Initial set: Box around nominal temperatures
X0 = Hyperrectangle(
    low=[T_hot_nominal - δT, T_cold_nominal - δT, T_cold_nominal - δT],
    high=[T_hot_nominal + δT, T_cold_nominal + δT, T_cold_nominal + δT]
)

println("\n=== Attempting TMJets (Taylor Models) ===")
println("This is FORMAL reachability analysis...")

# Create IVP with @taylorize version
prob = @ivp(x' = two_loop_taylor!(x), dim: 3, x(0) ∈ X0)

try
    # Try TMJets - this should work since all parameters are constants!
    sol = solve(prob, T=30.0, alg=TMJets(abstol=1e-10, orderT=7, orderQ=2))

    println("✓ TMJets SUCCESS! Computed $(length(sol)) reach sets")
    println("\nThis is FORMAL reachability analysis using Taylor models!")

    # Plot results
    println("\n=== Creating Plots ===")

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

    # Plot 2: T_cold1 reach tube
    p2 = plot(sol, vars=(0, 2),
             xlabel="Time (s)", ylabel="T_cold1 (°F)",
             title="Cold Leg 1",
             lw=0, alpha=0.5, color=:blue,
             lab="Reach tube")

    # Plot 3: T_cold2 reach tube
    p3 = plot(sol, vars=(0, 3),
             xlabel="Time (s)", ylabel="T_cold2 (°F)",
             title="Cold Leg 2",
             lw=0, alpha=0.5, color=:green,
             lab="Reach tube")

    # 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",
             lw=0, alpha=0.5, color=:purple,
             lab="Reach tube")

    plot_combined = plot(p1, p2, p3, p4, layout=(2, 2), size=(1400, 1000),
                        plot_title="Formal Reachability: Initial Temperature Uncertainty ±$(δT)°F")
    savefig(plot_combined, "temp_uncertainty_reachability.png")
    println("Saved: temp_uncertainty_reachability.png")

    # Compute T_avg for each reach set
    println("\n=== Computing T_avg Reach Tube ===")

    times = Float64[]
    T_avg_min = Float64[]
    T_avg_max = Float64[]
    T_hot_vals = Float64[]
    T_cold1_vals = Float64[]
    T_cold2_vals = Float64[]

    for (i, reach_set) in enumerate(sol)
        if i % 5 == 0  # Sample every 5th set
            t = sup(reach_set.t_start)
            set = reach_set.X

            # Sample points from the set
            n_samples = 100
            T_avg_samples = Float64[]

            for _ in 1:n_samples
                pt = sample(set)
                T_h = pt[1]
                T_c1 = pt[2]
                T_c2 = pt[3]

                # Compute T_avg = μ*T_hot + (1-μ)*(T_cold1 + T_cold2)/2
                T_avg = μ * T_h + (1 - μ) * (T_c1 + T_c2) / 2

                push!(T_avg_samples, T_avg)

                if i == length(sol)  # Store final points for phase portrait
                    push!(T_hot_vals, T_h)
                    push!(T_cold1_vals, T_c1)
                    push!(T_cold2_vals, T_c2)
                end
            end

            push!(times, t)
            push!(T_avg_min, minimum(T_avg_samples))
            push!(T_avg_max, maximum(T_avg_samples))
        end
    end

    # Plot T_avg reach tube
    p_tavg = plot(times, T_avg_min, fillrange=T_avg_max,
                 xlabel="Time (s)", ylabel="T_avg (°F)",
                 title="Average Primary Temperature\n(Derived from uncertain T_hot, T_cold1, T_cold2)",
                 fillalpha=0.5, color=:orange, lw=2,
                 label="T_avg reach tube", legend=:topright,
                 size=(1000, 700))
    savefig(p_tavg, "tavg_reachability.png")
    println("Saved: tavg_reachability.png")

    # Phase portrait: T_hot vs T_avg
    if length(T_hot_vals) > 0
        p_phase = scatter(T_avg_samples, T_hot_vals,
                         xlabel="T_avg (°F)", ylabel="T_hot (°F)",
                         title="Phase Portrait: T_hot vs T_avg\n(Final reachable set at t=30s)",
                         markersize=3, alpha=0.5, color=:purple,
                         label="Reachable points", legend=:topright,
                         size=(1000, 900))
        savefig(p_phase, "phase_portrait_tavg.png")
        println("Saved: phase_portrait_tavg.png")
    end

    println("\n✓ Complete!")
    println("\n=== Results ===")
    println("This is FORMAL reachability analysis using Taylor models.")
    println("The reach tubes show ALL possible trajectories starting from")
    println("initial temperature uncertainty: T ∈ [T_nominal ± $(δT)°F]")
    println("\nFinal T_avg range: [$(round(minimum(T_avg_min), digits=1)), $(round(maximum(T_avg_max), digits=1))] °F")
    println("Initial uncertainty ±$(δT)°F → Final spread $(round(maximum(T_avg_max) - minimum(T_avg_min), digits=1))°F")

catch e
    println("❌ TMJets failed: $e")
    println("\n=== Falling back to parameter sweep ===")

    # Fine parameter sweep as backup
    include("two_loop_fine_sweep.jl")
end