colors.py

"""
An example of the determinism of pymunk by coloring balls according to their 
position, and then respawning them to verify each ball ends up in the same 
place. Inspired by Pymunk user Nam Dao.
"""


import os
import random

import pygame

import pymunk
import pymunk.pygame_util


def new_space():
    space = pymunk.Space()
    space.gravity = 0, 900
    static_body = space.static_body
    walls = [
        pymunk.Segment(static_body, (0, -500), (600, -500), 20),
        pymunk.Segment(static_body, (20, -500), (0, 600), 20),
        pymunk.Segment(static_body, (0, 600), (600, 600), 20),
        pymunk.Segment(static_body, (600, 600), (580, -500), 20),
    ]
    for wall in walls:
        wall.elasticity = 0.9

    space.add(*walls)

    random.seed(0)
    return space


def set_colors(color_dict, logo_img, space):
    color_dict.clear()
    w = logo_img.get_width()
    h = logo_img.get_height()
    logo_img.lock()
    for shape in space.shapes:
        if not isinstance(shape, pymunk.Circle):
            continue
        r = shape.body.position.x / 600 * 255
        g = max((shape.body.position.y - 400) / 200 * 255, 0)
        if r < 0 or r > 255 or g < 0 or g > 255:
            print(shape.body.position)
            exit()
        shape.color = (r, g, 150, 255)
        p = shape.body.position
        x = int(p.x) - (600 - w) // 2
        y = int(p.y - 600 + h + 10)
        if x >= 0 and x < w and y > 0 and y < h:
            color = logo_img.get_at([x, y])
            if color.a > 200:
                shape.color = color.r, color.g, color.b, 255
        color_dict[shape.data] = shape.color
    logo_img.unlock()


pygame.init()
screen = pygame.display.set_mode((600, 600))
clock = pygame.time.Clock()
font = pygame.font.SysFont("Arial", 14)
text = font.render(
    "Press r to reset and respawn all balls."
    " Press c to set color of each ball according to its position.",
    True,
    pygame.Color("darkgray"),
)

logo_img = pygame.image.load(
    os.path.join(os.path.dirname(os.path.abspath(__file__)), "pymunk_logo_sphinx.png")
)
draw_options = pymunk.pygame_util.DrawOptions(screen)

color_dict = {}
space = new_space()
cnt = max_cnt = 6000

while True:
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            exit()
        elif event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE:
            exit()
        elif event.type == pygame.KEYDOWN and event.key == pygame.K_p:
            pygame.image.save(screen, "colors.png")
        elif event.type == pygame.KEYDOWN and event.key == pygame.K_r:
            if not color_dict:
                set_colors(color_dict, logo_img, space)
            space = new_space()
            cnt = max_cnt
        elif event.type == pygame.KEYDOWN and event.key == pygame.K_c:
            set_colors(color_dict, logo_img, space)

    if cnt > 0:
        for _ in range(15):
            cnt -= 1
            body = pymunk.Body()
            x = random.randint(550, 570)
            y = random.randint(30, 100)
            body.position = x, y
            shape = pymunk.Circle(body, 2.5)
            shape.mass = 1
            shape.data = cnt
            shape.elasticity = 0.9
            if color_dict != None and cnt in color_dict:
                shape.color = color_dict[cnt]
            space.add(body, shape)

    ### Update physics
    dt = 1.0 / 60.0
    for _ in range(1):
        space.step(dt / 1)

    ### Draw stuff
    screen.fill(pygame.Color("white"))

    color = pygame.Color("blue")
    for shape in space.shapes:
        if not isinstance(shape, pymunk.Circle):
            continue
        if hasattr(shape, "color"):
            color = shape.color
        # Draw the circles as little squares to make the end result nicer.
        d = int(shape.radius * 2) + 1
        pygame.draw.rect(
            screen,
            color,
            (
                shape.body.position.x - 1,
                shape.body.position.y - 1,
                d,
                d,
            ),
        )
    screen.blit(text, (25, 2))

    ### Flip screen
    pygame.display.flip()
    clock.tick(50)
    pygame.display.set_caption("fps: " + str(clock.get_fps()))