164 lines
4.3 KiB
Python
164 lines
4.3 KiB
Python
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()
|
|
|