Just a fun thing I made

wfc-test.py

#!/usr/bin/env python3
import random
import re
import sys
from collections.abc import Iterator

from enum import IntFlag
from string import ascii_letters, digits


class Direction(IntFlag):

    NONE = 0
    TOP = 1
    RIGHT = 2
    BOTTOM = 4
    LEFT = 8


SINGLE_DIRECTIONS = {Direction.TOP, Direction.RIGHT, Direction.BOTTOM, Direction.LEFT}


def get_opposite_direction(direction) -> Direction:
    if direction == direction.TOP:
        return direction.BOTTOM
    if direction == direction.RIGHT:
        return direction.LEFT
    if direction == direction.BOTTOM:
        return direction.TOP
    if direction == direction.LEFT:
        return direction.RIGHT

    return Direction.NONE


Row = list[None | Direction]
Grid = list[Row]
CHARACTERS = {  # uncomment lines to allow more chars
    # Direction.NONE: " ",
    # Direction.TOP: "╵",
    # Direction.RIGHT: "╶",
    # Direction.BOTTOM: "╷",
    # Direction.LEFT: "╴",
    Direction.TOP | Direction.RIGHT: "└",
    # Direction.TOP | Direction.BOTTOM: "│",
    Direction.TOP | Direction.LEFT: "┘",
    Direction.RIGHT | Direction.BOTTOM: "┌",
    # Direction.RIGHT | Direction.LEFT: "─",
    Direction.BOTTOM | Direction.LEFT: "┐",
    Direction.TOP | Direction.RIGHT | Direction.BOTTOM: "├",
    Direction.TOP | Direction.RIGHT | Direction.LEFT: "┴",
    Direction.TOP | Direction.BOTTOM | Direction.LEFT: "┤",
    Direction.RIGHT | Direction.BOTTOM | Direction.LEFT: "┬",
    # Direction.TOP | Direction.RIGHT | Direction.BOTTOM | Direction.LEFT: "┼",
}


def generate_neighbors(
    x: int, y: int, width: int, height: int
) -> dict[Direction, tuple[int, int]]:
    neighbors: dict[Direction, tuple[int, int]] = {}
    if x > 0:
        neighbors[Direction.LEFT] = (x - 1, y)
    if x < width - 1:
        neighbors[Direction.RIGHT] = (x + 1, y)
    if y > 0:
        neighbors[Direction.TOP] = (x, y - 1)
    if y < height - 1:
        neighbors[Direction.BOTTOM] = (x, y + 1)
    return neighbors


def coord_iterator(
    width: int, height: int, grid: Grid, border: bool
) -> Iterator[tuple[int, int]]:
    def entropy(_x: int, _y: int) -> int:
        valid_items = get_valid_items(grid, width, height, _x, _y, border=border)
        return len(valid_items) if valid_items else 0

    coords: dict[tuple[int, int], int] = {}

    for x in range(width):
        for y in range(height):
            coords[(x, y)] = entropy(x, y)

    current: tuple[int, int]
    while coords:
        sorted_ = sorted((ent, key) for key, ent in coords.items() if ent)
        if not sorted_:
            return

        current = sorted_.pop(0)[1]
        del coords[current]
        yield current
        for neighbor in generate_neighbors(
            *current, width=width, height=height
        ).values():
            if neighbor in coords:
                coords[neighbor] = entropy(*neighbor)


def get_valid_items(
    grid: Grid, width: int, height: int, x: int, y: int, border: bool
) -> tuple[Direction, ...] | None:
    connected = set()
    unconnected = set()
    neighbors = generate_neighbors(x, y, width, height)
    for dir_, (nx, ny) in neighbors.items():
        if (neighbor := grid[ny][nx]) is None:
            continue
        if get_opposite_direction(dir_) in neighbor:
            connected.add(dir_)
        else:
            unconnected.add(dir_)

    if border and len(neighbors) < len(SINGLE_DIRECTIONS):
        unconnected |= SINGLE_DIRECTIONS ^ set(neighbors.keys())

    if not connected and not unconnected:
        return tuple(CHARACTERS.keys())

    possible = [
        key
        for key in CHARACTERS
        if (
            (not unconnected or not any(spam in key for spam in unconnected))
            and (not connected or all(eggs in key for eggs in connected))
        )
    ]

    if not possible:
        return None  # bad

    return tuple(possible)


def row_as_str(row: Row) -> str:
    return "".join(("X" if dir_ is None else CHARACTERS[dir_]) for dir_ in row)


def grid_as_str(grid: Grid) -> str:
    return "\n".join(row_as_str(row) for row in grid)


def fix_missing(grid: Grid, width: int, height: int, border: bool) -> bool:
    change_counter = 0
    while any(None in row for row in grid):
        for y, row in enumerate(grid):
            while None in row:
                if change_counter > width * height:
                    return False
                x = row.index(None)
                grid[y][x] = random.choice(tuple(CHARACTERS.keys()))
                neighbors = generate_neighbors(x, y, width, height)
                for nx, ny in neighbors.values():
                    grid[ny][nx] = random.choice(
                        get_valid_items(grid, width, height, nx, ny, border) or (None,)
                    )
                change_field(grid, width, height, x, y, border)
                change_counter += 1
    return True


def change_field(
    grid: Grid, width: int, height: int, x: int, y: int, border: bool
) -> bool:
    valid_items = get_valid_items(grid, width, height, x, y, border)
    if valid_items:
        grid[y][x] = random.choice(valid_items)
    return bool(valid_items)


def create_grid(width: int, height: int, border: bool) -> Grid:
    grid: Grid = [[None] * width for _ in range(height)]

    for x, y in coord_iterator(width, height, grid, border):
        valid_items = get_valid_items(grid, width, height, x, y, border)
        if valid_items:
            grid[y][x] = random.choice(valid_items)

    fix_missing(grid, width, height, border)  # sometimes necessary

    return grid


def main() -> int | str:
    width, height = 20, 5

    border, print_command = False, False

    args = sys.argv[1:]
    if "--border" in args:
        args.remove("--border")
        border = True
    if "--print-command" in args:
        args.remove("--print-command")
        print_command = True
    seed: None | str = None
    if args:
        if "--help" in args or "-h" in args or "?" == args[0]:
            print(
                f"Usage: {sys.argv[0]} WIDTHxHEIGHT? --border? --print-command? seed..."
            )
            return 0
        if re.fullmatch(r"^\d+x\d+$", args[0]):
            width, height = [max(3, int(spam)) for spam in args[0].split("x")]
            seed = " ".join(args[1:])
        else:
            seed = " ".join(args)
    if seed:
        if print_command:
            print(" ".join(sys.argv))
    else:
        seed = "".join([random.choice(ascii_letters + digits) for _ in range(12)])
        if print_command:
            print(" ".join(sys.argv + [seed]))
    random.seed(seed)

    grid = create_grid(width, height, border)

    print(grid_as_str(grid))
    return 0


if __name__ == "__main__":
    sys.exit(main())