import numpy as np
import matplotlib.pyplot as plt

class Boid:
    def __init__(self, x, y, vx, vy):
        self.position = np.array([x, y])
        self.velocity = np.array([vx, vy])

class BoidsSimulation:
    def __init__(self, neighborhood_radius=2.0, separation_weight=2.5, alignment_weight=2.0,
                 cohesion_weight=3.0, max_speed=0.05, dt=0.01):
        self.boids = []
        self.neighborhood_radius = neighborhood_radius
        self.separation_weight = separation_weight
        self.alignment_weight = alignment_weight
        self.cohesion_weight = cohesion_weight
        self.max_speed = max_speed
        self.dt = dt

    def add_boid(self, boid):
        self.boids.append(boid)

    def separation(self, boid):
        steering = np.zeros(2)
        for other_boid in self.boids:
            if other_boid != boid:
                distance = np.linalg.norm(boid.position - other_boid.position)
                if distance < self.neighborhood_radius:
                    steering += (boid.position - other_boid.position) / distance
        return steering

    def alignment(self, boid):
        average_velocity = np.zeros(2)
        count = 0
        for other_boid in self.boids:
            if other_boid != boid:
                distance = np.linalg.norm(boid.position - other_boid.position)
                if distance < self.neighborhood_radius:
                    average_velocity += other_boid.velocity
                    count += 1
        if count > 0:
            average_velocity /= count
        return average_velocity - boid.velocity

    def cohesion(self, boid):
        center_of_mass = np.zeros(2)
        count = 0
        for other_boid in self.boids:
            if other_boid != boid:
                distance = np.linalg.norm(boid.position - other_boid.position)
                if distance < self.neighborhood_radius:
                    center_of_mass += other_boid.position
                    count += 1
        if count > 0:
            center_of_mass /= count
            return center_of_mass - boid.position
        return np.zeros(2)

    def update(self):
        for boid in self.boids:
            separation_force = self.separation(boid) * self.separation_weight
            alignment_force = self.alignment(boid) * self.alignment_weight
            cohesion_force = self.cohesion(boid) * self.cohesion_weight

            boid.velocity += (separation_force + alignment_force + cohesion_force) * self.dt
            speed = np.linalg.norm(boid.velocity)
            if speed > self.max_speed:
                boid.velocity = (boid.velocity / speed) * self.max_speed
            boid.position += boid.velocity * self.dt

    def simulate(self, steps=100):
        trajectories = [np.zeros((steps, 2)) for _ in self.boids]
        for i in range(steps):
            for j, boid in enumerate(self.boids):
                trajectories[j][i] = boid.position
            self.update()

        for traj in trajectories:
            plt.plot(traj[:, 0], traj[:, 1])
        plt.title('Smooth Boids Dynamics')
        plt.xlabel('X Position')
        plt.ylabel('Y Position')
        plt.grid(True)
        plt.legend([f'Boid {i+1}' for i in range(len(self.boids))])
        plt.show()

# Create an instance and run the simulation
sim = BoidsSimulation()

# Adding 5 boids with random positions and velocities
np.random.seed(42)  # for reproducibility
for _ in range(10):
    boid = Boid(np.random.rand(), np.random.rand(), (np.random.rand()-0.5)/10, (np.random.rand()-0.5)/10)
    sim.add_boid(boid)

sim.simulate(steps=3000)
