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

def dump(struct):
    s = f"{RL.ffi.typeof(struct)}: (".replace('<ctype ', '').replace('>', '')
    for field in dir(struct):
        data = struct.__getattribute__(field)
        if str(data).startswith("<cdata"):
            data = dump(data)
        s += f"{field}:{data} "
    s += ")"
    return s

screen_width = 1024
screen_height = 512
r = 25
k = 30
cell_size = r / math.sqrt(2)
cols = int(screen_width // cell_size)
rows = int(screen_height // cell_size)

@dataclass
class World:
    grid: List[Vec2]
    points: List[Vec2]
    active: List[Vec2]
    pidx = 0
    start = False
    point_count = 0
    frame_count: int = 0
    debug_draw: bool = False

def init() -> World:
    grid = [None] * rows * cols
    points = [Vec2(screen_width // 2, screen_height // 2)]
    # points = []
    active = points.copy()
    for p in points:
        idx = int(p.y // cell_size * cols + p.x // cell_size)
        grid[idx] = p
    return World(grid, points, active)

def player_input(w: World):
    if RL.is_mouse_button_pressed(0):
        p = RL.get_mouse_position()
        w.points.append(p)
        idx = int(p.y // cell_size * cols + p.x // cell_size)
        if idx < len(w.grid):
            w.grid[idx] = p
    if RL.is_key_pressed(RL.KEY_SPACE):
        w.debug_draw = not w.debug_draw
    if RL.is_key_pressed(RL.KEY_ENTER):
        w.start = True

def update(w: World):
    # if w.frame_count % 1 != 0:
    #     return
    if not w.start or not w.points or not w.active:
        return
    for _ in range(25):
        if not w.active:
            break
        point = w.active[-1]
        w.pidx = (w.pidx + 1) % len(w.points)
        found = False
        for i in range(k):
            r_angle = random.random() * math.pi * 2
            r_dist = random.uniform(r, 2 * r)
            new_point = Vec2(point.x + math.cos(r_angle) * r_dist, point.y + math.sin(r_angle) * r_dist)
            new_point_idx = int(new_point.y // cell_size * cols + new_point.x // cell_size)
            if (new_point_idx >= len(w.grid) or w.grid[new_point_idx] is not None
                or new_point.x < 0 or new_point.x > cell_size * cols
                or new_point.y < 0 or new_point.y > cell_size * rows):
                continue
            new_point_row = new_point_idx // cols
            new_point_col = new_point_idx % cols
            collides = False
            for i in range(-2, 3):
                for j in range(-2, 3):
                    if collides:
                        break
                    row = new_point_row + i
                    col = new_point_col + j
                    if row < 0 or row >= rows or col < 0 or col >= cols:
                        continue
                    idx = row * cols + col
                    if w.grid[idx] is not None:
                        # Check distance
                        dist_sqr = RL.vector_2distance_sqr(w.grid[idx], new_point)
                        if dist_sqr < r * r:
                            collides = True

            if collides:
                continue
            w.grid[new_point_idx] = new_point
            w.points.append(new_point)
            w.point_count += 1
            w.active.append(new_point)
            found = True
            break
        if not found:
            w.active.pop()


def draw_2d(w: World):
    # for p in w.points:
    #     prx = p.x // cell_size * cell_size
    #     pry = p.y // cell_size * cell_size
    #     rect = Rect(prx, pry, cell_size, cell_size)
    #     RL.draw_rectangle_rec(rect, Color(230, 230, 230, 255))
    if w.debug_draw:
        for i,c in enumerate(w.grid):
            if c != -1:
                cell_x = (i % cols) * cell_size
                cell_y = (i // cols) * cell_size
                rect = Rect(cell_x, cell_y, cell_size, cell_size)
                RL.draw_rectangle_rec(rect, Color(230, 230, 230, 255))

        for i in range(int((screen_width // cell_size) + 1) - 1):
            for j in range(int((screen_height // cell_size) + 1) - 1):
                rect = Rect(cell_size * i, cell_size * j, cell_size, cell_size)
                RL.draw_rectangle_lines_ex(rect, 2, RL.GREEN)
    for i,p in enumerate(w.points):
        red = min(200, i % 255)
        green = min(200, i * 2 % 255)
        blue = min(200, i * 3 % 255)
        RL.draw_circle_v(p, 2, Color(red, green, blue, 255))


RL.init_window(screen_width, screen_height, "Poisson-disk Sampling");
RL.set_target_fps(60)

w = init()
while not RL.window_should_close():
    player_input(w)
    update(w)

    # Drawing
    RL.begin_drawing()
    RL.clear_background(RL.WHITE)

    draw_2d(w)

    RL.end_drawing()
    w.frame_count += 1