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

plt.rcParams['text.usetex'] = True

# Constants
H0 = 67400 / (3.09 * 10**22) * (3600 * 24 * 365 * 10**9)  # Hubble constant in Gyr^-1

# Function for the differential system
def friedmann_system(t, y, Omega_m, Omega_r, Omega_l, Omega_k):
    a, a_dot = y
    da_dt = a_dot
    denominator = 2 * (H0**2 * (Omega_r / a**2 + Omega_l * a**2 + Omega_k))
    if denominator == 0 or a <= 0:
        da_dot_dt = 0  # Prevent division by zero or negative scale factors
    else:
        da_dot_dt = -0.5 * H0**2 * (Omega_m / a**3 + 2 * Omega_r / a**4 - 2 * Omega_l * a)
    return [da_dt, da_dot_dt]

# Event function to stop the integration when `a` reaches 0
def stop_event_a_zero(t, y):
    a, _ = y
    return a - 1e-8  # Stop when a approaches 0

# Event function to stop the integration when `a` reaches 4
def stop_event_a_limit(t, y):
    a, _ = y
    return 4 - a  # Stop when a approaches 4

# Mark both events as terminal
stop_event_a_zero.terminal = True
stop_event_a_zero.direction = -1  # Stop when decreasing to 0
stop_event_a_limit.terminal = True
stop_event_a_limit.direction = 1  # Stop when increasing to 4

# Initial conditions
a0 = 1.0
a_dot0_template = lambda Omega_m, Omega_r, Omega_l, Omega_k: H0 * (
    Omega_m / a0 + Omega_r / a0**2 + Omega_k + Omega_l * a0**2
)**0.5

scenarios = [
    {"Omega_m": 0.3, "Omega_r": 1e-4, "Omega_l": 0.0, "color": "blue", "label": r"$\Omega_m=0.3, \Omega_\Lambda=0, \Omega_r=10^{-4}, k=-1$"},
    {"Omega_m": 0.3, "Omega_r": 1e-4, "Omega_l": 0.7, "color": "red", "label": r"$\Omega_m=0.3, \Omega_\Lambda=0.7, \Omega_r=10^{-4}, \Omega_k \simeq 0$"},
    {"Omega_m": 5.0, "Omega_r": 1e-4, "Omega_l": 0.0, "color": "black", "label": r"$\Omega_m=5, \Omega_\Lambda=0, \Omega_r=10^{-4}, k=1$"},
    {"Omega_m": 1.0, "Omega_r": 1e-4, "Omega_l": 0.0, "color": "magenta", "label": r"$\Omega_m=1, \Omega_\Lambda=0, \Omega_r=10^{-4}, k=1$"},
]

# Solve and plot
plt.figure(figsize=(10, 8))  # Square plot

for scenario in scenarios:
    Omega_m = scenario["Omega_m"]
    Omega_r = scenario["Omega_r"]
    Omega_l = scenario["Omega_l"]
    Omega_k = 1 - Omega_m - Omega_r - Omega_l  # Compute curvature parameter

    # Compute initial a_dot
    a_dot0 = a_dot0_template(Omega_m, Omega_r, Omega_l, Omega_k)
    initial_conditions = [a0, a_dot0]

    # Solve backwards and forwards with increased resolution
    solution_backward = solve_ivp(
        lambda t, y: friedmann_system(t, y, Omega_m, Omega_r, Omega_l, Omega_k),
        [0, -15],
        initial_conditions,
        events=[stop_event_a_zero],
        dense_output=False,
        max_step=0.01  # Small step size for higher resolution
    )
    solution_forward = solve_ivp(
        lambda t, y: friedmann_system(t, y, Omega_m, Omega_r, Omega_l, Omega_k),
        [0, 30],
        initial_conditions,
        events=[stop_event_a_limit],
        dense_output=False,
        max_step=0.01  # Small step size for higher resolution
    )

    # Combine results
    t = np.concatenate((solution_backward.t[::-1], solution_forward.t))
    a = np.concatenate((solution_backward.y[0][::-1], solution_forward.y[0]))

    # Plot
    plt.plot(t, a, color=scenario["color"], label=scenario["label"])

# Finalize plot
plt.axhline(0, color="black", linewidth=0.5, linestyle="--")
plt.axvline(0, color="black", linewidth=0.5, linestyle="--")
plt.grid(True)
plt.xlabel("Cosmic Time (Gyr)", fontsize=15)  # Increase font size for x-axis
plt.ylabel("Relative size of the universe", fontsize=15)  # Increase font size for y-axis
plt.title("Relative size of the universe vs Cosmic Time", fontsize=17)  # Increase font size for title
plt.xlim(-15, 30)
plt.ylim(0, 4)
plt.legend(fontsize=15)  # Increase font size for legend
plt.tight_layout()
plt.savefig("friedmann_solving.pdf")
plt.show()