import pygame
import random
import math

WIDTH, HEIGHT = 800, 600

NUM_GROUPS = 3
BOIDS_PER_GROUP = 20
MAX_SPEED = 4
NEIGHBOR_RADIUS = 60
AVOID_RADIUS = 100

ALIGNMENT_WEIGHT = 1.0
COHESION_WEIGHT = 1.0
SEPARATION_WEIGHT = 3.0

GROUP_COLOR = [(200, 50, 50), (50, 200, 50), (50, 50, 200)]

class Boid:
    def __init__(self, x, y, group_id):
        self.position = pygame.Vector2(x, y)
        self.boid_vector = pygame.Vector2(0.5, 0.5).normalize()
        self.group_id = group_id

    def update(self, boids):
        alignment = self.align(boids) * ALIGNMENT_WEIGHT # dopasowanie
        cohesion = self.cohesion(boids) * COHESION_WEIGHT # kohesja
        separation = self.separation(boids) * SEPARATION_WEIGHT # separacja

        steering = alignment + cohesion + separation

        self.boid_vector += steering
        if self.boid_vector.length() > MAX_SPEED:
            self.boid_vector.scale_to_length(MAX_SPEED)

        self.position += self.boid_vector

        if self.position.x > WIDTH:
            self.position.x = WIDTH
            self.boid_vector.x *= -1
        if self.position.x < 0:
            self.position.x = 0
            self.boid_vector.x *= -1
        if self.position.y > HEIGHT:
            self.position.y = HEIGHT
            self.boid_vector.y *= -1
        if self.position.y < 0:
            self.position.y = 0
            self.boid_vector.y *= -1

    def align(self, boids):
        avg_boid_vector = pygame.Vector2(0, 0)
        count = 0
        for boid in boids:
            if boid == self or boid.group_id != self.group_id:
                continue
            if self.position.distance_to(boid.position) < NEIGHBOR_RADIUS:
                avg_boid_vector += boid.boid_vector
                count += 1
        if count > 0:
            avg_boid_vector = avg_boid_vector/count
            avg_boid_vector = avg_boid_vector.normalize() * MAX_SPEED
        return avg_boid_vector - self.boid_vector

    def cohesion(self, boids):
        center_of_group = pygame.Vector2(0, 0)
        count = 0
        for boid in boids:
            if boid == self or boid.group_id != self.group_id:
                continue
            if self.position.distance_to(boid.position) < NEIGHBOR_RADIUS:
                center_of_group += boid.position
                count += 1
        if count > 0:
            center_of_group = center_of_group/count
            return (center_of_group - self.position).normalize()
        return pygame.Vector2(0, 0)

    def separation(self, boids):
        avoid_vector = pygame.Vector2(0, 0)
        for boid in boids:
            if boid == self or boid.group_id != self.group_id:
                continue
            distance = self.position.distance_to(boid.position)
            if distance < AVOID_RADIUS:
                avoid_vector += (self.position - boid.position).normalize() / distance
        return avoid_vector
    def draw(self, screen):
        angle = math.atan2(self.boid_vector.y, self.boid_vector.x)
        points = [
            (self.position.x + math.cos(angle) * 10, self.position.y + math.sin(angle) * 10),
            (self.position.x + math.cos(angle + 2.5) * 5, self.position.y + math.sin(angle + 2.5) * 5),
            (self.position.x + math.cos(angle - 2.5) * 5, self.position.y + math.sin(angle - 2.5) * 5),
        ]
        pygame.draw.polygon(screen, GROUP_COLOR[self.group_id], points)

pygame.init()
screen = pygame.display.set_mode((WIDTH, HEIGHT))
clock = pygame.time.Clock()

boids = []
for group_id in range(NUM_GROUPS):
    for _ in range(BOIDS_PER_GROUP):
        x = random.randint(0, WIDTH)
        y = random.randint(0, HEIGHT)
        boids.append(Boid(x, y, group_id))

running = True
while running:
    screen.fill((30, 30, 30))

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            running = False

    for boid in boids:
        boid.update(boids)
        boid.draw(screen)

    pygame.display.flip()
    clock.tick(60)

pygame.quit()