# Re-importing libraries and redefining the Marchenko-Pastur animation process due to session reset
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

# Define the Marchenko-Pastur distribution function
def marchenko_pastur(x, lam):
    a = (1 - np.sqrt(lam))**2
    b = (1 + np.sqrt(lam))**2
    density = np.zeros_like(x)
    mask = (x >= a) & (x <= b)
    density[mask] = (1 / (2 * np.pi * x[mask] * lam)) * np.sqrt((b - x[mask]) * (x[mask] - a))
    return density

# Create the x range for the plot
x = np.linspace(0, 6, 1000)

# Initialize the figure and axis
fig, ax = plt.subplots()
line, = ax.plot([], [], lw=2)
ax.set_xlim(0, 6)
ax.set_ylim(0, 3)
ax.set_xlabel("x")
ax.set_ylabel("Density")
ax.set_title("Marchenko-Pastur Distribution (Varying λ)")

# Animation update function
def update(frame):
    lam = frame / 50 + 0.001  # λ varies from 0.1 to 5 over 100 frames
    y = marchenko_pastur(x, lam)
    line.set_data(x, y)
    ax.set_title(f"Marchenko-Pastur Distribution ($\\rho$ = {lam:.2f})")
    return line,

# Create the animation
ani = FuncAnimation(fig, update, frames=50, blit=False)

if True:
    ani.save("marchenko-pastur.mp4", writer="ffmpeg", fps=10)

plt.show()