import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation

np.random.seed(0)

# Define the spirograph function
def spirograph(t, R, k, l):
    x = R * ((1 - k) * np.cos(t) + l * k * np.cos((1 - k) / k * t))
    y = R * ((1 - k) * np.sin(t) - l * k * np.sin((1 - k) / k * t))
    return x, y

# Dynamically generate parameters for more diverse patterns
def generate_params(grid_size, R_range=(1, 4), k_range=(0.2, 0.8), l_range=(0.3, 0.9)):
    return [
        (4, np.random.uniform(*k_range), np.random.uniform(*l_range))
        for _ in range(grid_size**2)
    ]

# Create figure and axes for a customizable grid
grid_size = 4  # Define grid size
fig, axes = plt.subplots(grid_size, grid_size, figsize=(10, 10))
axes = axes.flatten()

# Set up plot limits, aesthetics, and line objects
lines = []
colors = plt.cm.viridis(np.linspace(0, 1, len(axes)))  # Generate a colormap
for ax, color in zip(axes, colors):
    ax.set_xlim(-5, 5)
    ax.set_ylim(-5, 5)
    ax.set_aspect('equal', 'box')
    ax.axis('off')
    line, = ax.plot([], [], lw=2, color=color)
    lines.append(line)

# Generate parameters for spirographs
params = generate_params(grid_size)

# Time values for the animation
t_vals = np.linspace(0, 25 * np.pi, 2000)

# Initialization function
def init():
    for line in lines:
        line.set_data([], [])
    return lines

# Update function for the animation
def update(frame):
    for i, (R, k, l) in enumerate(params):
        t = t_vals[:frame]  # Show up to the current frame
        x, y = spirograph(t, R, k, l)
        lines[i].set_data(x, y)
    return lines

# Create the animation
ani = FuncAnimation(fig, update, frames=len(t_vals), init_func=init, blit=True, interval=2)

if True:
    # Save the animation as an MP4 file
    ani.save('spirograph.mp4', writer='ffmpeg', fps=60)

# Save or display the animation
plt.show()