silly-noita-in-python.py

import math
import random
from typing import NamedTuple, Literal, Callable

import pygame
from PIL import Image
import numpy as np

type Pos = tuple[int, int]
type Color = tuple[int, int, int] | str
type BlockType = Literal['solid', 'sand', 'liquid']


class OutOfBounds:
    pass


class View(NamedTuple):
    x: int
    y: int
    width: int
    height: int


class Block:
    def __init__(self, color: Color, type: BlockType = 'solid'):
        self.color = color
        self.type = type


class World:
    def __init__(self, width: int, height: int):
        self.width = width
        self.height = height
        self.grid: list[list[None | Block]] = [[None] * self.width for _ in range(self.height)]
        self.build_world()

    def build_world(self):
        for x in range(self.width):
            self.grid[self.height - 1][x] = Block('white')
            self.grid[self.height - 2][x] = Block('white')

        for y in range(3, 20):
            self.grid[self.height - y][0] = Block('white')
            self.grid[self.height - y][self.width - 1] = Block('white')

    def spawn(self, factory: Callable[[], Block], pos: Pos, radius: int = 0):
        if radius < 1:
            raise ValueError('radius must be 1 or greater')

        x, y = pos
        if radius == 0:
            if self.get(x, y) is None:
                self.set(x, y, factory())
            return

        for dy in range(-radius, radius + 1):
            for dx in range(-radius, radius + 1):
                if math.dist((0, 0), (dx / radius, dy / radius)) <= 1.0:
                    nx, ny = x + dx, y + dy
                    if self.get(nx, ny) is None:
                        self.set(nx, ny, factory())

    def in_range(self, x: int, y: int):
        return 0 <= x < self.width and 0 <= y < self.height

    def get(self, x, y) -> None | Block | type[OutOfBounds]:
        if not self.in_range(x, y):
            return OutOfBounds
        return self.grid[y][x]

    def set(self, x, y, value: None | Block):
        self.grid[y][x] = value

    def update_blocks(self):
        sand_deltas = [(0, 1), (-1, 1), (1, 1)]
        water_deltas = [(0, 1), (-1, 1), (1, 1), (-1, 0), (1, 0)]

        skip = set()
        for y in range(self.height - 1, -1, -1):
            for x in range(self.width):
                if (x, y) in skip:
                    continue
                block = self.grid[y][x]
                if block is None or block.type == 'solid':
                    continue

                if block.type == 'sand':
                    a = sand_deltas[0:1]
                    b = sand_deltas[1:3]
                    random.shuffle(b)
                    deltas = a + b
                elif block.type == 'liquid':
                    a = water_deltas[0:1]
                    b = water_deltas[1:3]
                    c = water_deltas[3:5]
                    random.shuffle(b)
                    random.shuffle(c)
                    deltas = a + b + c
                else:
                    continue

                for dx, dy in deltas:
                    other_block = self.get(x + dx, y + dy)
                    if other_block is OutOfBounds:
                        continue
                    if other_block is None or block.type == 'sand' and other_block.type == 'liquid':
                        nx, ny = x + dx, y + dy
                        skip.add((nx, ny))
                        self.set(nx, ny, block)
                        self.set(x, y, other_block)

                        # if they only move left and rigth, there is a slight chance they disappear.
                        if dy == 0 and random.randint(0, 100) == 0:
                            self.set(nx, ny, None)

                        break

    def render(self, surf: pygame.Surface):
        for y in range(self.height - 1, -1, -1):
            for x in range(self.width):
                block = self.grid[y][x]
                if block is None:
                    continue
                surf.set_at((x, y), block.color)


def main() -> None:
    pygame.init()

    screen = pygame.display.set_mode((512, 512))
    clock = pygame.Clock()

    surf = pygame.Surface((128, 128))
    world = World(*surf.get_size())

    while True:
        delta = clock.tick(60) / 1000
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                return

            if event.type == pygame.MOUSEBUTTONDOWN:
                mx, my = event.pos
                mx = int(mx / (screen.get_width() / world.width))
                my = int(my / (screen.get_height() / world.height))
                if event.button == 1:
                    world.spawn(lambda: Block('yellow', 'sand'), (mx, my), 2)
                elif event.button == 3:
                    world.spawn(lambda: Block('blue', 'liquid'), (mx, my), 2)

        world.update_blocks()

        surf.fill('black')
        world.render(surf)
        screen.blit(pygame.transform.scale(surf, screen.get_size(), screen), (0, 0))
        pygame.display.flip()


if __name__ == '__main__':
    main()