"""
norvig.adversarial – AIMA Chapter 5
=====================================
Minimax and alpha-beta pruning are implemented directly (no appropriate lib).
MCTS uses the ``mcts`` package (pip install mcts) as the search engine.

    from ml_ai_library.norvig.adversarial import (
        Game,               # abstract game interface
        minimax_decision,   # AIMA Fig 5.3
        alphabeta_search,   # AIMA Fig 5.7
        mcts_search,        # uses mcts package
        TicTacToe,          # example game
    )
"""

from __future__ import annotations
import math
import random
from typing import Any, List, Optional

from mcts import mcts as _mcts_lib                              # mcts package


# ---------------------------------------------------------------------------
# Abstract game interface (AIMA §5.1)
# ---------------------------------------------------------------------------

class Game:
    """Abstract two-player zero-sum game."""

    def initial_state(self) -> Any:
        raise NotImplementedError

    def player(self, state: Any) -> int:
        """Return the player to move: +1 or -1."""
        raise NotImplementedError

    def actions(self, state: Any) -> List[Any]:
        raise NotImplementedError

    def result(self, state: Any, action: Any) -> Any:
        raise NotImplementedError

    def terminal_test(self, state: Any) -> bool:
        raise NotImplementedError

    def utility(self, state: Any, player: int) -> float:
        """Terminal utility from `player`'s perspective: +1 win, -1 loss, 0 draw."""
        raise NotImplementedError


# ---------------------------------------------------------------------------
# Minimax (AIMA Fig 5.3)
# ---------------------------------------------------------------------------

def minimax_decision(game: Game, state: Any) -> Any:
    """Return the minimax-optimal action for the current player."""
    player = game.player(state)

    def max_val(s):
        if game.terminal_test(s):
            return game.utility(s, player)
        return max(min_val(game.result(s, a)) for a in game.actions(s))

    def min_val(s):
        if game.terminal_test(s):
            return game.utility(s, player)
        return min(max_val(game.result(s, a)) for a in game.actions(s))

    return max(game.actions(state),
               key=lambda a: min_val(game.result(state, a)))


# ---------------------------------------------------------------------------
# Alpha-Beta Pruning (AIMA Fig 5.7)
# ---------------------------------------------------------------------------

def alphabeta_search(game: Game, state: Any,
                     depth_limit: int = 100, eval_fn=None) -> Any:
    """Alpha-beta pruned minimax.  eval_fn(state, player) used at depth limit."""
    player = game.player(state)
    INF    = math.inf

    def max_val(s, alpha, beta, depth):
        if game.terminal_test(s):
            return game.utility(s, player)
        if depth == 0 and eval_fn:
            return eval_fn(s, player)
        v = -INF
        for a in game.actions(s):
            v = max(v, min_val(game.result(s, a), alpha, beta, depth - 1))
            if v >= beta:
                return v
            alpha = max(alpha, v)
        return v

    def min_val(s, alpha, beta, depth):
        if game.terminal_test(s):
            return game.utility(s, player)
        if depth == 0 and eval_fn:
            return eval_fn(s, player)
        v = INF
        for a in game.actions(s):
            v = min(v, max_val(game.result(s, a), alpha, beta, depth - 1))
            if v <= alpha:
                return v
            beta = min(beta, v)
        return v

    best_score, best_action = -INF, None
    for a in game.actions(state):
        score = min_val(game.result(state, a), -INF, INF, depth_limit - 1)
        if score > best_score:
            best_score, best_action = score, a
    return best_action


# ---------------------------------------------------------------------------
# MCTS via the `mcts` package (AIMA §5.4)
# ---------------------------------------------------------------------------

def mcts_search(game: Game, state: Any,
                n_simulations: int = 1000,
                time_limit: Optional[float] = None) -> Any:
    """
    Monte Carlo Tree Search using the ``mcts`` package as the engine.

    Parameters
    ----------
    game          : Game instance
    state         : current game state
    n_simulations : iteration budget (used when time_limit is None)
    time_limit    : time budget in seconds (overrides n_simulations)

    Returns
    -------
    best_action

    Example
    -------
    >>> from ml_ai_library.norvig.adversarial import mcts_search, TicTacToe
    >>> game = TicTacToe()
    >>> move = mcts_search(game, game.initial_state(), n_simulations=500)
    """
    # Wrap our Game in the mcts package's State interface
    class _MCTSState:
        def __init__(self, s):
            self._s = s

        def getCurrentPlayer(self):
            return 1       # mcts package uses 1/-1; we normalise internally

        def getPossibleActions(self):
            return game.actions(self._s)

        def takeAction(self, action):
            return _MCTSState(game.result(self._s, action))

        def isTerminal(self):
            return game.terminal_test(self._s)

        def getReward(self):
            # mcts package expects reward in [0, 1] from current player perspective
            u = game.utility(self._s, game.player(state))  # use root player
            return (u + 1) / 2   # map [-1,1] -> [0,1]

    kwargs = {"timeLimit": int(time_limit * 1000)} if time_limit else \
             {"iterationLimit": n_simulations}
    searcher = _mcts_lib(**kwargs)
    return searcher.search(initialState=_MCTSState(state))


# ---------------------------------------------------------------------------
# Example game: Tic-Tac-Toe
# ---------------------------------------------------------------------------

class TicTacToe(Game):
    """
    3x3 Tic-Tac-Toe.  State: (board_tuple_9_ints, current_player).
    Board values: 0=empty, 1=X, -1=O.
    """

    _WINS = [(0,1,2),(3,4,5),(6,7,8),(0,3,6),(1,4,7),(2,5,8),(0,4,8),(2,4,6)]

    def initial_state(self):
        return (tuple([0] * 9), 1)

    def player(self, state):
        return state[1]

    def actions(self, state):
        board, _ = state
        return [i for i, v in enumerate(board) if v == 0]

    def result(self, state, action):
        board, p = state
        b = list(board)
        b[action] = p
        return (tuple(b), -p)

    def terminal_test(self, state):
        board, _ = state
        if any(abs(board[a] + board[b] + board[c]) == 3
               for a, b, c in self._WINS):
            return True
        return all(v != 0 for v in board)

    def utility(self, state, player):
        board, _ = state
        for a, b, c in self._WINS:
            s = board[a] + board[b] + board[c]
            if abs(s) == 3:
                return 1 if s // 3 == player else -1
        return 0
