Plinko game by Anthony Cook, adapted to use a Spatial Hash

plinko-spatial.py

# Created by Anthony Cook
# fb.com/groups/pygame

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
        cell = x // 64, y // 64
        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
        return self.grid.get((x // 64, y // 64), ())


# 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

start = time.time()
while True:
    # process events
    for e in pygame.event.get():
        if e.type == pygame.QUIT: # window close gadget
            exit = True
        elif e.type == pygame.MOUSEBUTTONDOWN:
            if e.button == 5: # mouse wheel down
                # increase the size of the viewport
                VIEWPORT.w += 18
                VIEWPORT.h += 10
            elif e.button == 4: # mouse wheel 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

    # exit the demo if ESC pressed or exit is True (set by pressing window x gadget)
    if pygame.key.get_pressed()[pygame.K_ESCAPE] or exit: break

    # get the distance the mouse has travelled since last get_rel() call
    rx, ry = pygame.mouse.get_rel()
    # if left mouse button has been pressed then you can drag the viewport about
    if pygame.mouse.get_pressed()[0]:
        # move the viewport
        VIEWPORT.x -= rx
        VIEWPORT.y -= ry
        # 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()
    os.write(1, f"FPS: {1 / (end - start)}\n".encode('ascii'))
    pygame.display.update()
    CLOCK.tick(FPS) # limit frames
    start = time.time()