import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import odeint

# Problem 1: Two-Loop Reactor System

print("=== Problem 1: Two-Loop Reactor with Steam Generators ===\n")

# Given Parameters
mu = None  # Reactor Water Mass Fraction (RWMF) - to be determined or specified

# Base parameters
C_0 = 33.33  # Base Heat Capacity [%-sec/°F]
tau_0 = 0.75 * C_0  # Base Time Constant [sec/(%-sec/°F)] * C_0 = [sec]

# Initial steady state temperature
T_hot_initial = 450  # °F (reactor outlet / hot leg temperature)
T_cold_1_initial = 450  # °F (cold leg 1 temperature)
T_cold_2_initial = 450  # °F (cold leg 2 temperature)

print("Base Parameters:")
print(f"C_0 (Base Heat Capacity) = {C_0:.2f} %-sec/°F")
print(f"tau_0 (Base Time Constant) = {tau_0:.2f} sec")
print(f"\nInitial Steady State:")
print(f"T_hot = {T_hot_initial} °F")
print(f"T_cold_1 = {T_cold_1_initial} °F")
print(f"T_cold_2 = {T_cold_2_initial} °F")


# System parameters as functions of mu (RWMF)
def calculate_system_parameters(mu):
    """
    Calculate system parameters based on Reactor Water Mass Fraction (mu)

    Parameters:
    mu: Reactor Water Mass Fraction (RWMF)

    Returns:
    dict: Dictionary containing all calculated parameters
    """
    # Reactor parameters
    tau_r = mu * tau_0  # Reactor time constant [sec]
    C_r = mu * C_0  # Reactor heat capacity [%-sec/°F]

    # Steam Generator parameters (each)
    tau_sg = (1 - mu) * tau_r / 2  # Steam Generator water time constant [sec]
    C_sg = (1 - mu) * C_0 / 2  # Steam Generator heat capacity [%-sec/°F]

    params = {
        "mu": mu,
        "tau_r": tau_r,
        "C_r": C_r,
        "tau_sg": tau_sg,
        "C_sg": C_sg,
        "C_0": C_0,
        "tau_0": tau_0,
    }

    return params


# Example: Calculate parameters for mu = 0.5 (50% of water in reactor)
if mu is None:
    mu_example = 0.5
    print(f"\n--- Example with mu = {mu_example} ---")
    params = calculate_system_parameters(mu_example)

    print(f"\nReactor Parameters:")
    print(f"  tau_r (Reactor time constant) = {params['tau_r']:.2f} sec")
    print(f"  C_r (Reactor heat capacity) = {params['C_r']:.2f} %-sec/°F")

    print(f"\nSteam Generator Parameters (each):")
    print(f"  tau_sg (SG water time constant) = {params['tau_sg']:.2f} sec")
    print(f"  C_sg (SG heat capacity) = {params['C_sg']:.2f} %-sec/°F")

# System description
print(f"\n" + "=" * 60)
print("System Description:")
print("=" * 60)
print(
    """
Two-Loop Reactor System:
- Reactor (hot leg) at T_hot
- Steam Generator 1 with cold leg at T_cold_1
- Steam Generator 2 with cold leg at T_cold_2
- Each loop: pump → reactor → steam generator → back to pump
- Lumped reactor volume with combined hot branches
- Equal mass flow rates in each loop
- Ignore loop transport times

Key Assumptions:
- Each SG + cold loop has same water mass and flow rate
- Reactor and hot branches lumped into single water volume
- System initially at equilibrium (100% power, balanced)
- Constant average temperature maintained
"""
)

# Part A: Differential Equations
print(f"\n" + "=" * 60)
print("Part A: Differential Equations")
print("=" * 60)


def get_differential_equations(mu):
    """
    Develop differential equations for T_hot, T_cold1, T_cold2

    Key insight: Flow heat capacity rate W = C_0 / (2 * tau_0)
    This is the (mass flow rate * specific heat) for each loop
    """
    params = calculate_system_parameters(mu)

    # Flow heat capacity rate for each loop (same for both loops)
    W = C_0 / (2 * tau_0)

    print(f"\nFor mu = {mu}:")
    print(f"Flow heat capacity rate (each loop): W = {W:.4f} %-sec/°F")
    print(f"C_r = {params['C_r']:.2f}, C_sg = {params['C_sg']:.2f}")

    print(f"\nDifferential Equations:")
    print(f"  C_r * dT_hot/dt = P_r - W*(T_hot - T_cold1) - W*(T_hot - T_cold2)")
    print(f"  C_sg * dT_cold1/dt = W*(T_hot - T_cold1) - Qdot1")
    print(f"  C_sg * dT_cold2/dt = W*(T_hot - T_cold2) - Qdot2")
    print(f"\nwhere:")
    print(f"  P_r = reactor power (positive for heat generation)")
    print(f"  Qdot1, Qdot2 = SG heat removal rates (positive for heat removal)")
    print(f"  W = {W:.4f} %-sec/°F (flow heat capacity rate per loop)")

    # Matrix form: dT/dt = A*T + B
    print(f"\n--- Matrix Form: dT/dt = A*T + B ---")

    # Coefficient matrix A
    A = np.array(
        [
            [-2 * W / params["C_r"], W / params["C_r"], W / params["C_r"]],
            [W / params["C_sg"], -W / params["C_sg"], 0],
            [W / params["C_sg"], 0, -W / params["C_sg"]],
        ]
    )

    print(f"\nMatrix A (coefficients):")
    print(f"  [{A[0,0]:7.4f}  {A[0,1]:7.4f}  {A[0,2]:7.4f}]")
    print(f"  [{A[1,0]:7.4f}  {A[1,1]:7.4f}  {A[1,2]:7.4f}]")
    print(f"  [{A[2,0]:7.4f}  {A[2,1]:7.4f}  {A[2,2]:7.4f}]")

    # For general case with P_r, Qdot1, Qdot2
    print(f"\nVector B (forcing terms, depends on P_r, Qdot1, Qdot2):")
    print(f"  [P_r/C_r, -Qdot1/C_sg, -Qdot2/C_sg]^T")
    print(
        f"  [P_r/{params['C_r']:.2f}, -Qdot1/{params['C_sg']:.2f}, -Qdot2/{params['C_sg']:.2f}]^T"
    )

    return W, params


# Part D: Transient Simulation Setup
print(f"\n" + "=" * 60)
print("Part D: Transient Simulation (0 to 30 seconds)")
print("=" * 60)


def reactor_odes(y, t, W, C_r, C_sg, P_r, Q1, Q2):
    """
    ODE system for two-loop reactor
    y = [T_hot, T_cold1, T_cold2]

    Parameters:
    P_r: Reactor power (heat generation rate)
    Q1, Q2: Steam generator heat removal rates (Qdot_1, Qdot_2)
    """
    T_hot, T_cold1, T_cold2 = y

    # Reactor energy balance
    dT_hot_dt = (P_r - W * (T_hot - T_cold1) - W * (T_hot - T_cold2)) / C_r

    # Steam generator 1 energy balance (Q1 is heat removal rate)
    dT_cold1_dt = (W * (T_hot - T_cold1) - Q1) / C_sg

    # Steam generator 2 energy balance (Q2 is heat removal rate)
    dT_cold2_dt = (W * (T_hot - T_cold2) - Q2) / C_sg

    return [dT_hot_dt, dT_cold1_dt, dT_cold2_dt]


def simulate_transient(mu, Q1_percent, Q2_percent, t_max=30):
    """
    Simulate reactor transient

    Parameters:
    mu: Reactor Water Mass Fraction
    Q1_percent: Steam generator 1 heat removal rate (% of nominal)
    Q2_percent: Steam generator 2 heat removal rate (% of nominal)
    t_max: Simulation time (seconds)
    """
    params = calculate_system_parameters(mu)
    W = C_0 / (2 * tau_0)

    # Initial conditions (steady state at 100% power balanced)
    T_hot_0 = T_hot_initial
    T_cold1_0 = T_cold_1_initial
    T_cold2_0 = T_cold_2_initial
    y0 = [T_hot_0, T_cold1_0, T_cold2_0]

    # Power levels (normalized - 100% = 1.0 in appropriate units)
    P_r = 100.0  # Reactor power (heat generation rate)
    Q1 = Q1_percent  # SG1 heat removal rate
    Q2 = Q2_percent  # SG2 heat removal rate

    # Time array
    t = np.linspace(0, t_max, 500)

    # Solve ODE system
    solution = odeint(
        reactor_odes, y0, t, args=(W, params["C_r"], params["C_sg"], P_r, Q1, Q2)
    )

    T_hot = solution[:, 0]
    T_cold1 = solution[:, 1]
    T_cold2 = solution[:, 2]

    # Calculate average temperature (mass-weighted, depends on mu)
    T_ave = mu * T_hot + (1 - mu) * (T_cold1 + T_cold2) / 2

    return t, T_hot, T_cold1, T_cold2, T_ave


# Example: Show how to set up simulations
print("\nSimulation cases to run:")
print("1. mu=0.5, Q1=50%, Q2=50% (equally loaded)")
print("2. mu=0.75, Q1=50%, Q2=50% (equally loaded)")
print("3. mu=0.5, Q1=60%, Q2=40% (unequally loaded)")
print("4. mu=0.75, Q1=60%, Q2=40% (unequally loaded)")
print("\nNote: Percentages are fractions of reactor power (P_r = 100%)")

# Run example simulation
print("\n--- Running Example: mu=0.5, Equal Loading ---")
W_example, params_example = get_differential_equations(0.5)

# Run all simulations and create plots
print("\n" + "=" * 60)
print("Running Simulations and Creating Plots")
print("=" * 60)

# Case 1 & 2: Equally loaded (Q1=50%, Q2=50%)
fig1, axes1 = plt.subplots(2, 2, figsize=(14, 10))
fig1.suptitle(
    "Equally Loaded Steam Generators (Q1=50%, Q2=50%)", fontsize=14, fontweight="bold"
)

for idx, mu in enumerate([0.5, 0.75]):
    print(f"\nSimulating mu={mu}, Q1=50%, Q2=50%")
    t, T_hot, T_cold1, T_cold2, T_ave = simulate_transient(mu, 50.0, 50.0)

    # Plot temperatures
    ax = axes1[idx, 0]
    ax.plot(t, T_hot, "r-", linewidth=2, label="T_hot")
    ax.plot(t, T_cold1, "b-", linewidth=2, label="T_cold1")
    ax.plot(t, T_cold2, "g-", linewidth=2, label="T_cold2")
    ax.plot(t, T_ave, "k--", linewidth=2, label="T_ave")
    ax.set_xlabel("Time (sec)", fontsize=10)
    ax.set_ylabel("Temperature (°F)", fontsize=10)
    ax.set_title(f"mu = {mu} (RWMF)", fontsize=11, fontweight="bold")
    ax.grid(True, alpha=0.3)
    ax.legend(loc="best")

    # Plot temperature differences
    ax = axes1[idx, 1]
    ax.plot(t, T_hot - T_cold1, "b-", linewidth=2, label="ΔT1 = T_hot - T_cold1")
    ax.plot(t, T_hot - T_cold2, "g-", linewidth=2, label="ΔT2 = T_hot - T_cold2")
    ax.set_xlabel("Time (sec)", fontsize=10)
    ax.set_ylabel("Temperature Difference (°F)", fontsize=10)
    ax.set_title(f"Temperature Differences (mu = {mu})", fontsize=11, fontweight="bold")
    ax.grid(True, alpha=0.3)
    ax.legend(loc="best")

plt.tight_layout()
plt.savefig("problem1_equal_loading.png", dpi=300, bbox_inches="tight")
print("Plot saved: problem1_equal_loading.png")
plt.close()

# Case 3 & 4: Unequally loaded (Q1=60%, Q2=40%)
fig2, axes2 = plt.subplots(2, 2, figsize=(14, 10))
fig2.suptitle(
    "Unequally Loaded Steam Generators (Q1=60%, Q2=40%)", fontsize=14, fontweight="bold"
)

for idx, mu in enumerate([0.5, 0.75]):
    print(f"\nSimulating mu={mu}, Q1=60%, Q2=40%")
    t, T_hot, T_cold1, T_cold2, T_ave = simulate_transient(mu, 60.0, 40.0)

    # Plot temperatures
    ax = axes2[idx, 0]
    ax.plot(t, T_hot, "r-", linewidth=2, label="T_hot")
    ax.plot(t, T_cold1, "b-", linewidth=2, label="T_cold1")
    ax.plot(t, T_cold2, "g-", linewidth=2, label="T_cold2")
    ax.plot(t, T_ave, "k--", linewidth=2, label="T_ave")
    ax.set_xlabel("Time (sec)", fontsize=10)
    ax.set_ylabel("Temperature (°F)", fontsize=10)
    ax.set_title(f"mu = {mu} (RWMF)", fontsize=11, fontweight="bold")
    ax.grid(True, alpha=0.3)
    ax.legend(loc="best")

    # Plot temperature differences
    ax = axes2[idx, 1]
    ax.plot(t, T_hot - T_cold1, "b-", linewidth=2, label="ΔT1 = T_hot - T_cold1")
    ax.plot(t, T_hot - T_cold2, "g-", linewidth=2, label="ΔT2 = T_hot - T_cold2")
    ax.set_xlabel("Time (sec)", fontsize=10)
    ax.set_ylabel("Temperature Difference (°F)", fontsize=10)
    ax.set_title(f"Temperature Differences (mu = {mu})", fontsize=11, fontweight="bold")
    ax.grid(True, alpha=0.3)
    ax.legend(loc="best")

plt.tight_layout()
plt.savefig("problem1_unequal_loading.png", dpi=300, bbox_inches="tight")
print("Plot saved: problem1_unequal_loading.png")
plt.close()

print("\n" + "=" * 60)
print("Simulation Complete!")
print("=" * 60)