"""
norvig.game_theory – AIMA Chapter 5 (Game Theory Appendix)
============================================================
Uses ``nashpy`` for Nash equilibrium computation.

    from ml_ai_library.norvig.game_theory import (
        NormalFormGame,     # thin wrapper over nashpy.Game
        nash_equilibria,    # support enumeration via nashpy
        dominant_strategy,  # pure dominant strategy check
        maximin,            # minimax / maximin for zero-sum games
    )

AIMA chapter map:
  Nash equilibria     → §5.4 / §17.5
  Dominant strategy   → §5.4
  Maximin             → §5.4
"""

from __future__ import annotations
from typing import List, Optional, Tuple

import numpy as np
import nashpy as nash                           # Nash equilibrium engine


# ---------------------------------------------------------------------------
# Normal-Form Game (AIMA §17.5)
# ---------------------------------------------------------------------------

class NormalFormGame:
    """
    Two-player normal-form game backed by nashpy.Game.

    Parameters
    ----------
    A : (m, n) numpy array – row player payoffs
    B : (m, n) numpy array – column player payoffs (B = -A for zero-sum)

    Example – Prisoner's Dilemma
    ----------------------------
    >>> A = np.array([[-1, -3], [0, -2]])
    >>> B = np.array([[-1,  0], [-3, -2]])
    >>> g = NormalFormGame(A, B)
    >>> eqs = nash_equilibria(g)
    >>> eqs[0]          # (row_mixed, col_mixed)
    (array([0., 1.]), array([0., 1.]))
    """

    def __init__(self, A: np.ndarray, B: np.ndarray):
        self.A    = np.asarray(A, dtype=float)
        self.B    = np.asarray(B, dtype=float)
        self._game = nash.Game(self.A, self.B)

    @classmethod
    def zero_sum(cls, A: np.ndarray) -> "NormalFormGame":
        """Construct a zero-sum game from the row player's payoff matrix."""
        return cls(A, -np.asarray(A, dtype=float))

    def expected_payoff(self, p: np.ndarray, q: np.ndarray) -> Tuple[float, float]:
        """Expected payoff (u_row, u_col) for mixed strategies p, q."""
        p, q = np.asarray(p), np.asarray(q)
        return float(p @ self.A @ q), float(p @ self.B @ q)


def nash_equilibria(game: NormalFormGame) -> List[Tuple[np.ndarray, np.ndarray]]:
    """
    Compute all Nash equilibria using nashpy's support enumeration (AIMA §17.5).

    Returns list of (row_strategy, col_strategy) mixed-strategy pairs.
    """
    return list(game._game.support_enumeration())


def vertex_equilibria(game: NormalFormGame) -> List[Tuple[np.ndarray, np.ndarray]]:
    """Alternative enumeration via vertex enumeration (nashpy)."""
    return list(game._game.vertex_enumeration())


def dominant_strategy(A: np.ndarray) -> Optional[int]:
    """
    Return the strictly dominant row-action index for the row player, or None.
    Uses numpy broadcasting to compare each row against all others.
    """
    A = np.asarray(A)
    m = A.shape[0]
    for i in range(m):
        if all(np.all(A[i] > A[j]) for j in range(m) if j != i):
            return i
    return None


def maximin(A: np.ndarray) -> Tuple[int, float]:
    """
    Row player's maximin strategy for a zero-sum game (AIMA §5.4).
    Uses numpy: max over rows of (min over columns).
    Returns (best_action_index, maximin_value).
    """
    A   = np.asarray(A)
    row_mins = A.min(axis=1)           # worst case per row action
    idx = int(np.argmax(row_mins))
    return idx, float(row_mins[idx])
