"""
LFR benchmark graph generator for pySurprise.

Wraps :func:`networkx.generators.community.LFR_benchmark_graph` and converts
the output into pySurprise's native format:

- **edges**: ``list[tuple[str, str]]`` — the same ``.pairs`` format accepted
  by :func:`pysurprise.surprise` and every function in
  :mod:`pysurprise.algorithms`.
- **ground_truth**: ``dict[str, int]`` — mapping ``{node_label: community_id}``
  that can be compared against algorithm outputs.

Requirements
------------
``networkx >= 2.6`` (``pip install networkx``)

References
----------
    Lancichinetti A, Fortunato S, Radicchi F (2008)
    Benchmark graphs for testing community detection algorithms.
    Physical Review E 78(4): 046110
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, List, Optional, Sequence, Tuple

try:
    import networkx as nx
    from networkx.generators.community import LFR_benchmark_graph
except ImportError:
    raise ImportError(
        "The 'networkx' package is required for LFR benchmarks. "
        "Install it with:  pip install networkx"
    )

Edge = Tuple[str, str]
Partition = Dict[str, int]


# ---------------------------------------------------------------------------
# Result container
# ---------------------------------------------------------------------------

@dataclass
class LFRResult:
    """Container returned by :func:`lfr` and :func:`lfr_graph`.

    Attributes
    ----------
    edges : list of (str, str)
        Edge list in pySurprise ``.pairs`` format.
    ground_truth : dict[str, int]
        Ground-truth partition ``{node_label: community_id}``.
    graph : networkx.Graph
        The full NetworkX graph (nodes carry a ``"community"`` attribute with
        a frozenset of community ids).
    n_communities : int
        Number of ground-truth communities.
    """

    edges: List[Edge]
    ground_truth: Partition
    graph: "nx.Graph"
    n_communities: int

    def __repr__(self) -> str:
        return (
            f"LFRResult(nodes={self.graph.number_of_nodes()}, "
            f"edges={len(self.edges)}, "
            f"communities={self.n_communities})"
        )


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _nx_to_pysurprise(G: nx.Graph) -> Tuple[List[Edge], Partition]:
    """Convert a NetworkX LFR graph to pySurprise edge list + partition."""
    edges: List[Edge] = [(str(u), str(v)) for u, v in G.edges()]

    # NetworkX stores community membership as a set/frozenset on each node.
    # Assign a deterministic integer id to each unique community set.
    community_sets = {
        node: frozenset(G.nodes[node]["community"]) for node in G.nodes()
    }

    # Map each unique frozenset to a sequential integer
    unique = sorted(community_sets.values(), key=lambda s: min(s))
    seen: Dict[frozenset, int] = {}
    cid = 0
    for fs in unique:
        if fs not in seen:
            seen[fs] = cid
            cid += 1

    ground_truth: Partition = {
        str(node): seen[community_sets[node]] for node in G.nodes()
    }

    return edges, ground_truth


# ---------------------------------------------------------------------------
# Public API
# ---------------------------------------------------------------------------

def lfr_graph(
    n: int,
    tau1: float,
    tau2: float,
    mu: float,
    *,
    average_degree: Optional[float] = None,
    min_degree: Optional[int] = None,
    max_degree: Optional[int] = None,
    min_community: Optional[int] = None,
    max_community: Optional[int] = None,
    tol: float = 1.0e-7,
    max_iters: int = 500,
    seed: Optional[int] = None,
) -> LFRResult:
    """Generate an LFR benchmark graph and return an :class:`LFRResult`.

    This is the full-featured version that also keeps the NetworkX graph
    object.  For a lighter call that only returns ``(edges, ground_truth)``
    see :func:`lfr`.

    Parameters
    ----------
    n : int
        Number of nodes.
    tau1 : float
        Power-law exponent for the **degree** distribution (must be > 1).
    tau2 : float
        Power-law exponent for the **community size** distribution (must be > 1).
    mu : float
        Mixing parameter — fraction of each node's edges that are
        inter-community.  Must be in [0, 1].
    average_degree : float, optional
        Desired average degree.  Exactly one of *average_degree* and
        *min_degree* must be given.
    min_degree : int, optional
        Minimum node degree.  Exactly one of *average_degree* and
        *min_degree* must be given.
    max_degree : int, optional
        Maximum node degree.  Defaults to *n*.
    min_community : int, optional
        Minimum community size.  Defaults to *min_degree*.
    max_community : int, optional
        Maximum community size.  Defaults to *max_degree*.
    tol : float
        Tolerance for comparing floats when fitting the degree sequence.
    max_iters : int
        Maximum iterations for generating the community structure.
    seed : int, optional
        Random seed for reproducibility.

    Returns
    -------
    LFRResult
        Dataclass with ``edges``, ``ground_truth``, ``graph``, and
        ``n_communities``.

    Raises
    ------
    ImportError
        If ``networkx`` is not installed.
    networkx.exception.ExceededMaxIterations
        If the algorithm cannot produce a valid graph within *max_iters*.

    Examples
    --------
    >>> from pysurprise.benchmarks import lfr_graph
    >>> result = lfr_graph(250, tau1=3, tau2=1.5, mu=0.1,
    ...                    average_degree=5, min_community=20, seed=42)
    >>> result
    LFRResult(nodes=250, edges=..., communities=...)
    >>> result.edges[:2]          # pySurprise format
    [('0', '...'), ...]
    >>> result.ground_truth['0']  # community id for node 0
    ...
    """
    G = LFR_benchmark_graph(
        n,
        tau1,
        tau2,
        mu,
        average_degree=average_degree,
        min_degree=min_degree,
        max_degree=max_degree,
        min_community=min_community,
        max_community=max_community,
        tol=tol,
        max_iters=max_iters,
        seed=seed,
    )

    edges, ground_truth = _nx_to_pysurprise(G)

    n_communities = len({v for v in ground_truth.values()})

    return LFRResult(
        edges=edges,
        ground_truth=ground_truth,
        graph=G,
        n_communities=n_communities,
    )


def lfr(
    n: int = 250,
    tau1: float = 3.0,
    tau2: float = 1.5,
    mu: float = 0.1,
    *,
    average_degree: Optional[float] = 5.0,
    min_degree: Optional[int] = None,
    max_degree: Optional[int] = None,
    min_community: Optional[int] = 20,
    max_community: Optional[int] = None,
    tol: float = 1.0e-7,
    max_iters: int = 500,
    seed: Optional[int] = None,
) -> Tuple[List[Edge], Partition]:
    """Generate an LFR benchmark and return ``(edges, ground_truth)``.

    This is the convenience wrapper.  It uses sensible defaults so you can
    call ``lfr()`` with zero arguments for a quick test network.

    Parameters
    ----------
    n : int
        Number of nodes (default 250).
    tau1 : float
        Degree-distribution exponent (default 3.0).
    tau2 : float
        Community-size exponent (default 1.5).
    mu : float
        Mixing parameter (default 0.1 → strong communities).
    average_degree : float, optional
        Target average degree (default 5).
    min_degree : int, optional
        Alternative to *average_degree*.
    max_degree : int, optional
        Maximum node degree.
    min_community : int, optional
        Minimum community size (default 20).
    max_community : int, optional
        Maximum community size.
    tol : float
        Float comparison tolerance.
    max_iters : int
        Max iterations for internal generation loops.
    seed : int, optional
        Random seed for reproducibility.

    Returns
    -------
    edges : list of (str, str)
        Edge list in pySurprise ``.pairs`` format.
    ground_truth : dict[str, int]
        ``{node_label: community_id}``

    Examples
    --------
    >>> from pysurprise.benchmarks import lfr
    >>> edges, truth = lfr()
    >>> len(edges)
    ...
    >>> import pysurprise
    >>> pysurprise.surprise(edges, truth)
    ...
    """
    result = lfr_graph(
        n,
        tau1,
        tau2,
        mu,
        average_degree=average_degree,
        min_degree=min_degree,
        max_degree=max_degree,
        min_community=min_community,
        max_community=max_community,
        tol=tol,
        max_iters=max_iters,
        seed=seed,
    )
    return result.edges, result.ground_truth
