import pyray as RL
from pyray import (Rectangle as Rect)
import math
import pdb
import random
from typing import Optional, Tuple, List
from dataclasses import dataclass, field

screen_width = 1280
screen_height = 1024

ball_r = 6
ball_speed = 3.5
num_balls = 1000
qt_capacity = 4

@dataclass
class Ball:
    px: float
    py: float
    vx: float
    vy: float

@dataclass
class QNode:
    aabb: Rect
    points: List = field(default_factory=list)

@dataclass
class Quadtree:
    node: QNode
    subdivided = False
    nw: Optional['Quadtree'] = None
    ne: Optional['Quadtree'] = None
    sw: Optional['Quadtree'] = None
    se: Optional['Quadtree'] = None

@dataclass
class World:
    balls = []
    qt = {}
    tick = 0
    paused = False
    mouse_clicks = []

w = World()

def qt_split(qt: Quadtree):
    x, y, hw, hh = qt.node.aabb.x, qt.node.aabb.y, qt.node.aabb.width * 0.5, qt.node.aabb.height * 0.5
    nw = Rect(x     , y     , hw, hh)
    ne = Rect(x + hw, y     , hw, hh)
    sw = Rect(x     , y + hh, hw, hh)
    se = Rect(x + hw, y + hh, hw, hh)
    qt.nw = Quadtree(QNode(nw))
    qt.ne = Quadtree(QNode(ne))
    qt.sw = Quadtree(QNode(sw))
    qt.se = Quadtree(QNode(se))
    qt.subdivided = True
    
def qt_insert(qt: Quadtree, p):
    if not RL.check_collision_point_rec(p, qt.node.aabb):
        return False
    if qt.subdivided:
        inserted = False
        if not inserted: inserted = qt_insert(qt.nw, p)
        if not inserted: inserted = qt_insert(qt.ne, p)
        if not inserted: inserted = qt_insert(qt.sw, p)
        if not inserted: inserted = qt_insert(qt.se, p)
        return inserted
    if len(qt.node.points) + 1 >= qt_capacity:
        qt_split(qt)
        qt.node.points.append(p)
        inserted = False
        for p in qt.node.points:
            if   qt_insert(qt.nw, p): pass
            elif qt_insert(qt.ne, p): pass
            elif qt_insert(qt.sw, p): pass
            elif qt_insert(qt.se, p): pass
        qt.node.points.clear()
        return True
    else:
        qt.node.points.append(p)
        return True

def construct_quadtree(points):
    root_node = QNode(Rect(0, 0, screen_width, screen_height))
    qt = Quadtree(root_node)
    for p in points:
        qt_insert(qt, p)
    return qt

def rect_values(r: Rect):
    return r.x, r.y, r.w, r.h

def init():
    for n in range(num_balls):
        px = random.randrange(ball_r, 50)
        py = random.randrange(ball_r, 50)
        # px = random.randrange(ball_r, screen_width  - ball_r)
        # py = random.randrange(ball_r, screen_height - ball_r)
        angle = random.uniform(0, 360)
        vx = math.cos(angle) * ball_speed * random.uniform(1, 3)
        vy = math.sin(angle) * ball_speed * random.uniform(1, 3)
        w.balls.append(Ball(px, py, vx, vy))

def player_input():
    if RL.is_key_pressed(RL.KEY_SPACE):
        w.paused = not w.paused
    if RL.is_mouse_button_pressed(0):
        print(RL.get_mouse_position())
        w.mouse_clicks.append(RL.get_mouse_position())
    
def update():
    # Recontruct quadtree
    if w.paused:
        return
    points = []
    for b in w.balls:
        points.append((b.px, b.py))
    w.qt = construct_quadtree(points)
    
    for ball in w.balls:
        ball.px += ball.vx
        ball.py += ball.vy
        if ball.px - ball_r <= 0 or ball.px + ball_r >= screen_width:
            # Reset position to make sure it's clamped
            ball.px = RL.clamp(ball.px, ball_r + 0.1, screen_width - ball_r - 0.1)
            ball.vx *= -1
        if ball.py - ball_r <= 0 or ball.py + ball_r > screen_height:
            # Reset position to make sure it's clamped
            ball.py = RL.clamp(ball.py, ball_r + 0.1, screen_height - ball_r - 0.1)
            ball.vy *= -1

def draw_qt_dfs(qt: Quadtree):
    if not qt:
        return
    draw_qt_dfs(qt.nw)
    draw_qt_dfs(qt.ne)
    draw_qt_dfs(qt.se)
    draw_qt_dfs(qt.sw)
    RL.draw_rectangle_lines_ex(qt.node.aabb, 0.5, RL.BLACK)

def draw():
    RL.begin_drawing()

    RL.clear_background(RL.WHITE)
    draw_qt_dfs(w.qt)
    for ball in w.balls:
        RL.draw_circle_lines_v((ball.px, ball.py), ball_r, RL.BLACK)
    for mc in w.mouse_clicks:
        RL.draw_circle_v(mc, 5, RL.RED)

    RL.end_drawing()

RL.init_window(screen_width, screen_height, "Quadtree");
RL.set_target_fps(60)

init()
while not RL.window_should_close():
    player_input()
    update()
    draw()