import numpy as np
import time
import matplotlib
matplotlib.use('Agg') # Safe headless mode for Linux/ModernGL interop
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation, FFMpegWriter
from vgpu_cache import VGPUCache

def simulate_double_pendulum_swarm(num_pendulums=1500, steps=300, dt=0.015):
    c = VGPUCache()
    g = np.float32(9.81)
    
    # 1. Initialize states on a fine mesh (tiny angular variation creates the chaos)
    # Theta1 initialized slightly away from vertical, theta2 has fractional offsets
    t1 = np.ones(num_pendulums, dtype=np.float32) * (np.pi / 2.0)
    t2 = np.ones(num_pendulums, dtype=np.float32) * (np.pi / 2.0)
    # Add an ultra-fine variation between 0 and 0.001 radians
    t2 += np.linspace(0.0, 0.001, num_pendulums, dtype=np.float32)
    
    # Angular velocities
    w1 = np.zeros(num_pendulums, dtype=np.float32)
    w2 = np.zeros(num_pendulums, dtype=np.float32)
    
    # Storage arrays to record every frame's Cartesian coordinates
    history_x1 = np.zeros((steps, num_pendulums), dtype=np.float32)
    history_y1 = np.zeros((steps, num_pendulums), dtype=np.float32)
    history_x2 = np.zeros((steps, num_pendulums), dtype=np.float32)
    history_y2 = np.zeros((steps, num_pendulums), dtype=np.float32)
    
    print(f"Simulating {num_pendulums} double pendulums across {steps} steps on GPU...")
    t0 = time.time()
    
    for step in range(steps):
        # Cache angular differences and trig ops on the GPU
        delta = c.subtract(t1, t2)
        sin_d = c.sin(delta)
        cos_d = c.cos(delta)
        
        sin_t1 = c.sin(t1)
        sin_t2 = c.sin(t2)
        
        # --- CALCULATE ANGULAR ACCELERATIONS (α1, α2) VIA ELEMENTWISE GPU OPERATIONS ---
        # Denominator: 2 - cos(t1 - t2)^2
        cos_d_sq = c.multiply(cos_d, cos_d)
        den = c.subtract(np.float32(2.0), cos_d_sq)
        
        # Numerator 1: -g*(2*sin(t1)) - g*sin(t1-2*t2) - 2*sin(t1-t2)*(w2^2 + w1^2*cos(t1-t2))
        w1_sq = c.multiply(w1, w1)
        w2_sq = c.multiply(w2, w2)
        
        w1_sq_cos = c.multiply(w1_sq, cos_d)
        w2_w1_cos = c.add(w2_sq, w1_sq_cos)
        two_sin_d = c.multiply(np.float32(2.0), sin_d)
        rhs1 = c.multiply(two_sin_d, w2_w1_cos)
        
        two_g = np.float32(2.0 * 9.81)
        g_sin_t1_2 = c.multiply(two_g, sin_t1)
        
        # Approximated Explicit Lagrangian coupling terms
        num1_part = c.add(g_sin_t1_2, c.multiply(g, sin_t2))
        num1 = c.negative(c.add(num1_part, rhs1))
        alpha1 = c.divide(num1, den)
        
        # Numerator 2: 2*sin(t1-t2)*(2*w1^2 + g*cos(t1) + w2^2*cos(t1-t2))
        two_w1_sq = c.multiply(np.float32(2.0), w1_sq)
        g_cos_t1 = c.multiply(g, c.cos(t1))
        w2_sq_cos = c.multiply(w2_sq, cos_d)
        
        num2_part = c.add(two_w1_sq, g_cos_t1)
        num2_part = c.add(num2_part, w2_sq_cos)
        num2 = c.multiply(two_sin_d, num2_part)
        alpha2 = c.divide(num2, den)
        
        # --- TIME INTEGRATION (Euler-Cromer) ---
        dt_gpu = np.float32(dt)
        w1 = c.add(w1, c.multiply(alpha1, dt_gpu))
        w2 = c.add(w2, c.multiply(alpha2, dt_gpu))
        t1 = c.add(t1, c.multiply(w1, dt_gpu))
        t2 = c.add(t2, c.multiply(w2, dt_gpu))
        
        # --- CONVERT TO CARTESIAN COORDINATES ---
        x1 = c.sin(t1)
        y1 = c.negative(c.cos(t1))
        x2 = c.add(x1, c.sin(t2))
        y2 = c.subtract(y1, c.cos(t2))
        
        # Store positions safely to memory
        history_x1[step] = x1
        history_y1[step] = y1
        history_x2[step] = x2
        history_y2[step] = y2
        
    print(f"Simulation done in {time.time() - t0:.4f}s. Generating Animation...")
    print(c.report())
    return history_x1, history_y1, history_x2, history_y2

if __name__ == "__main__":
    steps = 250
    h_x1, h_y1, h_x2, h_y2 = simulate_double_pendulum_swarm(num_pendulums=1500, steps=steps)
    
    # ---- ANIMATION GENERATION ----
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.set_xlim(-2.2, 2.2)
    ax.set_ylim(-2.2, 2.2)
    ax.set_aspect('equal')
    ax.axis('off')
    fig.patch.set_facecolor('black')
    ax.set_facecolor('black')
    
    # Use a scatter plot to represent the 1,500 secondary bobs moving in unison
    # Jet/rainbow colormap beautifully colors them based on their divergence index!
    colors = plt.cm.jet(np.linspace(0, 1, h_x2.shape[1]))
    #scat = ax.scatter([], [], c=colors, s=1.5, alpha=0.7)
    # CHANGE THIS LINE:
    # scat = ax.scatter([], [], c=colors, s=1.5, alpha=0.7)
    
    # TO THIS (Initialize with the first frame's positions):
    scat = ax.scatter(h_x2[0], h_y2[0], c=colors, s=1.5, alpha=0.7)
    
    # Optional: Draw a single reference stick line for the first pendulum in the swarm
    line, = ax.plot([], [], 'w-', lw=1.5, alpha=0.5)
    
    def update(frame):
        # Update positions for all 1,500 points at current time-step
        x_data = h_x2[frame]
        y_data = h_y2[frame]
        scat.set_offsets(np.c_[x_data, y_data])
        
        # Track the arm of the first pendulum item
        line.set_data([0, h_x1[frame, 0], h_x2[frame, 0]], [0, h_y1[frame, 0], h_y2[frame, 0]])
        return scat, line

    anim = FuncAnimation(fig, update, frames=steps, interval=25, blit=True)
    
    output_gif = "double_pendulum_swarm.gif"
    # Save as high quality GIF image
    anim.save(output_gif, writer='pillow', fps=40)
    print(f"✅ Awesome! Swarm animation saved successfully as '{output_gif}'")