"""
Phutball (Philosopher's Football) game logic.

Board is always stored in canonical form: the current player's goal is always
at row (rows-1). This means MCTS and the NN never need to flip anything —
they always see the same orientation. When the active player changes, next_state
flips the board before returning.

Rules:
  - On your turn you may place a man on any empty square (including goal rows),
    OR start a jump sequence.
  - A jump sends the ball over one or more consecutive men in a straight line
    (8 directions). Jumped men are removed.
  - After a jump you may continue jumping (if men are adjacent) or end your turn.
  - Landing ON a goal row does not end the game while a jump sequence can
    continue. The game ends only if the jump sequence ends there.
  - Jumping PAST a goal row (off the board) also ends the game immediately.
  - end_turn is only needed mid-jump when further jumps exist but the player
    chooses to stop. next_state automatically ends the turn when no further
    jumps are available.

Action encoding (in canonical space):
  0  ..  rows*cols-1 : place a man at (action // cols, action % cols)
  rows*cols  ..  +7  : jump in direction DIRECTIONS[i]
  rows*cols + 8      : end turn (only valid during mid-jump)
"""

from __future__ import annotations

import numpy as np
from typing import List, Optional, Tuple

# ── Cell values ───────────────────────────────────────────────────────────────

EMPTY = 0
MAN   = 1
BALL  = 2

# ── Directions: N, NE, E, SE, S, SW, W, NW ───────────────────────────────────

DIRECTIONS: List[Tuple[int, int]] = [
    (-1,  0),  # N
    (-1,  1),  # NE
    ( 0,  1),  # E
    ( 1,  1),  # SE
    ( 1,  0),  # S
    ( 1, -1),  # SW
    ( 0, -1),  # W
    (-1, -1),  # NW
]
NUM_DIRS = 8

# Permutation of direction indices under a left-right (horizontal) flip.
# N→N, NE→NW, E→W, SE→SW, S→S, SW→SE, W→E, NW→NE
LR_FLIP_DIRS = [0, 7, 6, 5, 4, 3, 2, 1]


# ── State ─────────────────────────────────────────────────────────────────────

class PhutballState:
    """
    Immutable-by-convention game state.

    board         -- (rows, cols) int8 array, canonical orientation:
                     current player's goal is always at row (rows-1).
    is_mid_jump   -- True while the current player is in a jump sequence.
    current_player -- 1 or -1, used for terminal_value and display.
    """

    __slots__ = ("board", "is_mid_jump", "current_player")

    def __init__(self, board: np.ndarray, is_mid_jump: bool, current_player: int):
        self.board = board
        self.is_mid_jump = is_mid_jump
        self.current_player = current_player

    def copy(self) -> PhutballState:
        return PhutballState(self.board.copy(), self.is_mid_jump, self.current_player)


# ── Game ──────────────────────────────────────────────────────────────────────

class PhutballGame:
    """
    Complete Phutball game interface consumed by MCTS and the training pipeline.

    Parameters
    ----------
    rows, cols      Board dimensions for this instance.
    max_rows, max_cols  Tensor dimensions (>= rows/cols). Allows a single NN
                    to train across board sizes (KataGo curriculum idea):
                    smaller boards are padded to max size.
    num_channels    Number of input channels in canonical_tensor(). Start with 1
                    (raw board); add engineered features later via features.py
                    without changing this interface.
    """

    def __init__(
        self,
        rows: int = 9,
        cols: int = 9,
        max_rows: int = 9,
        max_cols: int = 9,
        num_channels: int = 1,
    ):
        assert max_rows >= rows and max_cols >= cols, \
            "max dimensions must be >= board dimensions"

        self.rows = rows
        self.cols = cols
        self.max_rows = max_rows
        self.max_cols = max_cols
        self.num_channels = num_channels

        # Action-space layout (depends on board size)
        self._n_placements  = rows * cols
        self._jump_offset   = self._n_placements          # first jump action index
        self._end_turn      = self._n_placements + NUM_DIRS  # last action index
        self._n_actions     = self._n_placements + NUM_DIRS + 1

    # ── Properties ────────────────────────────────────────────────────────────

    @property
    def action_size(self) -> int:
        """Total number of distinct actions for this board size."""
        return self._n_actions

    @property
    def n_placements(self) -> int:
        """Number of placement actions at the front of the action vector."""
        return self._n_placements

    @property
    def jump_offset(self) -> int:
        """Index of the first jump action in the action vector."""
        return self._jump_offset

    @property
    def end_turn_action(self) -> int:
        """Action index used to voluntarily end a jump sequence."""
        return self._end_turn

    @property
    def tensor_shape(self) -> Tuple[int, int, int]:
        """Shape of canonical_tensor(): (C, max_rows, max_cols)."""
        return (self.num_channels, self.max_rows, self.max_cols)

    # ── Core interface ────────────────────────────────────────────────────────

    def initial_state(self) -> PhutballState:
        """Ball at center, player 1 to move."""
        board = np.zeros((self.rows, self.cols), dtype=np.int8)
        board[self.rows // 2, self.cols // 2] = BALL
        return PhutballState(board=board, is_mid_jump=False, current_player=1)

    def current_player(self, state: PhutballState) -> int:
        """Return the player to act in this state: 1 or -1."""
        return state.current_player

    def legal_moves(self, state: PhutballState) -> np.ndarray:
        """
        Binary mask of shape (action_size,) in canonical action space.

        During mid-jump: only jump actions and end_turn are valid.
        Otherwise:      placements on any empty square + any available jump.
        """
        valid = np.zeros(self._n_actions, dtype=np.int8)
        if self.is_terminal(state):
            return valid

        br, bc = self._find_ball(state.board)

        if state.is_mid_jump:
            for i, (dr, dc) in enumerate(DIRECTIONS):
                if self._can_jump(state.board, br, bc, dr, dc) is not None:
                    valid[self._jump_offset + i] = 1
            valid[self._end_turn] = 1
        else:
            for r in range(self.rows):
                for c in range(self.cols):
                    if state.board[r, c] == EMPTY:
                        valid[r * self.cols + c] = 1
            for i, (dr, dc) in enumerate(DIRECTIONS):
                if self._can_jump(state.board, br, bc, dr, dc) is not None:
                    valid[self._jump_offset + i] = 1

        return valid

    def next_state(self, state: PhutballState, action: int) -> PhutballState:
        """
        Apply a canonical action and return the new canonical state.

        Placement / end_turn: the active player changes, so the board is
        flipped vertically before returning (new canonical orientation).

        Jump (on-board): same player continues; board stays in same orientation.

        Jump (off-board): game ends immediately. If the ball exits past
        row (rows-1) the current player wins; if it exits past row 0 the
        opponent wins (the board is then flipped so the winner is current_player
        in the returned state, keeping terminal_value = +1 consistent).
        """
        new = state.copy()
        board = new.board

        if action < 0 or action >= self._n_actions:
            raise ValueError(f"Action {action} outside valid range [0, {self._n_actions})")

        if action < self._n_placements:
            # ── Placement ────────────────────────────────────────────────────
            if state.is_mid_jump:
                raise ValueError("Placement is not legal during a jump sequence")
            row, col = divmod(action, self.cols)
            if board[row, col] != EMPTY:
                raise ValueError(f"Placement target ({row}, {col}) is not empty")
            board[row, col] = MAN
            new.is_mid_jump = False
            new.current_player = -state.current_player
            new.board = np.flip(board, axis=0).copy()

        elif action < self._end_turn:
            # ── Jump ─────────────────────────────────────────────────────────
            dir_idx = action - self._jump_offset
            dr, dc = DIRECTIONS[dir_idx]
            br, bc = self._find_ball(board)
            jump = self._can_jump(board, br, bc, dr, dc)
            if jump is None:
                raise ValueError(f"Jump action {action} is not legal")
            land_r, land_c, men_jumped, off_board = jump

            # Remove ball and all jumped men
            board[br, bc] = EMPTY
            for mr, mc in men_jumped:
                board[mr, mc] = EMPTY

            if off_board:
                clamp_c = max(0, min(self.cols - 1, land_c))

                if land_r >= self.rows:
                    # Past current player's goal → current player wins
                    board[self.rows - 1, clamp_c] = BALL
                    new.is_mid_jump = False
                    # current_player unchanged: they won, terminal_value → +1

                else:
                    # land_r < 0: past opponent's goal → opponent wins
                    # Place ball at row 0, then flip so the winner sees it at rows-1
                    board[0, clamp_c] = BALL
                    new.is_mid_jump = False
                    new.current_player = -state.current_player
                    new.board = np.flip(board, axis=0).copy()

            else:
                board[land_r, land_c] = BALL
                has_jumps = any(
                    self._can_jump(board, land_r, land_c, dr, dc) is not None
                    for dr, dc in DIRECTIONS
                )
                if has_jumps:
                    # Further jumps available — player stays in mid-jump.
                    new.is_mid_jump = True
                else:
                    # No further jumps — forced end of turn.
                    new.is_mid_jump = False
                    new.current_player = -state.current_player
                    new.board = np.flip(board, axis=0).copy()

        else:
            # ── End turn ─────────────────────────────────────────────────────
            if not state.is_mid_jump:
                raise ValueError("end_turn is only legal during a jump sequence")
            new.is_mid_jump = False
            new.current_player = -state.current_player
            new.board = np.flip(board, axis=0).copy()

        return new

    def is_terminal(self, state: PhutballState) -> bool:
        """
        The game is over when the turn has ended (not mid-jump) and the ball
        is on a goal row. In canonical form both goal rows are row 0 and
        row (rows-1), and the current player's goal is always rows-1.
        """
        if state.is_mid_jump:
            return False
        br, _ = self._find_ball(state.board)
        return br == 0 or br == self.rows - 1

    def terminal_value(self, state: PhutballState) -> float:
        """
        +1.0 if current_player has won, -1.0 if they have lost.

        In canonical form the current player's goal is always row (rows-1),
        so ball at rows-1 → win, ball at 0 → loss.
        Only meaningful when is_terminal() is True.
        """
        br, _ = self._find_ball(state.board)
        return 1.0 if br == self.rows - 1 else -1.0

    def canonical_tensor(self, state: PhutballState) -> np.ndarray:
        """
        Returns a float32 tensor of shape (num_channels, max_rows, max_cols).

        Channel 0 holds the raw board values (EMPTY=0, MAN=1, BALL=2).
        Additional channels (e.g. engineered features) can be appended later
        without changing the game/search interface.

        The board is already canonical (no extra flip needed here).
        Smaller boards are zero-padded to (max_rows, max_cols).
        """
        out = np.zeros((self.num_channels, self.max_rows, self.max_cols), dtype=np.float32)
        out[0, :self.rows, :self.cols] = state.board
        return out

    def hash(self, state: PhutballState) -> bytes:
        """Compact, collision-resistant key for MCTS node lookup."""
        return (
            state.board.tobytes()
            + bytes([int(state.is_mid_jump), state.current_player & 0xFF])
        )

    def symmetries(
        self,
        state: PhutballState,
        policy: np.ndarray,
        aux_policy: Optional[np.ndarray] = None,
    ) -> List[Tuple[PhutballState, np.ndarray, Optional[np.ndarray]]]:
        """
        Data augmentation via left-right (horizontal) symmetry.

        Returns [(original, ...), (flipped, ...)] so the caller can add
        both to the training buffer. aux_policy is the KataGo auxiliary
        target (opponent's next-move policy); it is flipped identically.
        """
        original = (state, policy, aux_policy)

        flipped_board  = np.flip(state.board, axis=1).copy()
        flipped_state  = PhutballState(flipped_board, state.is_mid_jump, state.current_player)
        flipped_policy = self._flip_policy_lr(policy)
        flipped_aux    = self._flip_policy_lr(aux_policy) if aux_policy is not None else None

        return [original, (flipped_state, flipped_policy, flipped_aux)]

    def display(self, state: PhutballState) -> None:
        """
        Pretty-print the board in canonical orientation (current player's goal
        at the bottom). Useful for debugging in the terminal.
        """
        syms = {EMPTY: ".", MAN: "o", BALL: "O"}
        mid  = " (mid-jump)" if state.is_mid_jump else ""
        print(f"\nPlayer {state.current_player:+d} to move{mid}")
        print("   " + "  ".join(str(c) for c in range(self.cols)))
        for r in range(self.rows):
            row_str = "  ".join(syms[state.board[r, c]] for c in range(self.cols))
            tag = " ← opponent goal" if r == 0 else (" ← YOUR goal" if r == self.rows - 1 else "")
            print(f"{r:2d} {row_str}{tag}")
        print()

    # ── Private helpers ───────────────────────────────────────────────────────

    def _find_ball(self, board: np.ndarray) -> Tuple[int, int]:
        pos = np.argwhere(board == BALL)
        if len(pos) == 0:
            raise ValueError("No ball found on board")
        return int(pos[0, 0]), int(pos[0, 1])

    def _can_jump(
        self,
        board: np.ndarray,
        br: int,
        bc: int,
        dr: int,
        dc: int,
    ) -> Optional[Tuple[int, int, List[Tuple[int, int]], bool]]:
        """
        Check whether the ball at (br, bc) can jump in direction (dr, dc).

        Returns (land_r, land_c, men_jumped, off_board) or None if invalid.

        off_board=True means the ball exits the board past a goal row:
          land_r >= rows  →  current player's goal side  (win)
          land_r <  0     →  opponent's goal side         (own goal / loss)

        Exiting off the left/right side only (column out of range, row still
        in range) is never valid.

        Corner jumps are handled naturally: a diagonal that simultaneously
        exits a goal row and a side column is valid because the goal-row exit
        takes precedence (r check runs before c check).
        """
        r, c = br + dr, bc + dc

        # The first cell in the jump direction must be a man on the board
        if not (0 <= r < self.rows and 0 <= c < self.cols):
            return None
        if board[r, c] != MAN:
            return None

        # Collect all consecutive men
        men_jumped: List[Tuple[int, int]] = []
        while 0 <= r < self.rows and 0 <= c < self.cols and board[r, c] == MAN:
            men_jumped.append((r, c))
            r += dr
            c += dc

        # Determine landing
        if r < 0 or r >= self.rows:
            # Off-board past a goal row — winning (or own-goal) jump
            return (r, c, men_jumped, True)

        if c < 0 or c >= self.cols:
            # Off the side but not past a goal row — invalid
            return None

        if board[r, c] == EMPTY:
            return (r, c, men_jumped, False)

        return None  # Blocked by man or ball

    def _flip_policy_lr(self, policy: np.ndarray) -> np.ndarray:
        """Mirror a policy vector under a left-right board flip."""
        flipped = np.empty_like(policy)

        # Placement actions: mirror each column
        for r in range(self.rows):
            for c in range(self.cols):
                flipped[r * self.cols + (self.cols - 1 - c)] = policy[r * self.cols + c]

        # Jump directions
        base = self._jump_offset
        for i, j in enumerate(LR_FLIP_DIRS):
            flipped[base + i] = policy[base + j]

        # End-turn action is invariant
        flipped[self._end_turn] = policy[self._end_turn]

        return flipped
