wordsearch/app/generator.py

"""Grid building and word placement."""

from __future__ import annotations

import random
import string
from dataclasses import dataclass, field

from .normaliser import Normalised, normalise_all

# Direction vectors: (drow, dcol)
RIGHT = (0, 1)
DOWN = (1, 0)
DOWN_RIGHT = (1, 1)
UP_RIGHT = (-1, 1)
LEFT = (0, -1)
UP = (-1, 0)
UP_LEFT = (-1, -1)
DOWN_LEFT = (1, -1)

PLACEMENT_RETRIES = 200


@dataclass
class Placement:
    word: Normalised
    row: int
    col: int
    direction: tuple[int, int]


@dataclass
class Puzzle:
    grid: list[list[str]]
    placements: list[Placement]
    skipped: list[Normalised] = field(default_factory=list)
    filtered_out: list[Normalised] = field(default_factory=list)
    warnings: list[str] = field(default_factory=list)
    diagonal: bool = False
    reversed_: bool = False
    allow_overlap: bool = False

    @property
    def size(self) -> int:
        return len(self.grid)

    @property
    def option_labels(self) -> list[str]:
        labels = []
        if self.diagonal:
            labels.append("Diagonal")
        if self.reversed_:
            labels.append("Reversed")
        if self.allow_overlap:
            labels.append("Overlapping")
        return labels


def active_directions(diagonal: bool, reversed_: bool) -> list[tuple[int, int]]:
    base = [RIGHT, DOWN]
    if diagonal:
        base += [DOWN_RIGHT, UP_RIGHT]
    if reversed_:
        base = base + [(-dr, -dc) for (dr, dc) in base]
    return base


def _can_place(grid: list[list[str]], word: str, row: int, col: int,
               direction: tuple[int, int], allow_overlap: bool) -> bool:
    size = len(grid)
    dr, dc = direction
    for i, ch in enumerate(word):
        r, c = row + dr * i, col + dc * i
        if not (0 <= r < size and 0 <= c < size):
            return False
        existing = grid[r][c]
        if existing == "":
            continue
        if allow_overlap and existing == ch:
            continue
        return False
    return True


def _commit(grid: list[list[str]], word: str, row: int, col: int,
            direction: tuple[int, int]) -> None:
    dr, dc = direction
    for i, ch in enumerate(word):
        grid[row + dr * i][col + dc * i] = ch


def _start_bounds(size: int, length: int,
                  direction: tuple[int, int]) -> tuple[range, range]:
    """Valid starting (row, col) ranges so the whole word fits."""
    dr, dc = direction
    if dr >= 0:
        rows = range(0, size - (length - 1) * dr)
    else:
        rows = range((length - 1) * (-dr), size)
    if dc >= 0:
        cols = range(0, size - (length - 1) * dc)
    else:
        cols = range((length - 1) * (-dc), size)
    return rows, cols


def generate(
    *,
    words: list[str],
    grid_size: int,
    word_count: int,
    min_length: int,
    max_length: int,
    diagonal: bool,
    reversed_: bool,
    allow_overlap: bool,
    rng: random.Random | None = None,
) -> Puzzle:
    if rng is None:
        rng = random.Random()

    if grid_size < 5 or grid_size > 25:
        raise ValueError("grid size must be between 5 and 25")
    if word_count < 1 or word_count > 30:
        raise ValueError("word count must be between 1 and 30")
    if min_length < 1:
        raise ValueError("min length must be at least 1")
    if max_length < min_length:
        raise ValueError("max length must be greater than or equal to min length")

    effective_max = min(max_length, grid_size)

    normalised = normalise_all(words)
    eligible: list[Normalised] = []
    filtered_out: list[Normalised] = []
    skipped: list[Normalised] = []
    for n in normalised:
        if n.skipped:
            skipped.append(n)
            continue
        length = len(n.grid)
        if min_length <= length <= effective_max:
            eligible.append(n)
        else:
            filtered_out.append(n)

    rng.shuffle(eligible)
    directions = active_directions(diagonal, reversed_)

    grid: list[list[str]] = [["" for _ in range(grid_size)]
                              for _ in range(grid_size)]
    placements: list[Placement] = []
    warnings: list[str] = []

    for word in eligible:
        if len(placements) >= word_count:
            break
        placed = False
        for _ in range(PLACEMENT_RETRIES):
            direction = rng.choice(directions)
            rows, cols = _start_bounds(grid_size, len(word.grid), direction)
            if not rows or not cols:
                continue
            row = rng.choice(rows)
            col = rng.choice(cols)
            if _can_place(grid, word.grid, row, col, direction, allow_overlap):
                _commit(grid, word.grid, row, col, direction)
                placements.append(Placement(word=word, row=row, col=col,
                                            direction=direction))
                placed = True
                break
        if not placed:
            warnings.append(
                f"could not place word {word.grid!r} after "
                f"{PLACEMENT_RETRIES} attempts; skipping"
            )

    if filtered_out:
        warnings.append(
            f"{len(filtered_out)} word(s) excluded by length filter "
            f"(min={min_length}, max={effective_max})"
        )

    if len(placements) < word_count:
        warnings.append(
            f"asked for {word_count} word(s); placed {len(placements)}"
        )

    # Fill remaining empty cells with random uppercase letters.
    for r in range(grid_size):
        for c in range(grid_size):
            if grid[r][c] == "":
                grid[r][c] = rng.choice(string.ascii_uppercase)

    return Puzzle(grid=grid, placements=placements, skipped=skipped,
                  filtered_out=filtered_out, warnings=warnings,
                  diagonal=diagonal, reversed_=reversed_,
                  allow_overlap=allow_overlap)