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

screen_width = 1200
screen_height = 960

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[Tuple[float, float]] = field(default_factory=list)

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

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

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), parent=qt, direction='NW')
    qt.ne = Quadtree(QNode(ne), parent=qt, direction='NE')
    qt.sw = Quadtree(QNode(sw), parent=qt, direction='SW')
    qt.se = Quadtree(QNode(se), parent=qt, direction='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 qt_find_nearest_point(qt: Quadtree, point) -> Tuple[float, float]:
    containing_qt = qt
    while containing_qt.subdivided:
        if RL.check_collision_point_rec(point, containing_qt.nw.node.aabb):
            containing_qt = containing_qt.nw
        elif RL.check_collision_point_rec(point, containing_qt.ne.node.aabb):
            containing_qt = containing_qt.ne
        elif RL.check_collision_point_rec(point, containing_qt.sw.node.aabb):
            containing_qt = containing_qt.sw
        elif RL.check_collision_point_rec(point, containing_qt.se.node.aabb):
            containing_qt = containing_qt.se

    def search_for_nearest(qt: Quadtree, direction = ''):
        nonlocal closest_point, closest_dist
        contains_point = RL.check_collision_point_rec(point, qt.node.aabb)
        if not contains_point:
            px, py = point.x, point.y
            dx, dy = 0,0
            if px < qt.node.aabb.x:
                dx = qt.node.aabb.x - px
            elif px > qt.node.aabb.x + qt.node.aabb.width:
                dx = px - (qt.node.aabb.x + qt.node.aabb.width)
            if py < qt.node.aabb.y:
                dy = qt.node.aabb.y - py
            elif py > qt.node.aabb.y + qt.node.aabb.height:
                dy = py - (qt.node.aabb.y + qt.node.aabb.height)
            dist = RL.vector2_length(Vec2(dx, dy))
            if dist >= closest_dist:
                return
        if qt.subdivided:
            if direction != 'NW': search_for_nearest(qt.nw)
            if direction != 'NE': search_for_nearest(qt.ne)
            if direction != 'SW': search_for_nearest(qt.sw)
            if direction != 'SE': search_for_nearest(qt.se)

        w.visited_quadrants.append(qt)
        for p in qt.node.points:
            d = RL.vector_2distance(point, Vec2(p[0], p[1]))
            if d < closest_dist:
                closest_point = p
                closest_dist = d
    closest_point = None
    closest_dist = float('inf')
    previous_direction = ''
    while containing_qt is not None:
        search_for_nearest(containing_qt, previous_direction)
        previous_direction = containing_qt.direction
        containing_qt = containing_qt.parent

    return closest_point
                

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):
        mouse_pos = RL.get_mouse_position()
        nearest = qt_find_nearest_point(w.qt, mouse_pos)
        w.nearest_pairs.append((mouse_pos, nearest))
        w.paused = True
    if RL.is_key_pressed(RL.KEY_ENTER):
        w.current_visited += 1
    
def update():
    # Recontruct quadtree
    if w.paused:
        return
    
    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

    points = []
    for b in w.balls:
        points.append((b.px, b.py))
    w.qt = construct_quadtree(points)

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 i in range(w.current_visited):
        RL.draw_rectangle_rec(w.visited_quadrants[i].node.aabb, RL.LIGHTGRAY)
    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)
    for np in w.nearest_points:
        RL.draw_circle_lines_v(Vec2(np[0], np[1]), ball_r, RL.GREEN)
    for mc,(px,py) in w.nearest_pairs:
        pos = Vec2(px, py)
        RL.draw_circle_v(mc, 5, RL.RED)
        RL.draw_circle_lines_v(pos, ball_r, RL.GREEN)
        RL.draw_line_v(mc, pos, RL.BLUE)

    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()