using DifferentialEquations
using Plots

"""
Two-Loop Reactor System - FINE Parameter Sweep over μ
Using ONLY constants from Python script
"""

println("=== Two-Loop Reactor: Fine μ Sweep ===\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]

# Initial conditions from Python script
const T_initial = 450.0  # All temperatures start at 450°F

# Power levels from Python script
const P_r = 100.0   # Reactor power
const Q_sg = 50.0   # Heat removal per steam generator (equal loading)

println("=== Parameters (from Python script) ===")
println("C_0 = $C_0 %-sec/°F")
println("τ_0 = $τ_0 sec")
println("P_r = $P_r")
println("Q_sg = $Q_sg per SG")

# Two-loop system ODEs
function two_loop!(du, u, p, t)
    T_hot, T_cold1, T_cold2 = u
    μ = p[1]

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

    # Energy balances (exactly from Python script)
    # C_r * dT_hot/dt = P_r - W*(T_hot - T_cold1) - W*(T_hot - T_cold2)
    du[1] = (P_r - W * (T_hot - T_cold1) - W * (T_hot - T_cold2)) / C_r

    # C_sg * dT_cold1/dt = W*(T_hot - T_cold1) - Q1
    du[2] = (W * (T_hot - T_cold1) - Q_sg) / C_sg

    # C_sg * dT_cold2/dt = W*(T_hot - T_cold2) - Q2
    du[3] = (W * (T_hot - T_cold2) - Q_sg) / C_sg

    return du
end

# Initial conditions (exact, from Python script)
u0 = [T_initial, T_initial, T_initial]

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

# Fine μ sweep
μ_values = range(0.5, 0.75, length=50)  # VERY FINE sweep
println("\n=== μ Sweep ===")
println("μ range: [0.5, 0.75]")
println("Number of cases: $(length(μ_values))")

# Time span
tspan = (0.0, 30.0)
t_save = range(0, 30, length=300)

# Solve for all μ values
println("\nSolving $(length(μ_values)) cases...")
all_sols = []
for (i, μ) in enumerate(μ_values)
    if i % 10 == 0
        print(".")
    end
    p = [μ]
    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 for each μ
all_T_ave = []
for (i, μ) in enumerate(μ_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 envelopes
T_hot_min = [minimum([all_T_hot[i][j] for i in 1:length(μ_values)]) for j in 1:length(times)]
T_hot_max = [maximum([all_T_hot[i][j] for i in 1:length(μ_values)]) for j in 1:length(times)]
T_ave_min = [minimum([all_T_ave[i][j] for i in 1:length(μ_values)]) for j in 1:length(times)]
T_ave_max = [maximum([all_T_ave[i][j] for i in 1:length(μ_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 Reach Tube",
         fillalpha=0.4, color=:red, lw=2,
         label="Reach tube", legend=:bottomright)

# Plot 2: T_ave reach tube
p2 = plot(times, T_ave_min, fillrange=T_ave_max,
         xlabel="Time (s)", ylabel="T_avg (°F)",
         title="Average Temperature Reach Tube",
         fillalpha=0.4, color=:orange, lw=2,
         label="Reach tube", legend=:bottomright)

# Plot 3: Phase portrait - just the reach tube boundary
p3 = plot(xlabel="T_avg (°F)", ylabel="T_hot (°F)",
         title="Phase Portrait: T_hot vs T_avg",
         legend=:bottomright, size=(700, 700))

# Plot the envelope
plot!(p3, [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)

# Add boundary traces
plot!(p3, T_ave_min, T_hot_min, color=:blue, lw=2, label="μ = 0.5")
plot!(p3, T_ave_max, T_hot_max, color=:red, lw=2, label="μ = 0.75")

# Plot 4: ΔT reach tube
ΔT_min = [minimum([all_T_hot[i][j] - all_T_cold1[i][j] for i in 1:length(μ_values)]) for j in 1:length(times)]
ΔT_max = [maximum([all_T_hot[i][j] - all_T_cold1[i][j] for i in 1:length(μ_values)]) for j in 1:length(times)]

p4 = plot(times, ΔT_min, fillrange=ΔT_max,
         xlabel="Time (s)", ylabel="ΔT = T_hot - T_cold (°F)",
         title="Loop Temperature Difference",
         fillalpha=0.4, color=:green, lw=2,
         label="Reach tube", legend=:bottomright)

plot_combined = plot(p1, p2, p3, p4, layout=(2, 2), size=(1400, 1000),
                    plot_title="Two-Loop Reactor: μ Uncertainty [0.5, 0.75]")
savefig(plot_combined, "two_loop_fine_sweep.png")
println("Saved: two_loop_fine_sweep.png")

# Detailed phase portrait
p_phase = plot(xlabel="T_avg (°F)", ylabel="T_hot (°F)",
              title="Phase Portrait: T_hot vs T_avg\n(μ uncertainty: [0.5, 0.75])",
              size=(1000, 900), legend=:bottomright)

# Plot all trajectories (lightly)
for (i, μ) in enumerate(μ_values)
    color_val = (μ - minimum(μ_values)) / (maximum(μ_values) - minimum(μ_values))
    plot!(p_phase, all_T_ave[i], all_T_hot[i],
          color=cgrad(:viridis)[color_val], lw=1, alpha=0.2, label="")
end

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

# Mark extremes
plot!(p_phase, all_T_ave[1], all_T_hot[1],
      color=:blue, lw=3, label="μ = 0.5", linestyle=:dash)
plot!(p_phase, all_T_ave[end], all_T_hot[end],
      color=:red, lw=3, label="μ = 0.75", linestyle=:dash)

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

println("\n✓ Complete!")
println("\nResults:")
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 range: [$(round(minimum(ΔT_min), digits=1)), $(round(maximum(ΔT_max), digits=1))] °F")