using DifferentialEquations
using Plots

"""
Two-Loop Reactor: SG1 Heat Removal Uncertainty

Simple, realistic scenario:
- SG2 efficiency is known: Q2 = 50.0
- SG1 efficiency is uncertain: Q1 ∈ [45, 55] (±10% due to fouling/conditions)
- Everything else is exact!

This shows how uncertainty in ONE component propagates through the system.
"""

println("=== Two-Loop Reactor: SG1 Heat Removal Uncertainty ===\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)")

# Initial conditions (non-equilibrium for interesting dynamics)
const T_hot_0 = 455.0    # Hot leg slightly elevated
const T_cold1_0 = 450.0  # Cold legs at reference
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")

# Two-loop ODEs with Q1 as parameter
function two_loop!(du, u, p, t)
    T_hot, T_cold1, T_cold2 = u
    Q1 = p[1]  # SG1 heat removal (UNCERTAIN)

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

    # 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  # Q1 varies!
    du[3] = (W * (T_hot - T_cold2) - Q2) / C_sg  # Q2 fixed

    return du
end

# Q1 uncertainty range
Q1_values = range(45.0, 55.0, length=40)  # Fine sweep over ±10%

println("\n=== Q1 Uncertainty (SG1 Heat Removal) ===")
println("Q1 range: [45.0, 55.0] (±10% from nominal 50)")
println("Number of cases: $(length(Q1_values))")
println("\nPhysical meaning:")
println("  Q1 = 45: SG1 is 10% less efficient (fouling, poor heat transfer)")
println("  Q1 = 55: SG1 is 10% more efficient (better conditions)")

# Solve for all Q1 values
u0 = [T_hot_0, T_cold1_0, T_cold2_0]
tspan = (0.0, 30.0)
t_save = range(0, 30, length=300)

println("\nSolving...")
all_sols = []
for Q1 in Q1_values
    p = [Q1]
    prob = ODEProblem(two_loop!, u0, tspan, p)
    sol = solve(prob, Tsit5(), saveat=t_save)
    push!(all_sols, sol)
end
println("Done!")

# Extract data
times = all_sols[1].t
all_T_hot = [[sol.u[i][1] for i in 1:length(sol)] for sol in all_sols]
all_T_cold1 = [[sol.u[i][2] for i in 1:length(sol)] for sol in all_sols]
all_T_cold2 = [[sol.u[i][3] for i in 1:length(sol)] for sol in all_sols]

# Compute T_ave
all_T_ave = []
for (i, Q1) in enumerate(Q1_values)
    T_ave = μ * all_T_hot[i] .+ (1 - μ) * (all_T_cold1[i] .+ all_T_cold2[i]) ./ 2
    push!(all_T_ave, T_ave)
end

# Compute reach tube envelopes
T_hot_min = [minimum([all_T_hot[i][j] for i in 1:length(Q1_values)]) for j in 1:length(times)]
T_hot_max = [maximum([all_T_hot[i][j] for i in 1:length(Q1_values)]) for j in 1:length(times)]
T_cold1_min = [minimum([all_T_cold1[i][j] for i in 1:length(Q1_values)]) for j in 1:length(times)]
T_cold1_max = [maximum([all_T_cold1[i][j] for i in 1:length(Q1_values)]) for j in 1:length(times)]
T_ave_min = [minimum([all_T_ave[i][j] for i in 1:length(Q1_values)]) for j in 1:length(times)]
T_ave_max = [maximum([all_T_ave[i][j] for i in 1:length(Q1_values)]) for j in 1:length(times)]

println("\n=== Creating Plots ===")

# Plot 1: T_hot reach tube
p1 = plot(times, T_hot_min, fillrange=T_hot_max,
         xlabel="Time (s)", ylabel="T_hot (°F)",
         title="Hot Leg Temperature",
         fillalpha=0.4, color=:red, lw=2.5,
         label="Reach tube", legend=:topright)

# Plot 2: T_cold1 reach tube (affected by Q1 uncertainty)
p2 = plot(times, T_cold1_min, fillrange=T_cold1_max,
         xlabel="Time (s)", ylabel="T_cold1 (°F)",
         title="Cold Leg 1 (SG1 - Uncertain)",
         fillalpha=0.4, color=:blue, lw=2.5,
         label="Reach tube", legend=:topright)

# Plot 3: T_ave reach tube
p3 = plot(times, T_ave_min, fillrange=T_ave_max,
         xlabel="Time (s)", ylabel="T_avg (°F)",
         title="Average Primary Temperature",
         fillalpha=0.4, color=:orange, lw=2.5,
         label="Reach tube", legend=:topright)

# Plot 4: Phase portrait T_hot vs T_ave
p4 = plot(xlabel="T_avg (°F)", ylabel="T_hot (°F)",
         title="Phase Portrait: T_hot vs T_avg",
         legend=:topright)

# Plot envelope
plot!(p4, [T_ave_min; reverse(T_ave_max)], [T_hot_min; reverse(T_hot_max)],
      seriestype=:shape, fillalpha=0.3, color=:purple,
      label="Reachable region", lw=2.5)

# Mark extremes
plot!(p4, T_ave_min, T_hot_min, color=:blue, lw=2.5, label="Q1 = 45", linestyle=:dash)
plot!(p4, T_ave_max, T_hot_max, color=:red, lw=2.5, label="Q1 = 55", linestyle=:dash)

plot_combined = plot(p1, p2, p3, p4, layout=(2, 2), size=(1400, 1000),
                    plot_title="SG1 Heat Removal Uncertainty: Q1 ∈ [45, 55]")
savefig(plot_combined, "sg1_uncertainty.png")
println("Saved: sg1_uncertainty.png")

# Detailed phase portrait
p_phase = plot(xlabel="T_avg (°F)", ylabel="T_hot (°F)",
              title="Phase Portrait: T_hot vs T_avg\n(SG1 efficiency uncertainty: Q1 ∈ [45, 55])",
              size=(1000, 900), legend=:topright)

# Plot all trajectories
for (i, Q1) in enumerate(Q1_values)
    color_val = (Q1 - minimum(Q1_values)) / (maximum(Q1_values) - minimum(Q1_values))
    plot!(p_phase, all_T_ave[i], all_T_hot[i],
          color=cgrad(:RdBu)[color_val], lw=1.5, alpha=0.3, label="")
end

# Add envelope
plot!(p_phase, [T_ave_min; reverse(T_ave_max)], [T_hot_min; reverse(T_hot_max)],
      seriestype=:shape, fillalpha=0.25, color=:purple,
      label="Reachable region", lw=3)

# Mark extremes
plot!(p_phase, all_T_ave[1], all_T_hot[1],
      color=:blue, lw=3.5, label="Q1 = 45 (low efficiency)", linestyle=:dash)
plot!(p_phase, all_T_ave[end], all_T_hot[end],
      color=:red, lw=3.5, label="Q1 = 55 (high efficiency)", linestyle=:dash)

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

println("\n✓ Complete!")
println("\n=== Results ===")
println("T_hot range: [$(round(minimum(T_hot_min), digits=1)), $(round(maximum(T_hot_max), digits=1))] °F")
println("T_avg range: [$(round(minimum(T_ave_min), digits=1)), $(round(maximum(T_ave_max), digits=1))] °F")
println("T_cold1 range: [$(round(minimum(T_cold1_min), digits=1)), $(round(maximum(T_cold1_max), digits=1))] °F")
println("\nΔT_hot = $(round(maximum(T_hot_max) - minimum(T_hot_min), digits=1)) °F")
println("ΔT_avg = $(round(maximum(T_ave_max) - minimum(T_ave_min), digits=1)) °F")