"""Bitboard engine for Battlesnake grid spatial operations.

Cell index = y * width + x.  Bit *i* of a Python int represents cell *i*.
All heavy BFS / flood-fill / territory ops run on plain integer arithmetic —
no sets, deques, or per-cell Python objects.

Typical 11×11 board → 121-bit integers.  Python big-int ops on these are
extremely fast (single C-level limb operations under the hood).
"""

from __future__ import annotations

class BitBoard:
  """Pre-computed masks and fast spatial primitives for a fixed grid size."""

  __slots__ = (
    "width", "height", "size", "board_mask",
    "_not_rightcol", "_not_leftcol",
    "_neighbor_masks",
  )

  def __init__(self, width: int, height: int) -> None:
    self.width = width
    self.height = height
    self.size = width * height
    self.board_mask = (1 << self.size) - 1

    # Column masks — prevent bit-shift wrap-around at row boundaries
    rightcol = 0
    leftcol = 0
    for y in range(height):
      rightcol |= 1 << (y * width + width - 1)
      leftcol |= 1 << (y * width)
    self._not_rightcol = self.board_mask & ~rightcol
    self._not_leftcol = self.board_mask & ~leftcol

    # Per-cell neighbour bitmask (4-connected)
    nb = [0] * self.size
    for idx in range(self.size):
      x, y = idx % width, idx // width
      mask = 0
      if x > 0:
        mask |= 1 << (idx - 1)
      if x < width - 1:
        mask |= 1 << (idx + 1)
      if y > 0:
        mask |= 1 << (idx - width)
      if y < height - 1:
        mask |= 1 << (idx + width)
      nb[idx] = mask
    self._neighbor_masks = nb

  # ── Coordinate helpers ────────────────────────────────────────────────────

  def idx(self, x: int, y: int) -> int:
    """(x, y) → flat index."""
    return y * self.width + x

  def coord(self, flat: int) -> tuple[int, int]:
    """Flat index → (x, y)."""
    return flat % self.width, flat // self.width

  def pt_bit(self, x: int, y: int) -> int:
    """Single-cell bitmask for (x, y)."""
    return 1 << (y * self.width + x)

  def set_to_bits(self, points: set[tuple[int, int]]) -> int:
    """Convert a set of (x, y) tuples to a bitmask."""
    w = self.width
    bits = 0
    for x, y in points:
      bits |= 1 << (y * w + x)
    return bits

  def in_bounds(self, x: int, y: int) -> bool:
    return 0 <= x < self.width and 0 <= y < self.height

  # ── Core spatial primitives ───────────────────────────────────────────────

  def flood_fill(self, start_idx: int, blocked_bits: int) -> int:
    """Return bitmask of all cells reachable from *start_idx* (inclusive)."""
    free = self.board_mask & ~blocked_bits
    start_bit = 1 << start_idx
    # If start is blocked, return just itself
    if not (start_bit & free):
      return start_bit

    reachable = start_bit
    frontier = start_bit
    w = self.width
    nrc = self._not_rightcol
    nlc = self._not_leftcol

    while frontier:
      expanded = (
        ((frontier & nrc) << 1)
        | ((frontier & nlc) >> 1)
        | (frontier << w)
        | (frontier >> w)
      ) & free & ~reachable
      if not expanded:
        break
      reachable |= expanded
      frontier = expanded

    return reachable

  def flood_count(self, start_idx: int, blocked_bits: int) -> int:
    """Count of cells reachable from *start_idx*."""
    return self.flood_fill(start_idx, blocked_bits).bit_count()

  def open_neighbor_count(self, cell_idx: int, blocked_bits: int) -> int:
    """Number of free neighbours of *cell_idx*."""
    return (self._neighbor_masks[cell_idx] & ~blocked_bits & self.board_mask).bit_count()

  def neighbors_of(self, cell_idx: int) -> int:
    """Bitmask of 4-connected neighbours (may include blocked cells)."""
    return self._neighbor_masks[cell_idx]

  # ── Territory (dual-BFS expansion) ────────────────────────────────────────

  def territory(
    self,
    my_idx: int,
    enemy_indices: list[int],
    blocked_bits: int,
  ) -> int:
    """Simultaneous BFS from *my_idx* and all enemies.

    Returns (my_cells − enemy_cells).  Cells equidistant from both sides are
    counted for neither (contested).
    """
    if not enemy_indices:
      return 0

    free = self.board_mask & ~blocked_bits
    w = self.width
    nrc = self._not_rightcol
    nlc = self._not_leftcol

    my_front = 1 << my_idx
    my_terr = my_front

    en_front = 0
    for ei in enemy_indices:
      en_front |= 1 << ei
    en_terr = en_front

    remaining = free & ~my_terr & ~en_terr

    while (my_front or en_front) and remaining:
      # Expand both sides simultaneously (same BFS depth → ties go to neither)
      my_exp = 0
      if my_front:
        my_exp = (
          ((my_front & nrc) << 1)
          | ((my_front & nlc) >> 1)
          | (my_front << w)
          | (my_front >> w)
        ) & remaining

      en_exp = 0
      if en_front:
        en_exp = (
          ((en_front & nrc) << 1)
          | ((en_front & nlc) >> 1)
          | (en_front << w)
          | (en_front >> w)
        ) & remaining

      # Contested cells (reached by both at the same depth) → neither claims
      contested = my_exp & en_exp
      my_exp &= ~contested
      en_exp &= ~contested

      my_terr |= my_exp
      en_terr |= en_exp
      remaining &= ~(my_exp | en_exp | contested)

      my_front = my_exp
      en_front = en_exp

    return my_terr.bit_count() - en_terr.bit_count()

  # ── Partition sizes (for articulation-point detection) ────────────────────

  def partition_sizes(self, cut_idx: int, blocked_bits: int) -> list[int]:
    """Remove *cut_idx* from the free space and return sizes of each
    resulting connected component among its neighbours.

    Returns an empty list when the point is not a cut vertex (single component
    or ≤1 free neighbour).
    """
    test_blocked = blocked_bits | (1 << cut_idx)
    free_nb = self._neighbor_masks[cut_idx] & ~test_blocked & self.board_mask
    if free_nb.bit_count() <= 1:
      return []

    seen_all = 0
    sizes: list[int] = []

    temp = free_nb
    while temp:
      bit = temp & (-temp)   # lowest set bit
      temp ^= bit
      if bit & seen_all:
        continue
      component = self.flood_fill(bit.bit_length() - 1, test_blocked)
      seen_all |= component
      sizes.append(component.bit_count())

    return sizes if len(sizes) > 1 else []

  # ── BFS distance map (indexed by cell idx) ────────────────────────────────

  def distance_map(self, start_idx: int, blocked_bits: int) -> dict[int, int]:
    """BFS distance from *start_idx* to every reachable cell.

    Returns ``{cell_idx: distance}`` — same semantics as the original
    ``_distance_map`` but using bitboard expansion internally.
    """
    free = self.board_mask & ~blocked_bits
    start_bit = 1 << start_idx
    distances: dict[int, int] = {start_idx: 0}
    frontier = start_bit
    seen = frontier
    dist = 0
    w = self.width
    nrc = self._not_rightcol
    nlc = self._not_leftcol

    while frontier:
      dist += 1
      expanded = (
        ((frontier & nrc) << 1)
        | ((frontier & nlc) >> 1)
        | (frontier << w)
        | (frontier >> w)
      ) & free & ~seen

      if not expanded:
        break

      seen |= expanded
      # Extract individual bits
      temp = expanded
      while temp:
        bit = temp & (-temp)
        idx = bit.bit_length() - 1
        distances[idx] = dist
        temp ^= bit

      frontier = expanded

    return distances

  # ── Path distance (BFS to single target) ──────────────────────────────────

  def path_distance(
    self,
    start_idx: int,
    goal_idx: int,
    blocked_bits: int,
  ) -> int | None:
    """Shortest path length from *start_idx* to *goal_idx*, or ``None``."""
    # Unblock the goal cell so BFS can reach it
    free = (self.board_mask & ~blocked_bits) | (1 << goal_idx)
    start_bit = 1 << start_idx
    goal_bit = 1 << goal_idx

    if start_idx == goal_idx:
      return 0

    frontier = start_bit
    seen = frontier
    dist = 0
    w = self.width
    nrc = self._not_rightcol
    nlc = self._not_leftcol

    while frontier:
      dist += 1
      expanded = (
        ((frontier & nrc) << 1)
        | ((frontier & nlc) >> 1)
        | (frontier << w)
        | (frontier >> w)
      ) & free & ~seen

      if not expanded:
        break

      if expanded & goal_bit:
        return dist

      seen |= expanded
      frontier = expanded

    return None

  # ── Nearest-food BFS ──────────────────────────────────────────────────────

  def nearest_food(
    self,
    start_idx: int,
    food_bits: int,
    blocked_bits: int,
  ) -> tuple[int | None, int | None]:
    """BFS from *start_idx* to nearest food cell.

    Food cells are passable even if in *blocked_bits* (matching original
    ``_nearest_food_info`` semantics).

    Returns ``(distance, cell_idx)`` or ``(None, None)``.
    """
    if not food_bits:
      return None, None

    # Food tiles are always steppable
    free = (self.board_mask & ~blocked_bits) | food_bits
    start_bit = 1 << start_idx

    # Check start
    if start_bit & food_bits:
      return 0, start_idx

    frontier = start_bit
    seen = frontier
    dist = 0
    w = self.width
    nrc = self._not_rightcol
    nlc = self._not_leftcol

    while frontier:
      dist += 1
      expanded = (
        ((frontier & nrc) << 1)
        | ((frontier & nlc) >> 1)
        | (frontier << w)
        | (frontier >> w)
      ) & free & ~seen

      if not expanded:
        break

      hit = expanded & food_bits
      if hit:
        # Return the first (lowest-index) food cell found
        first_bit = hit & (-hit)
        return dist, first_bit.bit_length() - 1

      seen |= expanded
      frontier = expanded

    return None, None