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

# Parameters
N = 50                  # Number of particles
dim = 2                 # Dimension (2D space)
beta = 1.5              # Parameter controlling influence
dt = 0.5                # Time step
num_steps = 100         # Number of steps in the animation

# Cucker-Smale influence function
def psi(r, beta):
    return 1 / (1 + r**2)**beta

# Initialize positions and velocities
np.random.seed(42)
positions = np.random.rand(N, dim) * 10
velocities = (np.random.rand(N, dim) - 0.5) * 2

# Update function for the Cucker-Smale model
def update():
    global positions, velocities
    for i in range(N):
        influence = np.zeros(dim)
        for j in range(N):
            if i != j:
                distance = np.linalg.norm(positions[i] - positions[j])
                influence += psi(distance, beta) * (velocities[j] - velocities[i])
        velocities[i] += influence * dt
    positions += velocities * dt

# Setting up the plot with velocity vectors
fig, ax = plt.subplots(figsize=(8, 8))
ax.set_xlim(0, 10)
ax.set_ylim(0, 10)
scat = ax.scatter(positions[:, 0], positions[:, 1], color='blue')
quivers = ax.quiver(positions[:, 0], positions[:, 1], velocities[:, 0], velocities[:, 1], 
                    angles='xy', scale_units='xy', scale=0.05, color='red', headwidth=4, headlength=6)

# Animation function with velocity vectors
def animate_with_velocity(frame):
    update()
    scat.set_offsets(positions)
    # Remove and redraw quiver
    quivers.set_offsets(positions)
    quivers.set_UVC(velocities[:, 0], velocities[:, 1])
    return scat, quivers

# Create animation
anim = FuncAnimation(fig, animate_with_velocity, frames=num_steps, interval=50)

if True:
    anim.save('cucker-smale.mp4', writer='ffmpeg', fps=10)

plt.show()