import pygame as pg
from math import sqrt, ceil, floor
from random import random, randint
from lpgf import Pane

SCREEN_DIMENSIONS = (1500, 1000)
PADDING_CONSTANT = 0.75
FRAME_RATE = 144
N_BOIDS = 250

def clamp(x, min, max) -> int:
    if x < min:
        return min
    if x > max:
        return max
    return x

def surface_color(position: pg.Vector2, surf: pg.Surface) -> tuple[int,int,int]:
    rect = surf.get_rect()
    x = clamp(position.x, 0, rect.width)
    y = clamp(position.y, 0, rect.height)
    return (
        ceil(255 * (x / rect.width)),
        ceil(255 * (y / rect.height)),
        100
    )

def isoceles_triangle(dimensions: tuple[int, int], color = (0,0,0,255)) -> pg.Surface:
    surf = pg.Surface(dimensions, pg.SRCALPHA, 32)
    surf.fill((255,255,255,0))
    pg.draw.polygon(
        surf,
        color,
        (
            (dimensions[0], dimensions[1] / 2),
            (0, 0),
            (0, dimensions[1])
        )
    )
    return surf

class Boid:

    def __init__(
            self,
            position: pg.Vector2,
            velocity: pg.Vector2,
            search_radius: float,
            separation_radius: float,
            max_speed: float,
            cohesion: float,
            alignment: float,
            separation: float,
            bounds: pg.Rect,
            image: pg.Surface = isoceles_triangle((15,10)) #this is the same triangle instance on purpose
    ):
        self.position = position
        self.velocity = velocity
        self._acceleration = pg.Vector2(0,0)
        self._max_speed = max_speed
        self.cohesion = cohesion
        self.alignment = alignment
        self.separation = separation
        self.bounds = bounds
        self._radius_squared = search_radius**2
        self._separation_radius_squared = separation_radius**2
        self._image = image

    @property
    def search_radius(self):
        return sqrt(self._radius_squared)
    
    @search_radius.setter
    def search_radius(self, rad: float):
        self._radius_squared = rad**2

    @property
    def separation_radius(self):
        return sqrt(self._separation_radius_squared)

    @separation_radius.setter
    def separation_radius(self, speed: float):
        self._max_speed_squared = speed**2

    def calculate_next_timestep(self, boids: list['Boid']):
        centroid_target = pg.Vector2(0,0)
        velocity_target = pg.Vector2(0,0)
        nearby = 0

        for b in boids:
            if (self.position-b.position).length_squared() > self._radius_squared or b is self:
                continue
            centroid_target += b.position
            velocity_target += b.velocity
            nearby += 1
            if (self.position - b.position).length_squared() < self._separation_radius_squared:
                self._acceleration += self.separation * (self.position - b.position).normalize()

        if nearby > 0:
            centroid_target /= nearby
            velocity_target /= nearby
            self._acceleration += self.cohesion * (centroid_target - self.position).normalize() 
            self._acceleration += self.alignment * velocity_target.normalize()

        if self.position.y < self.bounds.top:
            self._acceleration += pg.Vector2(0,150)
        elif self.position.y > self.bounds.bottom:
            self._acceleration -= pg.Vector2(0,150)
        
        if self.position.x < self.bounds.left:
            self._acceleration += pg.Vector2(150,0)
        elif self.position.x > self.bounds.right:
            self._acceleration -= pg.Vector2(150,0)
    
    def attract(self, p: pg.Vector2, radius: float, strength: float) -> None:
        diff = p - self.position
        if diff.length_squared() < radius**2:
            self._acceleration += strength * diff.normalize()

    def update(self, dt: float) -> None:
        self.velocity += dt * self._acceleration
        if self.velocity.length_squared() > self._max_speed**2:
            self.velocity.scale_to_length(self._max_speed)
        self.position += self.velocity * dt
        self._acceleration = pg.Vector2(0,0)

    def draw(self, surface: pg.Surface):
        image = self._image.copy() #as per pygame docs, do not continuously rotate an image. Make a copy of a prefab image
        image.fill(surface_color(self.position, surface), special_flags=pg.BLEND_ADD)
        image = pg.transform.rotate(image, self.velocity.angle_to(pg.Vector2(1, 0)))
        surface.blit(image, self.position)

class BoidSim:



    def __init__(self, n_boids: int, simulation_pane: Pane):
        

def main():
    running = True
    pg.init()
    clock = pg.time.Clock()
    
    screen: Pane = Pane.master_pane(SCREEN_DIMENSIONS)
    boid_pane = Pane((SCREEN_DIMENSIONS[0], 0.90 * SCREEN_DIMENSIONS[1]), (0,0))
    bounds_rect = boid_pane.get_rect().scale_by(PADDING_CONSTANT)
    control_pane = Pane((SCREEN_DIMENSIONS[0], 0.1 * SCREEN_DIMENSIONS[1]), (0, 0.9 * SCREEN_DIMENSIONS[1]), blanking_color=(60,60,110))

    screen.add_subpane(boid_pane)
    screen.add_subpane(control_pane)

    boids: list[Boid] = []

    

    while running:
        screen.blank()
        dt = clock.tick(FRAME_RATE) / 1000

        events = pg.event.get()
        mouse_events = []

        for event in events:
            if event.type == pg.QUIT:
                running = False



        for boid in boids:
            boid.calculate_next_timestep(boids)
            boid.update(dt)
            boid.draw(boid_pane)

        screen.draw()
        pg.display.flip()

if __name__ == "__main__":
    main()