"""
High-level Pythonic API for computing Surprise.
"""

from __future__ import annotations

from typing import Dict, List, Sequence, Tuple, Union

from pysurprise._core import (
    compute_surprise,
    surprise_from_partition,
    surprise_from_partition_dict,
)


def surprise(
    edges: Sequence[Tuple[Union[int, str], Union[int, str]]],
    partition: Union[Sequence[int], Dict[Union[int, str], int]],
) -> float:
    """Compute the Surprise of a network partition.

    This is the main entry point. It accepts flexible input formats and
    dispatches to the optimised C++ implementation.

    Parameters
    ----------
    edges : sequence of (node, node) tuples
        Each element is an undirected edge.  Nodes can be integers (0-based
        indices) or strings (arbitrary labels).
    partition : list of int  **or**  dict mapping node -> community_id
        - **list**: ``partition[i]`` is the community id of node *i*.
          Nodes in *edges* must be ``int`` in ``[0, len(partition))``.
        - **dict**: maps every node label that appears in *edges* to its
          community id (``int``).

    Returns
    -------
    float
        The Surprise value (in −log₁₀ scale).  Higher values indicate
        a more surprising (i.e. better) partition.

    Examples
    --------
    >>> import pysurprise
    >>> edges = [(0, 1), (0, 2), (1, 2), (2, 3), (3, 4), (3, 5), (4, 5)]
    >>> partition = [0, 0, 0, 1, 1, 1]
    >>> pysurprise.surprise(edges, partition)
    3.212...

    Using string labels and a dict partition:

    >>> edges = [("a", "b"), ("a", "c"), ("b", "c"),
    ...          ("c", "d"), ("d", "e"), ("d", "f"), ("e", "f")]
    >>> partition = {"a": 0, "b": 0, "c": 0, "d": 1, "e": 1, "f": 1}
    >>> pysurprise.surprise(edges, partition)
    3.212...
    """
    # ── dict partition ──────────────────────────────────────────────
    if isinstance(partition, dict):
        str_edges: List[Tuple[str, str]] = [
            (str(u), str(v)) for u, v in edges
        ]
        str_partition: Dict[str, int] = {
            str(k): int(v) for k, v in partition.items()
        }
        return surprise_from_partition_dict(str_edges, str_partition)

    # ── list / sequence partition ───────────────────────────────────
    int_edges: List[Tuple[int, int]] = [(int(u), int(v)) for u, v in edges]
    int_partition: List[int] = [int(c) for c in partition]
    return surprise_from_partition(int_edges, int_partition)
