"""
# Communication constraints and restricted games.

This module currently provides the Myerson restriction induced by a communication
graph. It constructs a new TU game where a coalition's worth is the sum of the
worths of its connected components in the communication graph.
"""

from __future__ import annotations

from typing import Iterable

from ..base.coalition import all_coalitions
from ..base.game import Game
from ..base.exceptions import InvalidGameError, InvalidParameterError, NotSupportedError
from ..base.types import GameProtocol, require_tabular_game


def _validate_edges(n: int, edges: Iterable[tuple[int, int]]) -> list[int]:
    adj = [0] * n
    for u, v in edges:
        uu = int(u)
        vv = int(v)
        if uu < 0 or uu >= n or vv < 0 or vv >= n:
            raise InvalidParameterError("edge endpoints out of range")
        if uu == vv:
            continue
        adj[uu] |= 1 << vv
        adj[vv] |= 1 << uu
    return adj


def _connected_components(mask: int, *, adj: list[int]) -> list[int]:
    comps: list[int] = []
    rem = int(mask)
    while rem:
        seed = rem & -rem
        queue = seed
        comp = 0
        rem ^= seed
        while queue:
            bit = queue & -queue
            queue ^= bit
            comp |= bit
            i = bit.bit_length() - 1
            nbrs = adj[i] & rem
            queue |= nbrs
            rem ^= nbrs
        comps.append(int(comp))
    comps.sort()
    return comps


def myerson_restriction(
    game: GameProtocol,
    *,
    edges: Iterable[tuple[int, int]],
    require_complete: bool = True,
    max_players: int = 16,
) -> Game:
    """
    Compute the **Myerson restriction** of a TU game induced by a communication graph.

    In cooperative games with communication constraints, players can only cooperate
    effectively within **connected components** of the communication graph. Given
    an undirected graph $G=(N,E)$ and a coalition $S \\subseteq N$, let
    $\\{C_1, \\dots, C_k\\}$ be the connected components of the subgraph induced by $S$.
    The Myerson restricted game $v_G$ is defined by:

    $$
    v_G(S) = \\sum_{C \\in \\mathrm{components}_G(S)} v(C),
    $$

    with the TU convention $v_G(\\varnothing)=0$.

    Intuition:
        A coalition $S$ can only realize the value of the groups that can actually
        communicate (the connected components). The total value is the sum of the
        values generated by each communicating group.

    Parameters
    ----------
    game : Game
        Original TU game $v$ on the player set $N$.
    edges : Iterable[tuple[int, int]]
        Undirected edges of the communication graph $G$ as pairs of player indices.
        Self-loops are ignored. Duplicate edges are allowed.
    require_complete : bool, default=True
        If True, require `game` to specify a complete characteristic function,
        i.e. values for all $2^n$ coalitions. This restriction is useful to avoid
        silently treating missing coalitions as 0 when `Game.value()` is used.
    max_players : int, default=16
        Safety limit for the number of players. This routine enumerates all
        coalitions and is therefore exponential in `n_players`.

    Returns
    -------
    Game
        A new **complete** TU game with characteristic function $v_G$.

    Raises
    ------
    NotSupportedError
        If `n_players > max_players`.
    InvalidGameError
        If `require_complete=True` and the game does not specify all $2^n$ coalition values.
    InvalidParameterError
        If any edge endpoint is out of range.

    Notes
    -----
    - **Complexity:** This function enumerates all coalitions, so its runtime and
      memory scale as $\\Theta(2^n)$. Use `max_players` as a guardrail.
    - The resulting game is complete: it explicitly stores values for every coalition.
    - This function implements the *restricted game* used to define the **Myerson value**
      (the Shapley value of $v_G$) in communication situations.

    Examples
    --------
    Restrict a 3-player game with a path graph 0--1--2:

    >>> gG = myerson_restriction(g, edges=[(0, 1), (1, 2)])
    >>> gG.value(0)  # empty coalition
    0.0

    If players 0 and 2 cannot communicate directly (only through 1),
    then coalition {0,2} has two singleton components:

    >>> from tucoopy.base.coalition import mask_from_players
    >>> S = mask_from_players([0, 2])
    >>> gG.value(S) == g.value(mask_from_players([0])) + g.value(mask_from_players([2]))
    True
    """
    n = game.n_players
    if n > int(max_players):
        raise NotSupportedError(f"myerson_restriction is exponential; requires n<={max_players} (got n={n})")
    if require_complete:
        tabular = require_tabular_game(game, context="myerson_restriction")
        if len(tabular.v) < (1 << n):
            raise InvalidGameError("myerson_restriction requires a complete characteristic function (2^n coalition values)")

    adj = _validate_edges(n, edges)

    v = [0.0] * (1 << n)
    for S in all_coalitions(n):
        v[S] = float(game.value(S))

    vG: dict[int, float] = {0: 0.0}
    for S in range(1, 1 << n):
        total = 0.0
        for C in _connected_components(S, adj=adj):
            total += float(v[C])
        vG[int(S)] = float(total)

    labels = getattr(game, "player_labels", None)
    return Game(n_players=n, v=vG, player_labels=labels)


__all__ = ["myerson_restriction"]
