"""
LFR benchmark graph generator for pySurprise.
=============================================

Provides a clean Python API for generating LFR (Lancichinetti–Fortunato–
Radicchi) benchmark graphs and their *closed* variants as described in:

- Lancichinetti A, Fortunato S, Radicchi F (2008) *Benchmark graphs for
  testing community detection algorithms.* Phys Rev E 78: 046110.
- Aldecoa R, Marín I (2012) *Closed benchmarks for network community
  structure characterization.* Phys Rev E 85: 026109.
- Aldecoa R, Marín I (2013) *Surprise maximization reveals the community
  structure of complex networks.* Sci Rep 3: 1060.

Quick start::

    from pysurprise.benchmarks.lfr_benchmark import lfr_benchmark, LFR_PAPER_PRESETS

    # Generate a standard LFR network
    edges, gt = lfr_generate(1000, mu=0.3, average_degree=20, max_degree=50,
                             min_community=20, max_community=100, seed=42)

    # Generate + closed degradation at 30%
    edges, gt = lfr_benchmark(1000, mu=0.3, degradation=0.3, mode="closed",
                              average_degree=20, max_degree=50,
                              min_community=20, max_community=100, seed=42)

    # Use a paper preset
    edges, gt = lfr_benchmark(**LFR_PAPER_PRESETS["2013_scirep_N1000"],
                              mu=0.3, degradation=0.0, seed=42)

Backend
-------
Uses the **native C++ binary** from Lancichinetti et al., which produces
exact μ values and supports τ₂ = 1.  The binary is auto-discovered from the
repository layout, or you can set its path via ``set_lfr_binary(path)`` or
the ``PYSURPRISE_LFR_BINARY`` environment variable.

To build it::

    cd LFRbenchmarks-main/unweighted_undirected && make
"""

from __future__ import annotations

import logging
import os
import random
import shutil
import subprocess
import tempfile
from pathlib import Path
from typing import Dict, List, Optional, Tuple

logger = logging.getLogger(__name__)

# ── Type aliases ─────────────────────────────────────────────────────────────
Edge = Tuple[str, str]
Partition = Dict[str, int]

# ── Binary discovery ─────────────────────────────────────────────────────────
_lfr_binary_path: Optional[Path] = None


def _discover_binary() -> Optional[Path]:
    """Try to find the native LFR binary automatically."""
    # 1. Explicit env var
    env = os.environ.get("PYSURPRISE_LFR_BINARY")
    if env:
        p = Path(env)
        if p.is_file() and os.access(p, os.X_OK):
            return p

    # 2. Relative to this package (development layout)
    pkg_root = Path(__file__).resolve().parent.parent.parent.parent.parent
    candidate = pkg_root / "LFRbenchmarks-main" / "unweighted_undirected" / "benchmark"
    if candidate.is_file() and os.access(candidate, os.X_OK):
        return candidate

    # 3. On PATH
    which = shutil.which("lfr_benchmark") or shutil.which("benchmark")
    if which:
        return Path(which)

    return None


def set_lfr_binary(path: str) -> None:
    """Set the path to the native LFR benchmark binary.

    Parameters
    ----------
    path : str
        Absolute or relative path to the compiled ``benchmark`` executable
        from ``LFRbenchmarks-main/unweighted_undirected/``.
    """
    global _lfr_binary_path
    p = Path(path).resolve()
    if not p.is_file():
        raise FileNotFoundError(f"LFR binary not found: {p}")
    if not os.access(p, os.X_OK):
        raise PermissionError(f"LFR binary not executable: {p}")
    _lfr_binary_path = p


def get_lfr_binary() -> Optional[Path]:
    """Return the current LFR binary path, or None if not set/found."""
    global _lfr_binary_path
    if _lfr_binary_path is None:
        _lfr_binary_path = _discover_binary()
    return _lfr_binary_path


# ── Native binary wrapper ───────────────────────────────────────────────────

def _run_native_lfr(
    n: int,
    mu: float,
    *,
    average_degree: int = 20,
    max_degree: int = 50,
    tau1: float = 2.0,
    tau2: float = 1.0,
    min_community: int = 20,
    max_community: int = 100,
    seed: int = 0,
    timeout: int = 120,
) -> Tuple[List[Edge], Partition, List[int]]:
    """Run the native LFR binary and parse its output.

    Returns
    -------
    edges : list of (str, str)
    ground_truth : dict[str, int]  (0-based community IDs)
    community_sizes : list of int  (sorted ascending)
    """
    binary = get_lfr_binary()
    if binary is None:
        raise FileNotFoundError(
            "Native LFR binary not found. Either:\n"
            "  1. Build it: cd LFRbenchmarks-main/unweighted_undirected && make\n"
            "  2. Set path: set_lfr_binary('/path/to/benchmark')\n"
            "  3. Set env:  PYSURPRISE_LFR_BINARY=/path/to/benchmark"
        )

    td = tempfile.mkdtemp(prefix="lfr_")
    try:
        # Write seed file
        with open(os.path.join(td, "time_seed.dat"), "w") as f:
            f.write(str(seed) + "\n")

        cmd = [
            str(binary),
            "-N", str(n),
            "-k", str(average_degree),
            "-maxk", str(max_degree),
            "-mu", str(mu),
            "-t1", str(tau1),
            "-t2", str(tau2),
            "-minc", str(min_community),
            "-maxc", str(max_community),
        ]

        result = subprocess.run(
            cmd, cwd=td, capture_output=True, text=True,
            timeout=timeout, check=False,
        )

        if result.returncode != 0:
            raise RuntimeError(
                f"LFR binary failed (exit {result.returncode}):\n"
                f"  cmd: {' '.join(cmd)}\n"
                f"  stderr: {result.stderr.strip()}\n"
                f"  stdout: {result.stdout.strip()}"
            )

        # Parse network.dat → edges (1-indexed, listed both directions → deduplicate)
        edges: List[Edge] = []
        seen = set()
        net_path = os.path.join(td, "network.dat")
        with open(net_path) as f:
            for line in f:
                parts = line.split()
                if len(parts) >= 2:
                    u, v = int(parts[0]), int(parts[1])
                    if u < v:
                        key = (u, v)
                    else:
                        key = (v, u)
                    if key not in seen:
                        seen.add(key)
                        edges.append((str(key[0]), str(key[1])))

        # Parse community.dat → ground-truth partition
        gt_raw: Dict[str, int] = {}
        comm_path = os.path.join(td, "community.dat")
        with open(comm_path) as f:
            for line in f:
                parts = line.split()
                if len(parts) >= 2:
                    gt_raw[str(int(parts[0]))] = int(parts[1])

        # Remap community IDs to 0-based
        unique_comms = sorted(set(gt_raw.values()))
        remap = {old: new for new, old in enumerate(unique_comms)}
        gt: Partition = {k: remap[v] for k, v in gt_raw.items()}

        # Community sizes
        from collections import Counter
        comm_counts = Counter(gt.values())
        sizes = sorted(comm_counts.values())

        return edges, gt, sizes
    finally:
        shutil.rmtree(td, ignore_errors=True)



# ═══════════════════════════════════════════════════════════════════════════════
# PUBLIC API
# ═══════════════════════════════════════════════════════════════════════════════

def lfr_generate(
    n_nodes: int,
    mu: float,
    *,
    average_degree: int = 20,
    max_degree: int = 50,
    tau1: float = 2.0,
    tau2: float = 1.0,
    min_community: int = 20,
    max_community: int = 100,
    seed: Optional[int] = None,
) -> Tuple[List[Edge], Partition, List[int]]:
    """Generate an LFR benchmark graph.

    Parameters
    ----------
    n_nodes : int
        Number of nodes.
    mu : float
        Mixing parameter (fraction of inter-community edges per node).
        0 = perfect communities, 1 = random.
    average_degree : int
        Target average degree (default 20).
    max_degree : int
        Maximum node degree (default 50).
    tau1 : float
        Power-law exponent for degree distribution (default 2.0).
    tau2 : float
        Power-law exponent for community size distribution (default 1.0).
    min_community : int
        Minimum community size (default 20).
    max_community : int
        Maximum community size (default 100).
    seed : int, optional
        Random seed for reproducibility.

    Returns
    -------
    edges : list of (str, str)
        Edge list in pySurprise format.
    ground_truth : dict[str, int]
        ``{node_label: community_id}`` with 0-based IDs.
    community_sizes : list of int
        Community sizes sorted ascending.
    """
    if seed is None:
        seed = random.randint(0, 2**31 - 1)

    return _run_native_lfr(
        n_nodes, mu,
        average_degree=average_degree, max_degree=max_degree,
        tau1=tau1, tau2=tau2,
        min_community=min_community, max_community=max_community,
        seed=seed,
    )


# ── Closed degradation ──────────────────────────────────────────────────────
# Port of closed_benchmark.pl by Aldecoa & Marín (2012)

def lfr_degrade_closed(
    edges: List[Edge],
    ground_truth: Partition,
    degradation: float,
    *,
    seed: Optional[int] = None,
) -> Tuple[List[Edge], Partition]:
    """Apply closed degradation to an LFR network.

    This is a Python port of ``closed_benchmark.pl`` from Aldecoa & Marín
    (2012).  It permutes node labels to create a *target* community structure
    with the same sizes, then progressively swaps edges from the initial graph
    toward the target graph.

    The total number of edges is **preserved** (closed benchmark).

    Parameters
    ----------
    edges : list of (str, str)
        Initial edge list.
    ground_truth : dict[str, int]
        Initial community assignment ``{node_label: community_id}``.
    degradation : float
        Fraction of swappable edges to swap, in [0, 1].
    seed : int, optional
        Random seed.

    Returns
    -------
    degraded_edges : list of (str, str)
        Edge list after swapping.
    final_ground_truth : dict[str, int]
        The *target* (final) community assignment.  At D=1.0 the degraded
        graph perfectly matches this partition.
    """
    if not 0.0 <= degradation <= 1.0:
        raise ValueError(f"degradation must be in [0, 1], got {degradation}")

    rng = random.Random(seed)

    # ── Build adjacency set ──────────────────────────────────────────────
    labels = sorted(ground_truth.keys(), key=lambda x: int(x))
    n_nodes = len(labels)

    adj = set()
    for u, v in edges:
        a, b = (u, v) if u < v else (v, u)
        adj.add((a, b))

    # ── Create target communities by permuting node labels ───────────────
    # (Fisher-Yates shuffle matching Perl implementation)
    final_labels = list(labels)
    for i in range(len(final_labels) - 1, 0, -1):
        j = int(rng.random() * (i + 1))
        if i != j:
            final_labels[i], final_labels[j] = final_labels[j], final_labels[i]

    # Mapping: initial_label → final_label (and inverse)
    i2f = {labels[i]: final_labels[i] for i in range(n_nodes)}
    f2i = {final_labels[i]: labels[i] for i in range(n_nodes)}

    # Final ground truth: final_node gets the community of the initial node
    # it maps from.  i.e. if initial node A was in community C, and A maps
    # to final node F, then F belongs to community C in the final partition.
    final_gt: Partition = {}
    for orig_label in labels:
        final_label = i2f[orig_label]
        final_gt[final_label] = ground_truth[orig_label]

    if degradation == 0.0:
        return list(edges), final_gt

    # ── Find swappable edges ─────────────────────────────────────────────
    # An edge (A, B) in the initial graph maps to (i2f[A], i2f[B]) in the
    # final graph.  If (i2f[A], i2f[B]) is NOT already in the graph, then
    # (i2f[A], i2f[B]) is a "missing" edge that needs to be added.
    # We collect these as the pool of edges to add (@noGraph in Perl).
    edges_to_add = []  # Edges in final graph not in initial graph
    for i in range(n_nodes):
        for j in range(i + 1, n_nodes):
            a_init, b_init = labels[i], labels[j]
            a_final, b_final = i2f[a_init], i2f[b_init]

            # Check if (a_init, b_init) is in initial graph
            key_init = (min(a_init, b_init), max(a_init, b_init))
            if key_init not in adj:
                continue

            # Check if mapped edge (a_final, b_final) is NOT in initial graph
            key_final = (min(a_final, b_final), max(a_final, b_final))
            if key_final not in adj:
                edges_to_add.append(key_final)

    n_swappable = len(edges_to_add)

    if n_swappable == 0:
        # Final structure is identical to initial (very rare)
        return list(edges), final_gt

    # ── Working edge list (as list for random-access splice) ─────────────
    graph = list(adj)  # Current edges as sorted tuples

    # How many edges to swap at this degradation level
    target_swapped = int(round(degradation * n_swappable))
    if target_swapped == 0:
        return [(u, v) for u, v in graph], final_gt

    swapped = 0
    edges_to_add_list = list(edges_to_add)

    while swapped < target_swapped and edges_to_add_list:
        # Pick a random edge from current graph
        r = int(rng.random() * len(graph))
        edge = graph[r]

        # Check if this edge's inverse mapping is NOT in the initial graph
        # (i.e., it's an edge that "shouldn't" be there in the final graph)
        u_init = f2i.get(edge[0], edge[0])
        v_init = f2i.get(edge[1], edge[1])
        key_check = (min(u_init, v_init), max(u_init, v_init))

        if key_check not in adj:
            # This edge can be removed — swap it
            graph.pop(r)
            # Pick a random edge to add from the pool
            r2 = int(rng.random() * len(edges_to_add_list))
            new_edge = edges_to_add_list.pop(r2)
            graph.append(new_edge)
            swapped += 1

    return [(u, v) for u, v in graph], final_gt


# ── Main entry point ─────────────────────────────────────────────────────────

def lfr_benchmark(
    n_nodes: int,
    mu: float,
    *,
    degradation: float = 0.0,
    mode: str = "open",
    average_degree: int = 20,
    max_degree: int = 50,
    tau1: float = 2.0,
    tau2: float = 1.0,
    min_community: int = 20,
    max_community: int = 100,
    seed: Optional[int] = None,
) -> Tuple[List[Edge], Partition]:
    """Generate an LFR benchmark, optionally with closed degradation.

    Parameters
    ----------
    n_nodes : int
        Number of nodes.
    mu : float
        Mixing parameter (0 = perfect communities, 1 = random).
    degradation : float
        Closed degradation fraction in [0, 1].  0 = no degradation.
    mode : str
        ``"open"`` — the mixing parameter μ IS the degradation (no further
        degradation step; ``degradation`` parameter is ignored).
        ``"closed"`` — generate at μ then apply closed degradation.
    average_degree : int
        Target average degree (default 20).
    max_degree : int
        Maximum degree (default 50).
    tau1 : float
        Degree distribution exponent (default 2.0).
    tau2 : float
        Community size distribution exponent (default 1.0).
    min_community : int
        Minimum community size (default 20).
    max_community : int
        Maximum community size (default 100).
    seed : int, optional
        Random seed.

    Returns
    -------
    edges : list of (str, str)
    ground_truth : dict[str, int]
    """
    if seed is None:
        seed = random.randint(0, 2**31 - 1)

    edges, gt, _ = lfr_generate(
        n_nodes, mu,
        average_degree=average_degree,
        max_degree=max_degree,
        tau1=tau1, tau2=tau2,
        min_community=min_community,
        max_community=max_community,
        seed=seed,
    )

    if mode == "closed" and degradation > 0.0:
        edges, gt = lfr_degrade_closed(
            edges, gt, degradation, seed=seed + 1,
        )
    elif mode == "open":
        pass  # mu itself is the degradation
    else:
        if mode not in ("open", "closed"):
            raise ValueError(f"Unknown mode {mode!r}. Use 'open' or 'closed'.")

    return edges, gt


# ── Paper presets ────────────────────────────────────────────────────────────

LFR_PAPER_PRESETS: Dict[str, dict] = {
    "2011_N1000": {
        "n_nodes": 1000,
        "average_degree": 20,
        "max_degree": 50,
        "tau1": 2.0,
        "tau2": 1.0,
        "min_community": 20,
        "max_community": 100,
    },
    "2011_N5000": {
        "n_nodes": 5000,
        "average_degree": 20,
        "max_degree": 50,
        "tau1": 2.0,
        "tau2": 1.0,
        "min_community": 10,
        "max_community": 50,
    },
    "2012_closed_N1000": {
        "n_nodes": 1000,
        "average_degree": 20,
        "max_degree": 50,
        "tau1": 2.0,
        "tau2": 1.0,
        "min_community": 20,
        "max_community": 100,
    },
    "2013_scirep_N1000": {
        "n_nodes": 1000,
        "average_degree": 20,
        "max_degree": 50,
        "tau1": 2.0,
        "tau2": 1.0,
        "min_community": 20,
        "max_community": 100,
    },
    "2013_scirep_N5000": {
        "n_nodes": 5000,
        "average_degree": 20,
        "max_degree": 50,
        "tau1": 2.0,
        "tau2": 1.0,
        "min_community": 10,
        "max_community": 50,
    },
}
"""LFR parameter sets used in published papers.

Keys
----
- ``"2011_N1000"`` / ``"2011_N5000"`` — Aldecoa & Marín (2011), Fig. 2–3.
- ``"2012_closed_N1000"`` — Aldecoa & Marín (2012), closed LFR benchmarks.
- ``"2013_scirep_N1000"`` / ``"2013_scirep_N5000"`` — Aldecoa & Marín (2013).
"""


def describe(name: Optional[str] = None) -> str:
    """Describe LFR paper presets.

    Parameters
    ----------
    name : str, optional
        Preset name.  If ``None``, list all presets.

    Returns
    -------
    str
        Human-readable description.
    """
    descriptions = {
        "2011_N1000": (
            "Aldecoa & Marín (2011) — N=1000, k=20, maxk=50, τ₁=2, τ₂=1, "
            "minc=20, maxc=100.  Open LFR benchmark with varying μ."
        ),
        "2011_N5000": (
            "Aldecoa & Marín (2011) — N=5000, k=20, maxk=50, τ₁=2, τ₂=1, "
            "minc=10, maxc=50.  Open LFR benchmark with varying μ."
        ),
        "2012_closed_N1000": (
            "Aldecoa & Marín (2012) — N=1000, k=20, maxk=50, τ₁=2, τ₂=1, "
            "minc=20, maxc=100.  Closed LFR benchmark (constant edge count)."
        ),
        "2013_scirep_N1000": (
            "Aldecoa & Marín (2013) Sci Rep — N=1000, k=20, maxk=50, τ₁=2, "
            "τ₂=1, minc=20, maxc=100.  Open LFR with μ sweep."
        ),
        "2013_scirep_N5000": (
            "Aldecoa & Marín (2013) Sci Rep — N=5000, k=20, maxk=50, τ₁=2, "
            "τ₂=1, minc=10, maxc=50.  Open LFR with μ sweep."
        ),
    }

    if name is not None:
        if name not in descriptions:
            raise KeyError(
                f"Unknown preset {name!r}. "
                f"Available: {sorted(descriptions)}"
            )
        params = LFR_PAPER_PRESETS[name]
        lines = [descriptions[name], "", "Parameters:"]
        for k, v in params.items():
            lines.append(f"  {k} = {v}")
        return "\n".join(lines)

    lines = ["Available LFR paper presets:", ""]
    for key in sorted(descriptions):
        lines.append(f"  {key!r}:")
        lines.append(f"    {descriptions[key]}")
        lines.append("")
    return "\n".join(lines)


# ── Backward-compatible aliases ──────────────────────────────────────────────

def lfr(
    n: int = 1000,
    tau1: float = 2.0,
    tau2: float = 1.0,
    mu: float = 0.1,
    *,
    average_degree: int = 20,
    max_degree: int = 50,
    min_community: int = 20,
    max_community: int = 100,
    seed: Optional[int] = None,
    **_kwargs,
) -> Tuple[List[Edge], Partition]:
    """Backward-compatible wrapper.  Prefer :func:`lfr_generate`."""
    edges, gt, _ = lfr_generate(
        n, mu,
        average_degree=average_degree, max_degree=max_degree,
        tau1=tau1, tau2=tau2,
        min_community=min_community, max_community=max_community,
        seed=seed,
    )
    return edges, gt
