import pygame
import random
import math

# Screen dimensions and simulation parameters
WIDTH, HEIGHT = 800, 600
NUM_BOIDS = 30
MAX_SPEED = 4
MAX_FORCE = 0.05

NEIGHBOR_RADIUS = 50
SEPARATION_RADIUS = 20

# Debug mode flag (toggle with "D")
DEBUG_MODE = True

class Boid:
    def __init__(self, x, y):
        self.position = pygame.math.Vector2(x, y)
        angle = random.uniform(0, 2 * math.pi)
        self.velocity = pygame.math.Vector2(math.cos(angle), math.sin(angle))
        self.velocity.scale_to_length(random.uniform(1, MAX_SPEED))
        self.acceleration = pygame.math.Vector2(0, 0)
    
    def edges(self):
        # Wrap-around behavior for screen edges
        if self.position.x > WIDTH:
            self.position.x = 0
        elif self.position.x < 0:
            self.position.x = WIDTH
        if self.position.y > HEIGHT:
            self.position.y = 0
        elif self.position.y < 0:
            self.position.y = HEIGHT

    def update(self):
        # Update velocity and position
        self.velocity += self.acceleration
        if self.velocity.length() > MAX_SPEED:
            self.velocity.scale_to_length(MAX_SPEED)
        self.position += self.velocity
        self.acceleration *= 0

    def apply_force(self, force):
        self.acceleration += force

    def flock(self, boids):
        # Calculate steering vectors from the three flocking rules
        alignment = self.align(boids)
        cohesion = self.cohere(boids)
        separation = self.separate(boids)

        # Weighing the forces
        alignment *= 1.0
        cohesion *= 0.8
        separation *= 1.5

        self.apply_force(alignment)
        self.apply_force(cohesion)
        self.apply_force(separation)

    def align(self, boids):
        steering = pygame.math.Vector2(0, 0)
        total = 0
        for other in boids:
            if other != self and self.position.distance_to(other.position) < NEIGHBOR_RADIUS:
                steering += other.velocity
                total += 1
        if total > 0:
            steering /= total
            if steering.length() > 0:
                steering.scale_to_length(MAX_SPEED)
            steering -= self.velocity
            if steering.length() > MAX_FORCE:
                steering.scale_to_length(MAX_FORCE)
        return steering

    def cohere(self, boids):
        steering = pygame.math.Vector2(0, 0)
        total = 0
        for other in boids:
            if other != self and self.position.distance_to(other.position) < NEIGHBOR_RADIUS:
                steering += other.position
                total += 1
        if total > 0:
            steering /= total
            steering -= self.position
            if steering.length() > 0:
                steering.scale_to_length(MAX_SPEED)
            steering -= self.velocity
            if steering.length() > MAX_FORCE:
                steering.scale_to_length(MAX_FORCE)
        return steering

    def separate(self, boids):
        steering = pygame.math.Vector2(0, 0)
        total = 0
        for other in boids:
            distance = self.position.distance_to(other.position)
            if other != self and distance < SEPARATION_RADIUS:
                diff = self.position - other.position
                if distance > 0:
                    diff /= distance  # weight by distance
                steering += diff
                total += 1
        if total > 0:
            steering /= total
            if steering.length() > 0:
                steering.scale_to_length(MAX_SPEED)
            steering -= self.velocity
            if steering.length() > MAX_FORCE:
                steering.scale_to_length(MAX_FORCE)
        return steering

    def draw(self, screen):
        # Calculate orientation for drawing the boid as a triangle
        angle = self.velocity.angle_to(pygame.math.Vector2(1, 0))
        # Triangle points: front and two rear corners
        p1 = self.position + pygame.math.Vector2(10, 0).rotate(-angle)
        p2 = self.position + pygame.math.Vector2(-5, 5).rotate(-angle)
        p3 = self.position + pygame.math.Vector2(-5, -5).rotate(-angle)
        pygame.draw.polygon(screen, (255, 255, 255), [p1, p2, p3])
        
        # If debug is enabled, draw additional information (velocity vector)
        if DEBUG_MODE:
            end_pos = self.position + self.velocity * 10
            pygame.draw.line(screen, (0, 255, 0), self.position, end_pos, 2)

def draw_debug_panel(screen, boids, clock):
    # Draw a semi-transparent panel in the top-left corner

    font = pygame.font.SysFont("Arial", 14)
    # Collect debug information
    fps_text = font.render(f"FPS: {int(clock.get_fps())}", True, (255, 255, 255))
    boids_text = font.render(f"Boids: {len(boids)}", True, (255, 255, 255))
    max_speed_text = font.render(f"Max Speed: {MAX_SPEED}", True, (255, 255, 255))
    neighbor_text = font.render(f"Neighbor Radius: {NEIGHBOR_RADIUS}", True, (255, 255, 255))
    separation_text = font.render(f"Separation Radius: {SEPARATION_RADIUS}", True, (255, 255, 255))
    
    # Blit debug info on the panel
    screen.blit(fps_text, (10, 10))
    screen.blit(boids_text, (10, 30))
    screen.blit(max_speed_text, (10, 50))
    screen.blit(neighbor_text, (10, 70))
    screen.blit(separation_text, (10, 90))
    
    # Additionally, display details of the first boid for deeper debugging
    if boids:
        boid = boids[0]
        pos_text = font.render(f"Boid0 Pos: ({int(boid.position.x)}, {int(boid.position.y)})", True, (255, 255, 255))
        vel_text = font.render(f"Boid0 Vel: ({int(boid.velocity.x)}, {int(boid.velocity.y)})", True, (255, 255, 255))
        screen.blit(pos_text, (10, 110))
        screen.blit(vel_text, (10, 130))

def main():
    pygame.init()
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    pygame.display.set_caption("Boids Simulation with Debug Menu")
    clock = pygame.time.Clock()

    # Create boids at random positions
    boids = [Boid(random.randint(0, WIDTH), random.randint(0, HEIGHT)) for _ in range(NUM_BOIDS)]
    
    global DEBUG_MODE
    debug_toggle_key = pygame.K_d  # Press D to toggle the debug menu

    running = True
    while running:
        screen.fill((0, 0, 0))
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False
            elif event.type == pygame.KEYDOWN:
                if event.key == debug_toggle_key:
                    DEBUG_MODE = not DEBUG_MODE

        # Update and draw each boid
        for boid in boids:
            boid.flock(boids)
            boid.update()
            boid.edges()
            boid.draw(screen)
        
        # Draw debug panel if debug mode is on
        if DEBUG_MODE:
            draw_debug_panel(screen, boids, clock)
        
        pygame.display.flip()
        clock.tick(60)

    pygame.quit()

if __name__ == "__main__":
    main()
