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

# Problem 5: Xenon-135 Transient Analysis

# Given parameters
gamma_I = 0.057          # I-135 fission yield (5.7%)
gamma_Xe = 0.003         # Xe-135 fission yield (0.3%)
lambda_I = 2.87e-5       # I-135 decay constant [sec^-1] (t_1/2 = 6.7 hr)
lambda_Xe = 2.09e-5      # Xe-135 decay constant [sec^-1] (t_1/2 = 9.2 hr)
R_max = 7.34e-5          # Full power burnout factor [sec^-1]
K = 4.56                 # Power constant [pcm*sec^-1]
Xe_eq_reactivity = -2900 # Full power equilibrium Xe reactivity [pcm]

print("=== Problem 5: Xenon-135 Transient ===")
print(f"\nGiven Parameters:")
print(f"gamma_I = {gamma_I} (I-135 fission yield)")
print(f"gamma_Xe = {gamma_Xe} (Xe-135 fission yield)")
print(f"lambda_I = {lambda_I:.2e} sec^-1 (t_1/2 = 6.7 hr)")
print(f"lambda_Xe = {lambda_Xe:.2e} sec^-1 (t_1/2 = 9.2 hr)")
print(f"R_max = {R_max:.2e} sec^-1 (full power burnout)")
print(f"K = {K} pcm*sec^-1")
print(f"Initial Xe reactivity = {Xe_eq_reactivity} pcm (at 100% power)")

# Power history (normalized to 1.0 = 100% power)
# Time ranges in hours
power_history = [
    (0, 5, 1.0),      # 0-5 hr: 100% power
    (5, 15, 0.0),     # 5-15 hr: Shutdown
    (15, 50, 1.0),    # 15-50 hr: 100% power
    (50, 80, 0.4),    # 50-80 hr: 40% power
    (80, 100, 0.0),   # 80-100 hr: Shutdown
    (100, 150, 1.0),  # 100-150 hr: 100% power
]

def get_power(t_hours):
    """Return normalized power at time t (in hours)"""
    for t_start, t_end, power in power_history:
        if t_start <= t_hours < t_end:
            return power
    return 0.0  # Default to shutdown

# Calculate equilibrium I and Xe at full power
# At equilibrium: dI/dt = 0, dX/dt = 0, rho = 1.0
# I_eq = gamma_I * Sigma_f * phi / lambda_I
# X_eq = (lambda_I * I_eq + gamma_Xe * Sigma_f * phi) / (lambda_Xe + sigma_a^Xe * phi)

# From the reactivity equation and given parameters, we can compute equilibrium
# Xe_reactivity = -K * X / rho
# At full power equilibrium: -2900 = -K * X_eq / 1.0
# X_eq = 2900 / K (in relative units)

# But we need to be careful about units. Let's work in terms of production rates.
# Production rate at full power: P_I = gamma_I * (production source)
# At equilibrium: I_eq = P_I / lambda_I
# At equilibrium: X_eq = (lambda_I * I_eq + gamma_Xe * P_source) / (lambda_Xe + R_max)

# Let's define production in terms that give correct equilibrium
# From B = rho*K*[gamma_I, gamma_Xe], production scales with power

# For now, let's compute initial conditions at full power equilibrium
rho_initial = 1.0  # Full power

# At equilibrium with rho=1:
# dI/dt = 0 = -lambda_I * I + rho * gamma_I * P0
# dX/dt = 0 = -lambda_Xe * X - rho * R_max * X + lambda_I * I + rho * gamma_Xe * P0

# Let P0 be a production constant we need to determine
# From equilibrium: I_eq = rho * gamma_I * P0 / lambda_I
# X_eq = (lambda_I * I_eq + rho * gamma_Xe * P0) / (lambda_Xe + rho * R_max)

# We know Xe_reactivity = -K * X at full power
# So X_eq = 2900 / K at rho = 1

X_eq_fullpower = abs(Xe_eq_reactivity) / K  # Equilibrium Xe "concentration" at full power
print(f"\nEquilibrium Xe concentration (full power) = {X_eq_fullpower:.2f} [arbitrary units]")

# From equilibrium equation at full power (rho=1):
# X_eq = (lambda_I * I_eq + gamma_Xe * P0) / (lambda_Xe + R_max)
# And: I_eq = gamma_I * P0 / lambda_I

# Substituting:
# X_eq = (lambda_I * gamma_I * P0 / lambda_I + gamma_Xe * P0) / (lambda_Xe + R_max)
# X_eq = P0 * (gamma_I + gamma_Xe) / (lambda_Xe + R_max)

# Solving for P0:
P0 = X_eq_fullpower * (lambda_Xe + R_max) / (gamma_I + gamma_Xe)
I_eq_fullpower = gamma_I * P0 / lambda_I

print(f"Equilibrium I-135 concentration (full power) = {I_eq_fullpower:.2f} [arbitrary units]")
print(f"Production constant P0 = {P0:.2e}")

# Define ODE system for I-135 and Xe-135
def xenon_ode(y, t, power_func):
    """
    ODE system for I-135 and Xe-135
    y = [I, X] where I is I-135, X is Xe-135
    t is time in seconds
    """
    I, X = y
    t_hours = t / 3600  # Convert to hours for power lookup
    rho = power_func(t_hours)

    # dI/dt = -lambda_I * I + rho * gamma_I * P0
    dI_dt = -lambda_I * I + rho * gamma_I * P0

    # dX/dt = -lambda_Xe * X - rho * R_max * X + lambda_I * I + rho * gamma_Xe * P0
    dX_dt = -lambda_Xe * X - rho * R_max * X + lambda_I * I + rho * gamma_Xe * P0

    return [dI_dt, dX_dt]

# Initial conditions at t=0 (full power equilibrium)
I0 = I_eq_fullpower
X0 = X_eq_fullpower
y0 = [I0, X0]

print(f"\nInitial conditions at t=0 (full power equilibrium):")
print(f"I-135 = {I0:.2f}")
print(f"Xe-135 = {X0:.2f}")
print(f"Xe reactivity = {-K * X0:.1f} pcm")

# Part A: Solve and sketch the xenon transient
print(f"\n=== Part A: Xenon Transient Sketch ===")

# Time array: 0 to 150 hours
t_hours = np.linspace(0, 150, 2000)
t_seconds = t_hours * 3600

# Solve ODE
solution = odeint(xenon_ode, y0, t_seconds, args=(get_power,))
I_transient = solution[:, 0]
X_transient = solution[:, 1]

# Calculate xenon reactivity
Xe_reactivity_transient = -K * X_transient

# Plot results
fig, (ax1, ax2, ax3) = plt.subplots(3, 1, figsize=(12, 10))

# Plot power history
power_values = [get_power(t) for t in t_hours]
ax1.plot(t_hours, power_values, 'b-', linewidth=2)
ax1.set_ylabel('Normalized Power', fontsize=12)
ax1.set_title('Power History', fontsize=14, fontweight='bold')
ax1.grid(True, alpha=0.3)
ax1.set_ylim(-0.1, 1.2)

# Plot I-135
ax2.plot(t_hours, I_transient, 'g-', linewidth=2, label='I-135')
ax2.set_ylabel('I-135 Concentration', fontsize=12)
ax2.set_title('I-135 Transient', fontsize=14, fontweight='bold')
ax2.grid(True, alpha=0.3)
ax2.legend()

# Plot Xe reactivity
ax3.plot(t_hours, Xe_reactivity_transient, 'r-', linewidth=2)
ax3.axhline(y=-2900, color='k', linestyle='--', alpha=0.5, label='Initial Eq (-2900 pcm)')
ax3.set_xlabel('Time (hours)', fontsize=12)
ax3.set_ylabel('Xenon Reactivity (pcm)', fontsize=12)
ax3.set_title('Xenon Reactivity Transient', fontsize=14, fontweight='bold')
ax3.grid(True, alpha=0.3)
ax3.legend()

plt.tight_layout()
plt.savefig('problem5_xenon_transient.png', dpi=300, bbox_inches='tight')
print("Xenon transient plot saved as 'problem5_xenon_transient.png'")
plt.close()

# Part B: Find the peak xenon after first shutdown (at t=5 hours)
print(f"\n=== Part B: First Xenon Peak Analysis ===")

# Focus on the first shutdown period (5 to 15 hours)
mask = (t_hours >= 5) & (t_hours <= 15)
t_peak_range = t_hours[mask]
Xe_react_peak_range = Xe_reactivity_transient[mask]

# Find the peak (most negative reactivity)
peak_idx = np.argmin(Xe_react_peak_range)
t_peak = t_peak_range[peak_idx]
Xe_peak_reactivity = Xe_react_peak_range[peak_idx]

print(f"\nFirst shutdown: t = 5 hours to t = 15 hours")
print(f"Peak xenon occurs at t = {t_peak:.2f} hours")
print(f"Time after shutdown = {t_peak - 5:.2f} hours")
print(f"Peak xenon reactivity = {Xe_peak_reactivity:.1f} pcm")
print(f"\nChange from equilibrium: {Xe_peak_reactivity - Xe_eq_reactivity:.1f} pcm")