Spatial Hash Plinko Board - Python Pygame

# Created by Anthony Cook & Daniel Pope
# fb.com/groups/pygame


# NOTES:
# Requires Python 2.0.0dev8
# use: pip install pygame==2.0.0dev8
# Use WSAD to move the viewport & UP/DOWN arrow keys to zoom in and out

import pygame
from pygame import gfxdraw
import random
import os
import time
from itertools import product


class SpatialHash:
    def __init__(self):
        self.grid = {}

    def insert_point(self, point):
        x, y = point // 64
        cell = (x, y)
        items = self.grid.get(cell)
        if items is None:
            self.grid[cell] = [point]
        else:
            items.append(point)

    def insert(self, rect, value):
        for cell in self._rect_cells(rect):
            items = self.grid.get(cell)
            if items is None:
                self.grid[cell] = [value]
            else:
                items.append(value)

    def _rect_cells(self, rect):
        x1, y1 = rect.topleft
        x1 //= 64
        y1 //= 64
        x2, y2 = rect.bottomright
        x2 = x2 // 64 + 1
        y2 = y2 // 64 + 1
        return product(range(x1, x2), range(y1, y2))

    def query(self, rect):
        items = []
        for cell in self._rect_cells(rect):
            items.extend(self.grid.get(cell, ()))
        return items

    def query_point(self, pos):
        x, y = pos // 64
        return self.grid.get((x, y), ())


# define a marble class
class marble:
    def __init__(s, x, y):
        # pos is a vector (x, y)
        s.pos = pygame.math.Vector2(x, y)
        # velocity is also a vector with a starting x velocity of random (-1, 1) * 10
        s.velocity = pygame.math.Vector2(random.uniform(-1, 1), 0) * 10

    def update(s):
        global GRAVITY, FPS, STAGE

        # apply velocity to the marble's position
        s.pos += s.velocity

        # apply gravity to the marble's velocity
        s.velocity += GRAVITY / FPS

        # if a marble goes outside of the stage then it reappears
        # on the opposite side. For instance if the marble drops out of the
        # bottom of the stage then it will reappear at the top
        if s.pos.y >= STAGE.bottom:
            s.pos.y = 0
        if s.pos.x > STAGE.right:
            s.pos.x = 0
        if s.pos.x < 0:
            s.pos.x = STAGE.right

# pygame display settings
DR = pygame.Rect((0, 0, 1280, 720)) # Display Rectangle
HDW, HDH = DR.center # H = half
FPS = 60

# set up pygame
pygame.init()
PD = pygame.display.set_mode(DR.size) # primary display based of the size of Display Rect (800, 600)
CLOCK = pygame.time.Clock()

# set strength of gravity
GRAVITY = pygame.math.Vector2(0, 9.8)

# set up stage
SCALE = 10
STAGE = pygame.Rect((0, 0, DR.w * SCALE, DR.h * SCALE))

# Create a spatial hash for broad phase collision detection
spatial = SpatialHash()

# add 1000 randomly placed Plinko pins to quadtree
PIN_COUNT = 1000
pins = []
for index in range(PIN_COUNT):
    pos = pygame.math.Vector2(
        random.randint(0, STAGE.w),
        random.randint(0, STAGE.h)
    )
    bounds = pygame.Rect(pos.x - 30, pos.y - 30, 60, 60)
    spatial.insert(bounds, pos)
    pins.append(pos)

# create 1000 marbles
MARBLE_COUNT = 1000
MARBLE_SPACING = STAGE.w / MARBLE_COUNT
MARBLES = [marble(index * MARBLE_SPACING, 0) for index in range(MARBLE_COUNT)]

# the viewport is like a window that looks onto the stage
# this sets the location and size of the viewport which has a minimum size set to
# that of the primary display surface and starts in the top left of the stage
VIEWPORT = pygame.Rect(DR)

# exit the demo?
exit = False

fps = 0
frames = 1

while True:
    start = time.time()
    # process events
    for e in pygame.event.get():
        if e.type == pygame.QUIT: # window close gadget
            exit = True

    k = pygame.key.get_pressed()
    if k[pygame.K_ESCAPE] or exit: break

    if k[pygame.K_DOWN]:
        # increase the size of the viewport
        VIEWPORT.w += 18
        VIEWPORT.h += 10
    elif k[pygame.K_UP]:
        # decrease the size of the viewport
        VIEWPORT.w -= 18
        VIEWPORT.h -= 10
    # limit the mimium size of the viewport
    # to that of the primary display resolution
    if VIEWPORT.w < DR.w:
        VIEWPORT.w = DR.w
        VIEWPORT.h = DR.h

    if k[pygame.K_a]:
        VIEWPORT.x -= 10
    elif k[pygame.K_d]:
        VIEWPORT.x += 10
    if k[pygame.K_w]:
        VIEWPORT.y -= 10
    elif k[pygame.K_s]:
        VIEWPORT.y += 10
    # limit the viewport to stage's boundry
    if VIEWPORT.right > STAGE.w:
        VIEWPORT.x = STAGE.w - VIEWPORT.w
    if VIEWPORT.x < 0:
        VIEWPORT.x = 0
    if VIEWPORT.bottom > STAGE.h:
        VIEWPORT.y = STAGE.h - VIEWPORT.h
    if VIEWPORT.y < 0:
        VIEWPORT.y = 0

    # clear the primary display (fill it with black)
    PD.fill((0, 0, 0))

    # calculate the scale of the viewport against the size of the primary display
    scale = VIEWPORT.w / DR.w

    # draw all of the Plinko board's pins if they're within the viewport
    cull = VIEWPORT.inflate(20, 20)
    radius = 20 / scale
    if VIEWPORT.w < 200:
        visible_pins = spatial.query(VIEWPORT)
    else:
        visible_pins = [p for p in pins if cull.collidepoint(p)]
    for pin in visible_pins:
        pos = (pin - VIEWPORT.topleft) / scale
        pygame.draw.circle(PD, (255, 0, 0), pos, radius)

    # draw all of the marbles and update marble position
    radius = 6 / scale
    if radius < 1:
        color = (0, round(radius * 255), 0)
        radius = 1
    else:
        color = (0, 255, 0)
    for m in MARBLES:
        if cull.collidepoint(m.pos):
            pos = (m.pos - VIEWPORT.topleft) / scale
            pygame.draw.circle(PD, color, pos, radius)

        # update position of the marble (apply velocity)
        m.update()

        possible_hits = spatial.query_point(m.pos)

        # if there are pins then determine if the marble has collided with them
        for cp in possible_hits:
            # is the distance between the marble and the pin less than
            # their combined radius's?
            if cp.distance_to(m.pos) <= 26:
                # if yes then a collision has occurred and we need to calculate a new
                # trajectory for the marble. This basically the opposite direction to which
                # it was going
                angle = pygame.math.Vector2().angle_to(cp - m.pos)
                m.velocity = m.pos.rotate(angle - 180).normalize() * 10

    # update the primary display
    #end = time.time()
    #frames += 1
    #fps += end - start
    #os.write(1, f"FPS: {1 / (fps / frames)}, {scale}\n".encode('ascii'))
    pygame.display.update()
    CLOCK.tick(FPS) # limit frames
    start = time.time()