very simple wave function collapse thingy

wfc.py

from collections.abc import Iterable
import math
import heapq
import random
from collections import deque
import time

type WeightedOptions = list[tuple[int, int]]

UP = 0b1000
RIGHT = 0b0100
DOWN = 0b0010
LEFT = 0b0001


def entropy(weights: Iterable[int]) -> float:
    ws = sum(weights)
    weights = (w / ws for w in weights)
    return -sum(p * math.log2(p) if p > 0 else 0 for p in weights)


class Cell:
    options: WeightedOptions

    def __init__(self, index, options: WeightedOptions):
        self.index = index
        self._options = options
        self._entropy = None
        self.collapsed = False

    @property
    def options(self):
        return self._options

    @options.setter
    def options(self, value):
        self._options = value
        self._entropy = None

    @property
    def entropy(self) -> float:
        if self._entropy is None:
            self._entropy = entropy([w for _, w in self.options])
        return self._entropy

    def collapse(self, wave: 'Wave'):
        self.collapsed = True
        options, weights = zip(*self.options)
        option = wave.rng.choices(options, k=1, weights=weights)[0]
        self.options = [(option, 1)]

    def __lt__(self, other: 'Cell'):
        return self.entropy < other.entropy

    def __repr__(self):
        return (f'<{self.__class__.__name__} collapsed={self.collapsed} '
                f'index={self.index} entropy={self.entropy:.4f}>')


class Wave:
    # entropy, index, possible tiles
    grid: list[Cell]
    # maps index to cell.
    grid_map: dict[int, Cell]

    def __init__(self, width: int, height: int, options: WeightedOptions):
        self.width = width
        self.height = height
        self.options = options
        self.reset()
        self.rng = random.Random()

    def reset(self):
        self.grid = [Cell(i, self.options[:]) for i in range(self.width * self.height)]
        self.grid_map = {c.index: c for c in self.grid}
        heapq.heapify(self.grid)

    def get_minimum_entropy_tile(self) -> Cell | None:
        minvalues = []
        others = []
        heapq.heapify(self.grid)
        while self.grid:
            cell: Cell = heapq.heappop(self.grid)
            if cell.collapsed:
                others.append(cell)
            elif not minvalues:
                minvalues.append(cell)
            elif math.isclose(minvalues[0].entropy, cell.entropy):
                minvalues.append(cell)
            else:
                others.append(cell)
                break

        for cell in others:
            heapq.heappush(self.grid, cell)

        for cell in minvalues:
            heapq.heappush(self.grid, cell)

        if not minvalues:
            return None

        return self.rng.choice(minvalues)

    def collapse(self):
        while True:
            cell = self.get_minimum_entropy_tile()
            if cell is None:
                return
            if not self.collapse_cell(cell):
                return

    def collapse_cell(self, cell: Cell):
        cell.collapse(self)
        return self.propagate(cell)

    def propagate(self, cell: Cell):
        stack = deque([cell])
        while stack:
            cell = stack.popleft()
            for nexti, direction in self.neighbours(cell.index):
                other_cell = self.get(nexti)
                if other_cell.collapsed:
                    continue
                options_before = other_cell.options[:]
                options_after = self.reduce(cell, other_cell, direction)
                if options_before != options_after:
                    other_cell.options = options_after
                    if not other_cell.options:
                        return False
                    stack.append(other_cell)

        return True

    def neighbours(self, index: int):
        y, x = divmod(index, self.width)
        for dx, dy, direction in [(0, -1, UP), (1, 0, RIGHT), (0, 1, DOWN), (-1, 0, LEFT)]:
            nx, ny = dx + x, dy + y
            if self.in_range(nx, ny):
                ni = ny * self.width + nx
                yield ni, direction

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

    def get(self, x_or_i: int, y: int | None = None) -> Cell:
        if y is not None:
            index = y * self.width + x_or_i
        else:
            index = x_or_i

        return self.grid_map[index]

    @staticmethod
    def reduce(a: Cell, b: Cell, direction: int) -> WeightedOptions:
        """
        direction is from a to b. e.g. "UP" means "b" is above "a", so we
        have to compare a's "UP" side, and b's "DOWN" side.

        return new list of weighted options for "b".
        """
        if direction not in (UP, RIGHT, DOWN, LEFT):
            raise ValueError('Invalid value for direction.')

        other_direction = direction >> 2 if direction > 2 else direction << 2

        keep = []
        for tp_b, w in b.options:
            for tp_a, _ in a.options:
                if bool(tp_a & direction) == bool(tp_b & other_direction):
                    keep.append((tp_b, w))
                    break

        return keep


def dbg(lst: WeightedOptions):
    def _dbg_opt(opt):
        v, w = opt
        vals = []
        for s, l in ((UP, 'UP'), (RIGHT, 'RIGHT'), (DOWN, 'DOWN'), (LEFT, 'LEFT')):
            if v & s:
                vals.append(l)
        return '(' + ' | '.join(vals) + f', {w})'

    return f'[' + ', '.join(_dbg_opt(ow) for ow in lst) + ']'


def output(w: Wave):
    boxes: list[list[None | list[str]]] = [[None] * w.width for _ in range(w.height)]
    a = '\033[33;43m##\033[0m'
    b = '\033[32;42m  \033[0m'

    for y in range(w.height):
        for x in range(w.width):
            cell = w.get(x, y)
            if cell.collapsed:
                opt = cell.options[0][0]
                boxes[y][x] = [
                    b + a + b if opt & UP else b + b + b,
                    ''.join([
                        a if opt & LEFT else b,
                        a if opt else b,
                        a if opt & RIGHT else b,
                    ]),
                    b + a + b if opt & DOWN else b + b + b,
                ]
            else:
                boxes[y][x] = ['      ', f'{len(cell.options):^6d}', '      ']

    for y in range(w.height):
        for i in range(3):
            for x in range(w.width):
                print(boxes[y][x][i], end='')
            print()


def main() -> None:
    options = [
        (0, 10),
        (UP | DOWN, 10),
        (LEFT | RIGHT, 10),

        (LEFT | DOWN, 10),
        (LEFT | UP, 10),
        (RIGHT | DOWN, 10),
        (RIGHT | UP, 10),

        # T-cross roads (3-way)
        (UP | RIGHT | DOWN, 5),
        (RIGHT | DOWN | LEFT, 5),
        (DOWN | LEFT | UP, 5),
        (LEFT | UP | RIGHT, 5),

        # 4-way intersection
        (LEFT | RIGHT | UP | DOWN, 2),
    ]

    w = Wave(25, 15, options)

    print('\033[2J')

    while True:
        print('\033[H')
        w.reset()

        x1 = w.rng.choice(range(1, w.width - 1))
        x2 = w.rng.choice(range(1, w.width - 1))

        for x in range(w.width):
            c = w.get(x, 0)
            c.options = [(0 if x1 != x else UP | DOWN, 1)]
            w.collapse_cell(c)

            c = w.get(x, w.height - 1)
            c.options = [(0 if x2 != x else UP | DOWN, 1)]
            w.collapse_cell(c)

        for y in range(w.height):
            c = w.get(0, y)
            c.options = [(0, 1)]
            w.collapse_cell(c)

            c = w.get(w.width - 1, y)
            c.options = [(0, 1)]
            w.collapse_cell(c)

        w.collapse()
        output(w)
        time.sleep(1)


if __name__ == '__main__':
    main()