pf1e-simulator/src/pf1e_simulator/grid.py

"""Square grid: PF1e distances, movement costs, and Dijkstra pathing.

Diagonal movement alternates 1-2-1-2 squares (the 5-10-5-10 ft rule), so the
Dijkstra state carries the parity of diagonals used. A diagonal step into
difficult terrain costs a flat 3 squares and still counts as a diagonal for
the alternation (table convention). A diagonal step requires both flanking
orthogonal squares to be terrain-passable (no cutting wall corners); creature
occupancy never blocks a diagonal, and occupied squares cannot be entered in
Phase 0 (moving through allies is a later refinement).
"""

from __future__ import annotations

import heapq
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from pf1e_simulator.map import MapSpec, Pos, TerrainType

_DIFFICULT_TERRAIN_COST = 2  # rubble, water, etc. count as 2 squares (PF1e)

_DIRECTIONS: tuple[Pos, ...] = (
    (-1, -1),
    (-1, 0),
    (-1, 1),
    (0, -1),
    (0, 1),
    (1, -1),
    (1, 0),
    (1, 1),
)


class Grid:
    """Terrain view of a MapSpec with movement and distance queries.

    An optional elevation layer stores a height per cell; climbing one level
    costs one extra square of movement (PF1e), descending is free.
    """

    def __init__(
        self,
        cells: tuple[tuple[TerrainType, ...], ...],
        elevation: tuple[tuple[int, ...], ...] | None = None,
    ) -> None:
        self._cells = cells
        self._height = len(cells)
        self._width = len(cells[0]) if cells else 0
        self._elevation = elevation

    @classmethod
    def from_spec(cls, spec: MapSpec) -> Grid:
        cells = tuple(tuple(spec.legend[ch] for ch in row) for row in spec.terrain)
        elevation = None
        if spec.elevation is not None:
            height = len(cells)
            width = len(cells[0]) if cells else 0
            if len(spec.elevation) != height or (
                height and len(spec.elevation[0]) != width
            ):
                msg = "elevation layer dimensions do not match terrain"
                raise ValueError(msg)
            elevation = tuple(
                tuple(0 if ch == "." else int(ch) for ch in row)
                for row in spec.elevation
            )
        return cls(cells, elevation=elevation)

    @property
    def height(self) -> int:
        return self._height

    @property
    def width(self) -> int:
        return self._width

    def in_bounds(self, pos: Pos) -> bool:
        row, col = pos
        return 0 <= row < self._height and 0 <= col < self._width

    def terrain_at(self, pos: Pos) -> TerrainType:
        row, col = pos
        return self._cells[row][col]

    def passable(self, pos: Pos) -> bool:
        return self.in_bounds(pos) and self.terrain_at(pos).move_cost is not None

    def elevation_at(self, pos: Pos) -> int:
        """Height of a cell; 0 when no elevation layer or out of bounds."""
        if self._elevation is None or not self.in_bounds(pos):
            return 0
        row, col = pos
        return self._elevation[row][col]

    def distance(self, a: Pos, b: Pos) -> int:
        """Distance in squares under the 5-10-5-10 diagonal rule."""
        d_row = abs(a[0] - b[0])
        d_col = abs(a[1] - b[1])
        return max(d_row, d_col) + min(d_row, d_col) // 2

    def step_cost(self, frm: Pos, to: Pos, diagonals_used: int) -> int:
        """Cost in squares of one step; diagonals_used is the count so far.

        Climbing one elevation level adds one square per level; descending is
        free and never reduces the terrain cost.
        """
        cost = self.terrain_at(to).move_cost
        if cost is None:
            msg = f"cannot step into impassable square {to}"
            raise ValueError(msg)
        is_diagonal = frm[0] != to[0] and frm[1] != to[1]
        if not is_diagonal:
            return cost + max(0, self.elevation_at(to) - self.elevation_at(frm))
        if cost >= _DIFFICULT_TERRAIN_COST:
            return 3 + max(0, self.elevation_at(to) - self.elevation_at(frm))
        return 1 + (diagonals_used % 2) + max(
            0, self.elevation_at(to) - self.elevation_at(frm)
        )

    def diagonal_allowed(self, frm: Pos, to: Pos) -> bool:
        flank_a = (frm[0], to[1])
        flank_b = (to[0], frm[1])
        return self.passable(flank_a) and self.passable(flank_b)

    def reachable(self, start: Pos, budget: int | None, blocked: frozenset[Pos]) -> dict[Pos, int]:
        """Cheapest movement cost per reachable square, capped at budget.

        `blocked` holds creature-occupied squares: unenterable, but never
        affecting diagonal_allowed. The start square is included at cost 0.
        budget=None disables the cap (full cost field).
        """
        best_per_state: dict[tuple[Pos, int], int] = {(start, 0): 0}
        results: dict[Pos, int] = {start: 0}
        heap: list[tuple[int, Pos, int]] = [(0, start, 0)]
        while heap:
            cost, pos, parity = heapq.heappop(heap)
            if cost > best_per_state.get((pos, parity), cost):
                continue
            row, col = pos
            for d_row, d_col in _DIRECTIONS:
                nxt = (row + d_row, col + d_col)
                if nxt == start or nxt in blocked or not self.passable(nxt):
                    continue
                is_diagonal = d_row != 0 and d_col != 0
                if is_diagonal and not self.diagonal_allowed(pos, nxt):
                    continue
                new_cost = cost + self.step_cost(pos, nxt, parity)
                if budget is not None and new_cost > budget:
                    continue
                new_parity = (parity + 1) % 2 if is_diagonal else parity
                state = (nxt, new_parity)
                known = best_per_state.get(state)
                if known is None or new_cost < known:
                    best_per_state[state] = new_cost
                    current = results.get(nxt)
                    if current is None or new_cost < current:
                        results[nxt] = new_cost
                    heapq.heappush(heap, (new_cost, nxt, new_parity))
        return results