"""
Community detection algorithms bundled with SurpriseMe.

All seven original algorithms are exposed as Python functions that return
a partition dict ``{node_label: community_id}`` ready to be passed to
:func:`pysurprise.surprise`.

Algorithms
----------
CPM, Infomap, RB, RN, RNSC, SCluster, UVCluster

Each function receives a list of ``(str, str)`` edges (the same ``.pairs``
format used by SurpriseMe) and returns ``dict[str, int]``.

A convenience function :func:`run_all` executes every algorithm on the
same edge list and returns a dict of ``{name: partition}``.
"""

from __future__ import annotations

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

# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

import sys as _sys

# Directory inside the installed package that contains compiled binaries
_PKG_BIN_DIR: Path = Path(__file__).resolve().parent / "bin"

# Map from flat binary names used in the wheel to legacy repo-relative paths
_LEGACY_PATHS = {
    "computeSurprise": "src/Surprise/computeSurprise",
    "jerarca": "src/Jerarca/jerarca",
    "cpm_community": "src/CPM/bin/community",
    "cpm_slicer": "src/CPM/bin/slicer",
    "rb_community": "src/RB/bin/community",
    "rb_slicer": "src/RB/bin/slicer",
    "rnsc": "src/RNSC/rnsc",
    "rnscconvert": "src/RNSC/rnscconvert",
    "rnscfilter": "src/RNSC/rnscfilter",
    "rnMRA": "src/RN/rnMRA",
    "infomap": "src/Infomap/src/infomap",
}


def _find_repo_root() -> Optional[Path]:
    """Walk up from this file to locate the repository root (contains src/)."""
    p = Path(__file__).resolve()
    for parent in [p] + list(p.parents):
        if (parent / "src" / "Surprise").is_dir():
            return parent
        if (parent / "pysurprise" / "src" / "pysurprise").is_dir():
            return parent
    candidate = p.parent.parent.parent.parent
    if (candidate / "src" / "Surprise").is_dir():
        return candidate
    return None


_REPO: Optional[Path] = None


def _repo() -> Optional[Path]:
    global _REPO
    if _REPO is None:
        _REPO = _find_repo_root()
    return _REPO


def _bin(name: str) -> str:
    """Return absolute path to a compiled binary.

    Lookup order:
      1. Package ``bin/`` directory (pip-installed wheel)
      2. Legacy repo-relative path (development checkout with ``make``)
    """
    ext = ".exe" if _sys.platform == "win32" else ""

    # 1. Package bin/ (wheel install)
    pkg = _PKG_BIN_DIR / f"{name}{ext}"
    if pkg.exists():
        return str(pkg)

    # 2. Legacy repo layout
    legacy_rel = _LEGACY_PATHS.get(name)
    if legacy_rel:
        repo = _repo()
        if repo is not None:
            legacy = repo / legacy_rel
            if legacy.exists():
                return str(legacy)

    raise FileNotFoundError(
        f"Binary '{name}' not found.\n"
        f"  Checked package dir: {_PKG_BIN_DIR}\n"
        f"  Checked repo root:   {_repo()}\n"
        f"Install from wheel (pip install pysurprise) or run 'make' in the repo."
    )


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


def _write_pairs(edges: Sequence[Edge], path: Path) -> None:
    with open(path, "w") as f:
        for u, v in edges:
            f.write(f"{u} {v}\n")


def _read_part_file(path: Path) -> Partition:
    """Read a SurpriseMe ``.part`` file (header line + node\\tcommunity)."""
    part: Partition = {}
    with open(path) as f:
        lines = f.read().splitlines()
    for line in lines[1:]:  # skip header
        line = line.strip()
        if not line:
            continue
        cols = line.split()
        if len(cols) >= 2:
            part[cols[0]] = int(cols[1])
    return part


def _node_index_maps(edges: Sequence[Edge]):
    """Build node<->index mappings preserving insertion order."""
    node2idx: Dict[str, int] = {}
    idx2node: Dict[int, str] = {}
    idx = 0
    for u, v in edges:
        if u not in node2idx:
            node2idx[u] = idx
            idx2node[idx] = u
            idx += 1
        if v not in node2idx:
            node2idx[v] = idx
            idx2node[idx] = v
            idx += 1
    return node2idx, idx2node


def _write_indexed_edges(edges: Sequence[Edge], path: Path,
                         node2idx: Dict[str, int]) -> None:
    with open(path, "w") as f:
        for u, v in edges:
            f.write(f"{node2idx[u]}\t{node2idx[v]}\n")


def _write_slicer_input(edges: Sequence[Edge], path: Path,
                        node2idx: Dict[str, int]) -> None:
    """Write the weird CPM/RB slicer format."""
    n_nodes = len(node2idx)
    with open(path, "w") as f:
        f.write(">\n")
        for i in range(n_nodes):
            f.write(f"{i} {i}\n")
        f.write(">\n1 1\n>\n")
        for u, v in edges:
            f.write(f"{node2idx[u]} {node2idx[v]} 1 1\n")


def _write_gml(edges: Sequence[Edge], path: Path,
               node2idx: Dict[str, int]) -> None:
    with open(path, "w") as f:
        f.write("graph[\n")
        for label, idx in sorted(node2idx.items(), key=lambda x: x[1]):
            f.write(f"\tnode [\n\t\tid {idx}\n\t\tlabel {label}\n]\n")
        for u, v in edges:
            f.write(f"\tedge [\n\t\tsource {node2idx[u]}\n"
                    f"\t\ttarget {node2idx[v]}\n\t\tvalue 1\n\t]\n")
        f.write("]\n")


def _write_pajek(edges: Sequence[Edge], path: Path,
                  node2idx: Dict[str, int]) -> None:
    """Write Pajek format matching ``txt2pajek.pl`` (1-indexed).

    Builds its own consecutive 1-based index from lexicographically sorted
    labels, exactly as the Perl script does.
    """
    # txt2pajek.pl: sort keys %nodes → assign $index = 1, 2, 3, ...
    sorted_labels = sorted(node2idx.keys())
    paj_idx = {label: i + 1 for i, label in enumerate(sorted_labels)}
    n_nodes = len(paj_idx)

    with open(path, "w") as f:
        f.write(f"*Vertices {n_nodes}\n")
        for label in sorted_labels:
            f.write(f'{paj_idx[label]} "{label}"\n')
        f.write(f"*Edges {len(edges)}\n")
        for u, v in edges:
            f.write(f"{paj_idx[u]} {paj_idx[v]}\n")


def _run(cmd: List[str], cwd: str, timeout: int = 120) -> None:
    subprocess.run(
        cmd, cwd=cwd, timeout=timeout,
        stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
        check=False,
    )


# ---------------------------------------------------------------------------
# Algorithm wrappers
# ---------------------------------------------------------------------------

def cpm(edges: Sequence[Edge], *, timeout: int = 120) -> Partition:
    """CPM — Constant Potts Model (Traag et al. 2011).

    Scans the gamma parameter space and picks the partition that maximises
    Surprise, exactly as the original SurpriseMe CPM script does.
    """
    node2idx, idx2node = _node_index_maps(edges)
    slicer = _bin("cpm_slicer")
    community_bin = _bin("cpm_community")
    compute_s = _bin("computeSurprise")

    with tempfile.TemporaryDirectory() as td:
        pairs = os.path.join(td, "net.pairs")
        _write_pairs(edges, Path(pairs))
        idx_file = os.path.join(td, "idx.txt")
        _write_indexed_edges(edges, Path(idx_file), node2idx)
        slicer_in = os.path.join(td, "slicer.in")
        _write_slicer_input(edges, Path(slicer_in), node2idx)
        bin_f = os.path.join(td, "net.bin")
        conf_f = os.path.join(td, "net.conf")
        tmp_out = os.path.join(td, "tmp_out")

        _run([slicer, "-C", "-i", slicer_in, "-o", bin_f,
              "-c", conf_f, "-n", idx_file], cwd=td, timeout=timeout)

        best_part: Partition = {}
        max_s = -1.0
        step = 0.5
        times_limit = 100
        count = 0
        converged = False

        while not converged:
            gamma = 0.001
            while gamma <= 1.0:
                _run([community_bin, bin_f, conf_f, "-pp", str(gamma),
                      "-o", tmp_out], cwd=td, timeout=timeout)
                if os.path.exists(tmp_out):
                    comm: Partition = {}
                    with open(tmp_out) as f:
                        for line in f:
                            cols = line.strip().split()
                            if len(cols) >= 2:
                                nidx = int(cols[0])
                                if nidx in idx2node:
                                    comm[idx2node[nidx]] = int(cols[1])
                    # Write .part and compute Surprise
                    part_f = os.path.join(td, "cpm.part")
                    with open(part_f, "w") as pf:
                        pf.write("CPM\n")
                        for node in sorted(comm):
                            pf.write(f"{node}\t{comm[node]}\n")
                    result = subprocess.run(
                        [compute_s, pairs, part_f],
                        capture_output=True, text=True, cwd=td, timeout=timeout
                    )
                    s_val = 0.0
                    for rline in result.stdout.splitlines():
                        if "Surprise" in rline:
                            s_val = float(rline.split("=")[1].strip())
                    if s_val > max_s:
                        max_s = s_val
                        best_part = dict(comm)
                        count = 0
                    else:
                        count += 1
                        if count > times_limit:
                            converged = True
                            break
                gamma += step
            step /= 2

        return best_part


def infomap(edges: Sequence[Edge], *, attempts: int = 10,
            timeout: int = 120) -> Partition:
    """Infomap — map equation (Rosvall & Bergstrom 2008).

    Wraps the original ``infomap`` binary from the SurpriseMe distribution,
    matching ``runInfomap.pl`` exactly: convert to Pajek format, run binary
    with random seed and *attempts* trials, parse ``.tree`` output.
    """
    infomap_bin = _bin("infomap")
    node2idx, idx2node = _node_index_maps(edges)
    n_nodes = len(node2idx)

    with tempfile.TemporaryDirectory() as td:
        pajek_path = os.path.join(td, "net.pajek")
        _write_pajek(edges, Path(pajek_path), node2idx)

        seed = random.randint(0, 999_999_999)
        # Usage: infomap <seed> <pajek_file> <attempts>
        _run([infomap_bin, str(seed), pajek_path, str(attempts)],
             cwd=td, timeout=timeout)

        # Binary produces net.tree in the same directory
        tree_path = os.path.join(td, "net.tree")
        if not os.path.exists(tree_path):
            return {idx2node[i]: i for i in range(n_nodes)}

        # Parse .tree: lines like "1:1 0.1234 \"node_label\""
        # Skip lines containing "modules"
        part: Partition = {}
        with open(tree_path) as f:
            for line in f:
                if "modules" in line:
                    continue
                line = line.replace('"', '')
                cols = line.split()
                if len(cols) >= 3:
                    # cols[0] = "comm:subcomm", cols[1] = flow, cols[2] = label
                    comm_str = cols[0].split(":")[0]
                    label = cols[2]
                    try:
                        comm_id = int(comm_str)
                        part[label] = comm_id
                    except ValueError:
                        continue

        if not part:
            return {idx2node[i]: i for i in range(n_nodes)}
        return part


def rb(edges: Sequence[Edge], *, timeout: int = 120) -> Partition:
    """RB — Reichardt-Bornholdt (2006).

    Scans gamma and picks the partition maximising Surprise.
    """
    node2idx, idx2node = _node_index_maps(edges)
    slicer = _bin("rb_slicer")
    community_bin = _bin("rb_community")
    compute_s = _bin("computeSurprise")

    with tempfile.TemporaryDirectory() as td:
        pairs = os.path.join(td, "net.pairs")
        _write_pairs(edges, Path(pairs))
        idx_file = os.path.join(td, "idx.txt")
        _write_indexed_edges(edges, Path(idx_file), node2idx)
        slicer_in = os.path.join(td, "slicer.in")
        _write_slicer_input(edges, Path(slicer_in), node2idx)
        bin_f = os.path.join(td, "net.bin")
        conf_f = os.path.join(td, "net.conf")
        tmp_out = os.path.join(td, "tmp_out")

        _run([slicer, "-i", slicer_in, "-o", bin_f,
              "-c", conf_f, "-n", idx_file], cwd=td, timeout=timeout)

        best_part: Partition = {}
        max_s = -1.0
        step = 1.0
        times_limit = 100
        count = 0
        converged = False

        while not converged:
            gamma = 0.01
            while gamma < 30.0:
                _run([community_bin, bin_f, conf_f, "-pp", str(gamma),
                      "-o", tmp_out], cwd=td, timeout=timeout)
                if os.path.exists(tmp_out):
                    comm: Partition = {}
                    with open(tmp_out) as f:
                        for line in f:
                            cols = line.strip().split()
                            if len(cols) >= 2:
                                nidx = int(cols[0])
                                if nidx in idx2node:
                                    comm[idx2node[nidx]] = int(cols[1])
                    part_f = os.path.join(td, "rb.part")
                    with open(part_f, "w") as pf:
                        pf.write("RB\n")
                        for node in sorted(comm):
                            pf.write(f"{node}\t{comm[node]}\n")
                    result = subprocess.run(
                        [compute_s, pairs, part_f],
                        capture_output=True, text=True, cwd=td, timeout=timeout
                    )
                    s_val = 0.0
                    for rline in result.stdout.splitlines():
                        if "Surprise" in rline:
                            s_val = float(rline.split("=")[1].strip())
                    if s_val > max_s:
                        max_s = s_val
                        best_part = dict(comm)
                        count = 0
                    else:
                        count += 1
                        if count > times_limit:
                            converged = True
                            break
                gamma += step
            step /= 2

        return best_part


def rn(edges: Sequence[Edge], *, timeout: int = 120) -> Partition:
    """RN — Ronhovde-Nussinov (2009) via the ``rnMRA`` C++ binary.

    Calls the compiled ``rnMRA`` multiresolution solver, replicating the
    logic of ``src/scripts/runRN.pl``.  The binary is located via
    :func:`_bin`.  Falls back to the Python reimplementation if the
    binary is not found.
    """
    try:
        return _rn_binary(edges, timeout=timeout)
    except FileNotFoundError:
        from pysurprise._rn_solver import solve as _rn_solve
        return _rn_solve(edges)


def rn_python(edges: Sequence[Edge], *, timeout: int = 120) -> Partition:
    """RN — Ronhovde-Nussinov (2009), pure-Python reimplementation.

    Faithful Python reimplementation of the ``rnMRA`` greedy solver.
    See :mod:`pysurprise._rn_solver` for full algorithm documentation.
    """
    from pysurprise._rn_solver import solve as _rn_solve
    return _rn_solve(edges)


def _rn_binary(edges: Sequence[Edge], *, timeout: int = 300) -> Partition:
    """Run the rnMRA binary (Ronhovde & Nussinov 2009).

    Replicates the logic of ``src/scripts/runRN.pl``:
    1. Write edges as GML
    2. Run ``rnMRA -n:N -nsm:max(N,100) -rseed:SEED -inf:F.gml -outf:F``
    3. Parse ``FBest.txt`` output
    4. Map internal node IDs back to original labels
    """
    rnmra = _bin("rnMRA")
    node2idx, idx2node = _node_index_maps(edges)
    n_nodes = len(node2idx)
    nsm = max(n_nodes, 100)
    seed = random.randint(1, 999999999)

    with tempfile.TemporaryDirectory() as td:
        gml_path = os.path.join(td, "net.gml")
        out_root = os.path.join(td, "net")

        _write_gml(edges, Path(gml_path), node2idx)

        _run([rnmra, f"-n:{n_nodes}", f"-nsm:{nsm}",
              f"-rseed:{seed}", f"-inf:{gml_path}", f"-outf:{out_root}"],
             cwd=td, timeout=timeout)

        best_path = out_root + "Best.txt"
        if not os.path.exists(best_path):
            raise RuntimeError("rnMRA did not produce output")

        csv_path = out_root + ".csv"
        if os.path.exists(csv_path):
            os.remove(csv_path)

        part: Partition = {}
        with open(best_path) as f:
            for comm_id, line in enumerate(f):
                for node_idx_str in line.strip().split():
                    nidx = int(node_idx_str)
                    if nidx in idx2node:
                        part[idx2node[nidx]] = comm_id
        return part


def rnsc(edges: Sequence[Edge], *, timeout: int = 120) -> Partition:
    """RNSC — Restricted Neighbourhood Search Clustering (King et al. 2004)."""
    rnscconvert = _bin("rnscconvert")
    rnsc_bin = _bin("rnsc")
    rnscfilter = _bin("rnscfilter")

    with tempfile.TemporaryDirectory() as td:
        pairs = os.path.join(td, "net.pairs")
        _write_pairs(edges, Path(pairs))

        graph_f = os.path.join(td, "a.graph")
        func_f = os.path.join(td, "a.functions")

        _run([rnscconvert, "-g", graph_f, "-i", pairs, "-n", func_f],
             cwd=td, timeout=timeout)
        _run([rnsc_bin, "-g", graph_f], cwd=td, timeout=timeout)

        rnsc_out = os.path.join(td, "out.rnsc")
        _run([rnscfilter, "-g", graph_f, "-c", rnsc_out, "-n", func_f],
             cwd=td, timeout=timeout)

        output_txt = os.path.join(td, "output.txt")
        if not os.path.exists(output_txt):
            return {}

        part: Partition = {}
        with open(output_txt) as f:
            lines = f.read().splitlines()

        i = 0
        while i < len(lines):
            line = lines[i]
            if line.startswith("Cluster #"):
                import re
                m = re.search(r"#(\d+)", line)
                comm_id = int(m.group(1)) if m else i
                if i + 1 < len(lines):
                    node_line = lines[i + 1].strip()
                    nodes = [n for n in node_line.split("(U)") if n.strip()]
                    for n in nodes:
                        n = n.strip()
                        if n:
                            part[n] = comm_id
                i += 3  # skip blank line
            else:
                i += 1
        return part


def scluster(edges: Sequence[Edge], *, iterations: int = 10000,
             timeout: int = 120) -> Partition:
    """SCluster — hierarchical clustering from Jerarca (Aldecoa & Marin 2010)."""
    jerarca = _bin("jerarca")

    with tempfile.TemporaryDirectory() as td:
        pairs = os.path.join(td, "net.pairs")
        _write_pairs(edges, Path(pairs))

        _run([jerarca, pairs, "scluster", str(iterations)],
             cwd=td, timeout=timeout)

        part_file = os.path.join(td, "partition_scluster.txt")
        if not os.path.exists(part_file):
            return {}

        part: Partition = {}
        with open(part_file) as f:
            for line in f:
                line = line.strip()
                if not line or "scluster" in line.lower():
                    continue
                cols = line.split()
                if len(cols) >= 2:
                    part[cols[0]] = int(cols[1])
        return part


def uvcluster(edges: Sequence[Edge], *, iterations: Optional[int] = None,
              timeout: int = 120) -> Partition:
    """UVCluster — iterative cluster analysis from Jerarca (Arnau et al. 2005)."""
    jerarca = _bin("jerarca")
    nodes = {u for e in edges for u in e}
    if iterations is None:
        iterations = min(1000, len(nodes) * 10)

    with tempfile.TemporaryDirectory() as td:
        pairs = os.path.join(td, "net.pairs")
        _write_pairs(edges, Path(pairs))

        _run([jerarca, pairs, "uvcluster", str(iterations)],
             cwd=td, timeout=timeout)

        part_file = os.path.join(td, "partition_uvcluster.txt")
        if not os.path.exists(part_file):
            return {}

        part: Partition = {}
        with open(part_file) as f:
            for line in f:
                line = line.strip()
                if not line or "uvcluster" in line.lower():
                    continue
                cols = line.split()
                if len(cols) >= 2:
                    part[cols[0]] = int(cols[1])
        return part


def duch_arenas(edges: Sequence[Edge], *, timeout: int = 120) -> Partition:
    """DA — Duch & Arenas (2005) extremal optimization for modularity.

    Faithful Python reimplementation of the recursive bipartition algorithm.
    See :mod:`pysurprise._da_solver` for full algorithm documentation.
    """
    from pysurprise._da_solver import solve as _da_solve
    return _da_solve(edges)


# ---------------------------------------------------------------------------
# Surprise-guided greedy refinement (post-processor for any partition)
# ---------------------------------------------------------------------------

def surprise_refine(
    edges: Sequence[Edge],
    partition: Partition,
    *,
    max_iters: int = 100,
    allow_merge: bool = True,
) -> Partition:
    """Greedily refine any partition to maximise Surprise.

    Performs a Louvain-style node-relabelling pass: each node is tentatively
    moved to every other existing community (and optionally merged with every
    other community), keeping the move if it strictly increases Surprise.
    Passes repeat until no improvement is found or *max_iters* is reached.

    This is a **pure post-processor** — it is independent of the embedding
    method and can sharpen any initial partition.

    Parameters
    ----------
    edges : list of (str, str)
        Edge list in pySurprise pairs format.
    partition : dict[str, int]
        Initial community assignment ``{node: community_id}``.
    max_iters : int
        Maximum number of full sweeps (default 100).
    allow_merge : bool
        If ``True`` (default), also try merging pairs of communities after
        each node-relabelling pass.

    Returns
    -------
    dict[str, int]
        Refined partition with Surprise ≥ initial Surprise.
    """
    import pysurprise as _pys

    part = dict(partition)
    nodes = list(part.keys())
    best_s = _pys.surprise(edges, part)

    def _compact(p: Partition) -> Partition:
        """Re-label community IDs to 0, 1, 2, ... (avoids gaps)."""
        mapping = {old: new for new, old in enumerate(sorted(set(p.values())))}
        return {n: mapping[c] for n, c in p.items()}

    part = _compact(part)

    for _ in range(max_iters):
        improved = False

        # ---- node relabelling pass ----------------------------------------
        for node in nodes:
            current_c = part[node]
            communities = set(part.values())

            best_local_s = best_s
            best_local_c = current_c

            for candidate_c in communities:
                if candidate_c == current_c:
                    continue
                trial = dict(part)
                trial[node] = candidate_c
                trial = _compact(trial)
                s = _pys.surprise(edges, trial)
                if s > best_local_s:
                    best_local_s = s
                    best_local_c = candidate_c

            if best_local_c != current_c:
                part[node] = best_local_c
                part = _compact(part)
                best_s = best_local_s
                improved = True

        # ---- community merge pass -----------------------------------------
        if allow_merge:
            merge_improved = True
            while merge_improved:
                merge_improved = False
                communities = sorted(set(part.values()))
                for i in range(len(communities)):
                    for j in range(i + 1, len(communities)):
                        ca, cb = communities[i], communities[j]
                        trial = {n: ca if c == cb else c for n, c in part.items()}
                        trial = _compact(trial)
                        s = _pys.surprise(edges, trial)
                        if s > best_s:
                            part = trial
                            best_s = s
                            merge_improved = True
                            improved = True
                            break  # communities list is now stale, restart
                    if merge_improved:
                        break

        if not improved:
            break

    return _compact(part)


# ---------------------------------------------------------------------------
# Simulated Annealing Surprise optimiser
# ---------------------------------------------------------------------------

def surprise_anneal(
    edges: Sequence[Edge],
    init_partition: Optional[Partition] = None,
    *,
    T_start: float = 2.0,
    T_end: float = 1e-4,
    steps: int = 200_000,
    singleton_prob: float = 0.25,
    seed: Optional[int] = None,
) -> Partition:
    """Maximise Surprise via Simulated Annealing.

    Explores the partition space with a Metropolis–Hastings scheme: at each
    step a random node is moved to a random existing community **or** isolated
    into a new singleton (controlled by *singleton_prob*).  Moves that
    decrease Surprise are accepted with probability ``exp(ΔS / T)``.

    This is a **global optimiser** — unlike :func:`surprise_refine` it can
    escape local optima.  It is most useful for graphs where classical
    algorithms stall before reaching the optimum.

    Parameters
    ----------
    edges : list of (str, str)
        Edge list in pySurprise pairs format.
    init_partition : dict[str, int], optional
        Starting partition.  When ``None`` (default) each node starts in its
        own singleton community.
    T_start : float
        Initial temperature (default 2.0).
    T_end : float
        Final temperature (default 1e-4).  Cooling is exponential.
    steps : int
        Total number of Metropolis steps (default 200 000).
    singleton_prob : float
        Probability of proposing a singleton move at each step (default 0.25).
    seed : int, optional
        Random seed for reproducibility.

    Returns
    -------
    dict[str, int]
        Best partition found (highest Surprise seen during annealing).
    """
    import pysurprise as _pys
    import random as _random
    import math as _math

    rng = _random.Random(seed)

    # collect nodes from edges
    all_nodes = sorted({n for e in edges for n in e})
    if not all_nodes:
        return {}

    # initialise partition
    if init_partition is not None:
        part = dict(init_partition)
    else:
        part = {n: i for i, n in enumerate(all_nodes)}

    s_cur = _pys.surprise(edges, part)
    best_s, best_part = s_cur, dict(part)

    alpha = (_math.log(T_end) - _math.log(T_start)) / steps
    T = T_start

    for step in range(steps):
        T = T_start * _math.exp(alpha * step)

        node = rng.choice(all_nodes)
        old_c = part[node]
        comms = list(set(part.values()))

        if rng.random() < singleton_prob or len(comms) == 1:
            new_c = max(comms) + 1  # new singleton
        else:
            new_c = rng.choice(comms)

        if new_c == old_c:
            continue

        part[node] = new_c
        s_new = _pys.surprise(edges, part)
        delta = s_new - s_cur

        if delta >= 0 or rng.random() < _math.exp(delta / max(T, 1e-10)):
            s_cur = s_new
            if s_new > best_s:
                best_s = s_new
                best_part = dict(part)
        else:
            part[node] = old_c

    # final greedy polish
    return surprise_refine(edges, best_part, max_iters=50)


# ---------------------------------------------------------------------------
# DeepWalk
# ---------------------------------------------------------------------------

def deepwalk(
    edges: Sequence[Edge],
    *,
    n_clusters: Optional[int] = None,
    walk_length: int = 80,
    num_walks: int = 10,
    dimensions: int = 128,
    window: int = 5,
    workers: int = 1,
    seed: Optional[int] = None,
    refine: bool = False,
) -> Partition:
    """DeepWalk — graph embedding + K-Means clustering.

    Learns node embeddings via random walks and Word2Vec (Perozzi et al.
    2014), then clusters the embedding space with K-Means.

    Requires ``pip install gensim scikit-learn networkx``.

    Parameters
    ----------
    edges : list of (str, str)
        Edge list in pySurprise pairs format.
    n_clusters : int, optional
        Number of communities.  When ``None`` (default) an estimate is
        computed as ``max(2, int(sqrt(n_nodes / 2)))``.
    walk_length : int
        Number of nodes in each random walk (default 80).
    num_walks : int
        Number of random walks per node (default 10).
    dimensions : int
        Embedding dimensionality (default 128).
    window : int
        Word2Vec context window size (default 5).
    workers : int
        Number of threads for Word2Vec (default 1 for reproducibility).
    seed : int, optional
        Random seed for reproducibility.
    refine : bool
        If ``True``, apply :func:`surprise_refine` to the K-Means partition
        before returning (default ``False``).

    Returns
    -------
    dict[str, int]
        ``{node_label: community_id}``

    References
    ----------
        Perozzi B, Al-Rfou R, Skiena S (2014)
        DeepWalk: Online Learning of Social Representations.
        KDD 2014.
    """
    try:
        import networkx as nx
    except ImportError:
        raise ImportError(
            "The 'networkx' package is required for DeepWalk. "
            "Install it with: pip install networkx"
        )
    try:
        from gensim.models import Word2Vec
    except ImportError:
        raise ImportError(
            "The 'gensim' package is required for DeepWalk. "
            "Install it with: pip install gensim"
        )
    try:
        from sklearn.cluster import KMeans
    except ImportError:
        raise ImportError(
            "The 'scikit-learn' package is required for DeepWalk clustering. "
            "Install it with: pip install scikit-learn"
        )

    import random as _random

    rng = _random.Random(seed)

    # Build graph
    G = nx.Graph()
    G.add_edges_from(edges)
    nodes = list(G.nodes())
    n_nodes = len(nodes)

    if n_nodes == 0:
        return {}

    if n_clusters is None:
        n_clusters = max(2, int(n_nodes ** 0.5 // 2))

    # Generate random walks (each walk is a list of node-label strings)
    walks: List[List[str]] = []
    for _ in range(num_walks):
        shuffled = list(nodes)
        rng.shuffle(shuffled)
        for start in shuffled:
            walk = [start]
            for _ in range(walk_length - 1):
                neighbors = list(G.neighbors(walk[-1]))
                if neighbors:
                    walk.append(rng.choice(neighbors))
                else:
                    break
            walks.append([str(n) for n in walk])

    # Train Word2Vec
    w2v_seed = seed if seed is not None else 42
    model = Word2Vec(
        sentences=walks,
        vector_size=dimensions,
        window=window,
        min_count=0,
        sg=1,              # skip-gram (as in original DeepWalk)
        workers=workers,
        seed=w2v_seed,
        epochs=1,
    )

    # Build embedding matrix aligned to nodes
    node_labels = [str(n) for n in nodes]
    vectors = [model.wv[label] for label in node_labels if label in model.wv]
    valid_labels = [label for label in node_labels if label in model.wv]

    if len(valid_labels) < n_clusters:
        n_clusters = max(2, len(valid_labels))

    import numpy as _np
    X = _np.array(vectors)

    km_seed = seed if seed is not None else 42
    kmeans = KMeans(n_clusters=n_clusters, random_state=km_seed, n_init=10)
    labels = kmeans.fit_predict(X)

    part = {label: int(c) for label, c in zip(valid_labels, labels)}
    return surprise_refine(edges, part) if refine else part


# ---------------------------------------------------------------------------
# Node2Vec
# ---------------------------------------------------------------------------

def node2vec(
    edges: Sequence[Edge],
    *,
    n_clusters: Optional[int] = None,
    dimensions: int = 128,
    walk_length: int = 80,
    num_walks: int = 10,
    p: float = 1.0,
    q: float = 1.0,
    window: int = 10,
    workers: int = 1,
    seed: Optional[int] = None,
    refine: bool = False,
) -> Partition:
    """Node2Vec — biased random walk embedding + K-Means clustering.

    Learns node embeddings via second-order biased random walks (Grover &
    Leskovec 2016), then clusters the embedding space with K-Means.

    Setting ``p=1, q=1`` recovers DeepWalk.  Lower *q* favours DFS
    (community-like exploration); lower *p* favours BFS (structural roles).

    Requires ``pip install node2vec scikit-learn networkx``.

    Parameters
    ----------
    edges : list of (str, str)
        Edge list in pySurprise pairs format.
    n_clusters : int, optional
        Number of communities.  When ``None`` an estimate is computed as
        ``max(2, int(sqrt(n_nodes / 2)))``.
    dimensions : int
        Embedding dimensionality (default 128).
    walk_length : int
        Number of nodes in each random walk (default 80).
    num_walks : int
        Number of random walks per node (default 10).
    p : float
        Return parameter — controls likelihood of revisiting a node (default 1).
    q : float
        In-out parameter — controls BFS vs DFS behaviour (default 1).
    window : int
        Word2Vec context window size (default 10).
    workers : int
        Number of threads (default 1 for reproducibility; Windows requires 1).
    seed : int, optional
        Random seed for reproducibility.
    refine : bool
        If ``True``, apply :func:`surprise_refine` to the K-Means partition
        before returning (default ``False``).

    Returns
    -------
    dict[str, int]
        ``{node_label: community_id}``

    References
    ----------
        Grover A, Leskovec J (2016)
        node2vec: Scalable Feature Learning for Networks.
        KDD 2016.
    """
    try:
        import networkx as nx
    except ImportError:
        raise ImportError(
            "The 'networkx' package is required for Node2Vec. "
            "Install it with: pip install networkx"
        )
    try:
        from node2vec import Node2Vec as _Node2Vec
    except ImportError:
        raise ImportError(
            "The 'node2vec' package is required. "
            "Install it with: pip install node2vec"
        )
    try:
        from sklearn.cluster import KMeans
    except ImportError:
        raise ImportError(
            "The 'scikit-learn' package is required for Node2Vec clustering. "
            "Install it with: pip install scikit-learn"
        )

    import numpy as _np

    G = nx.Graph()
    G.add_edges_from(edges)
    nodes = list(G.nodes())
    n_nodes = len(nodes)

    if n_nodes == 0:
        return {}

    if n_clusters is None:
        n_clusters = max(2, int(n_nodes ** 0.5 // 2))

    n2v = _Node2Vec(
        G,
        dimensions=dimensions,
        walk_length=walk_length,
        num_walks=num_walks,
        p=p,
        q=q,
        workers=workers,
        seed=seed,
        quiet=True,
    )

    model = n2v.fit(
        window=window,
        min_count=1,
        batch_words=4,
        workers=workers,
        seed=seed if seed is not None else 42,
    )

    node_labels = [str(n) for n in nodes if str(n) in model.wv]
    vectors = [model.wv[label] for label in node_labels]

    if len(node_labels) < n_clusters:
        n_clusters = max(2, len(node_labels))

    X = _np.array(vectors)
    km_seed = seed if seed is not None else 42
    kmeans = KMeans(n_clusters=n_clusters, random_state=km_seed, n_init=10)
    labels = kmeans.fit_predict(X)

    part = {label: int(c) for label, c in zip(node_labels, labels)}
    return surprise_refine(edges, part) if refine else part


# ---------------------------------------------------------------------------
# VGAE — Variational Graph Auto-Encoder
# ---------------------------------------------------------------------------

def vgae(
    edges: Sequence[Edge],
    *,
    n_clusters: Optional[int] = None,
    hidden1: int = 64,
    hidden2: int = 32,
    epochs: int = 500,
    lr: float = 0.01,
    dropout: float = 0.0,
    kl_weight: float = 1.0,
    seed: Optional[int] = None,
    refine: bool = False,
) -> Partition:
    """VGAE — Variational Graph Auto-Encoder community detection.

    Learns node embeddings via a GCN-based VAE (Kipf & Welling 2016), then
    clusters the latent mean vectors with K-Means.

    Architecture::

        Input (identity features)
        → GCN layer 1  (ReLU, hidden1 units)
        → GCN μ head   (hidden2 units)   → mean embedding
        → GCN σ head   (hidden2 units)   → log variance
        → sample Z via reparameterisation
        → inner-product decoder → reconstructed adjacency
        ELBO loss = BCE(A, Â) − KL(N(μ,σ²) ‖ N(0,I))

    Requires ``pip install torch scipy scikit-learn``.

    Parameters
    ----------
    edges : list of (str, str)
        Edge list in pySurprise pairs format.
    n_clusters : int, optional
        Number of communities.  When ``None`` an estimate is computed as
        ``max(2, int(sqrt(n_nodes / 2)))``.
    hidden1 : int
        Hidden units in GCN layer 1 (default 64).
    hidden2 : int
        Latent embedding dimensionality (default 32).
    epochs : int
        Training iterations (default 500).
    lr : float
        Adam learning rate (default 0.01).
    dropout : float
        Dropout probability applied to GCN layer 1 (default 0.0).
    kl_weight : float
        Weight of the KL divergence term (β in β-VAE). Values < 1 give
        less regularisation, producing tighter clusters but less smooth
        latent space (default 1.0).
    seed : int, optional
        Random seed for reproducibility.
    refine : bool
        If ``True``, apply :func:`surprise_refine` to the K-Means partition
        before returning (default ``False``).

    Returns
    -------
    dict[str, int]
        ``{node_label: community_id}``

    References
    ----------
        Kipf T, Welling M (2016)
        Variational Graph Auto-Encoders.
        NeurIPS Workshop on Bayesian Deep Learning.
    """
    try:
        import torch
        import torch.nn as nn
        import torch.nn.functional as F
    except ImportError:
        raise ImportError(
            "PyTorch is required for VGAE. Install it with: pip install torch"
        )
    try:
        import scipy.sparse as sp
    except ImportError:
        raise ImportError(
            "scipy is required for VGAE. Install it with: pip install scipy"
        )
    try:
        import networkx as nx
    except ImportError:
        raise ImportError(
            "networkx is required for VGAE. Install it with: pip install networkx"
        )
    try:
        from sklearn.cluster import KMeans
    except ImportError:
        raise ImportError(
            "scikit-learn is required for VGAE clustering. "
            "Install it with: pip install scikit-learn"
        )

    import numpy as _np

    if seed is not None:
        torch.manual_seed(seed)
        _np.random.seed(seed)

    # ---- Build graph -------------------------------------------------------
    G = nx.Graph()
    G.add_edges_from(edges)
    node_labels = sorted(G.nodes(), key=lambda x: x)
    n = len(node_labels)
    node2idx = {v: i for i, v in enumerate(node_labels)}

    if n == 0:
        return {}
    if n_clusters is None:
        n_clusters = max(2, int(n ** 0.5 // 2))

    # ---- Normalised adjacency Â = D^-½(A+I)D^-½ ---------------------------
    adj = nx.adjacency_matrix(G, nodelist=node_labels).astype(float)
    adj_hat = adj + sp.eye(n)
    deg = _np.array(adj_hat.sum(axis=1)).flatten()
    d_inv_sqrt = sp.diags(deg ** -0.5)
    adj_norm = d_inv_sqrt @ adj_hat @ d_inv_sqrt  # sparse
    adj_norm_t = torch.FloatTensor(adj_norm.toarray())

    # ---- Target adjacency (original, without self-loops) -------------------
    adj_target = torch.FloatTensor(adj.toarray())

    # ---- Node features: identity matrix ------------------------------------
    X = torch.eye(n)

    # ---- Model definition ---------------------------------------------------
    class GCNLayer(nn.Module):
        def __init__(self, in_f: int, out_f: int):
            super().__init__()
            self.W = nn.Linear(in_f, out_f, bias=False)

        def forward(self, adj_n: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
            return self.W(adj_n @ x)

    class VGAEModel(nn.Module):
        def __init__(self, n_feat: int, h1: int, h2: int, drop: float):
            super().__init__()
            self.gc1 = GCNLayer(n_feat, h1)
            self.gc_mu = GCNLayer(h1, h2)
            self.gc_logvar = GCNLayer(h1, h2)
            self.drop = drop

        def encode(self, adj_n: torch.Tensor, x: torch.Tensor):
            h = F.relu(self.gc1(adj_n, x))
            h = F.dropout(h, p=self.drop, training=self.training)
            mu = self.gc_mu(adj_n, h)
            logvar = self.gc_logvar(adj_n, h)
            return mu, logvar

        def reparameterise(self, mu: torch.Tensor,
                           logvar: torch.Tensor) -> torch.Tensor:
            if self.training:
                std = torch.exp(0.5 * logvar)
                eps = torch.randn_like(std)
                return mu + eps * std
            return mu

        def decode(self, z: torch.Tensor) -> torch.Tensor:
            return z @ z.T  # raw logits — sigmoid applied in loss

        def forward(self, adj_n: torch.Tensor, x: torch.Tensor):
            mu, logvar = self.encode(adj_n, x)
            z = self.reparameterise(mu, logvar)
            return self.decode(z), mu, logvar

    model = VGAEModel(n, hidden1, hidden2, dropout)
    optimiser = torch.optim.Adam(model.parameters(), lr=lr)

    # ---- Training loop ------------------------------------------------------
    # Positive-weight rebalancing (sparse graphs have many more 0s than 1s)
    n_pos = float(adj_target.sum().item())
    n_neg = float(n * n - n_pos)
    pos_weight = torch.tensor([n_neg / max(n_pos, 1e-6)])
    norm = (n * n) / max(2.0 * n_neg, 1.0)

    model.train()
    for _ in range(epochs):
        optimiser.zero_grad()
        logits, mu, logvar = model(adj_norm_t, X)
        # Reconstruction loss: BCEWithLogits expects raw logits
        recon = norm * F.binary_cross_entropy_with_logits(
            logits, adj_target,
            pos_weight=pos_weight,
            reduction='mean',
        )
        # KL divergence: -0.5 * sum(1 + log(σ²) - μ² - σ²)
        kl = -0.5 / n * (1 + logvar - mu.pow(2) - logvar.exp()).sum()
        loss = recon + kl_weight * kl
        loss.backward()
        optimiser.step()

    # ---- Cluster latent means -----------------------------------------------
    model.eval()
    with torch.no_grad():
        mu_final, _ = model.encode(adj_norm_t, X)
    Z = mu_final.numpy()

    km_seed = seed if seed is not None else 42
    kmeans = KMeans(n_clusters=n_clusters, random_state=km_seed, n_init=10)
    labels = kmeans.fit_predict(Z)

    part = {lbl: int(c) for lbl, c in zip(node_labels, labels)}
    return surprise_refine(edges, part) if refine else part


# ---------------------------------------------------------------------------
# Convenience: run all
# ---------------------------------------------------------------------------

ALL_ALGORITHMS = {
    "CPM": cpm,
    "Infomap": infomap,
    "RB": rb,
    "RN": rn,
    "RNSC": rnsc,
    "SCluster": scluster,
    "UVCluster": uvcluster,
}


def run_all(
    edges: Sequence[Edge],
    *,
    algorithms: Optional[Sequence[str]] = None,
    timeout: int = 120,
) -> Dict[str, Partition]:
    """Run multiple community detection algorithms on the same edge list.

    Parameters
    ----------
    edges : list of (str, str)
        Edge list in SurpriseMe pairs format.
    algorithms : list of str, optional
        Subset of algorithm names to run.  Default: all seven.
    timeout : int
        Per-algorithm timeout in seconds.

    Returns
    -------
    dict[str, dict[str, int]]
        ``{algorithm_name: partition_dict}``
    """
    if algorithms is None:
        algorithms = list(ALL_ALGORITHMS.keys())

    results: Dict[str, Partition] = {}
    for name in algorithms:
        func = ALL_ALGORITHMS.get(name)
        if func is None:
            raise ValueError(
                f"Unknown algorithm '{name}'. "
                f"Available: {list(ALL_ALGORITHMS.keys())}"
            )
        try:
            results[name] = func(edges, timeout=timeout)
        except (FileNotFoundError, subprocess.TimeoutExpired) as exc:
            print(f"WARNING: {name} failed — {exc}")
            results[name] = {}
    return results


# ---------------------------------------------------------------------------
# Louvain Surprise maximiser (C++ accelerated)
# ---------------------------------------------------------------------------

def surprise_louvain(
    edges: Sequence[Edge],
    partition: Optional[Partition] = None,
    *,
    n_restarts: int = 8,
    seed: int = 42,
) -> Partition:
    """Maximise Surprise via a C++-accelerated Louvain-style algorithm.

    This is a direct drop-in replacement for :func:`surprise_refine` and
    :func:`surprise_anneal` with substantially better performance:

    * **Log-factorial precomputation** — each Surprise evaluation is O(1)
      per hypergeometric term instead of O(n) via ``sumRange``.
    * **Incremental ΔM / Δp** — community statistics are updated in O(1)
      without rescanning the full partition on every candidate move.
    * **Neighbour-restricted proposals** — only communities adjacent to the
      moved node are tried (O(degree) vs O(k) communities).
    * **Community merge pass** — pairs of communities are tested for exact-
      Surprise improvement after each node-movement sweep, escaping optima
      that are invisible at the single-node level.
    * **Pre_CS initialisation** (tang2019) — nodes are seeded by common-
      neighbour similarity, giving the optimiser a topology-aware start.
    * **Multi-restart** — *n_restarts* independent runs (alternating Pre_CS
      and all-singleton initialisation); the best partition is returned.

    Parameters
    ----------
    edges : list of (str, str)
        Edge list in pySurprise pairs format.
    partition : dict[str, int], optional
        Ignored — the C++ optimiser always finds its own starting point.
        Kept for API symmetry with the other optimisers.
    n_restarts : int
        Number of independent optimisation runs (default 8).
    seed : int
        RNG seed for reproducibility (default 42).

    Returns
    -------
    dict[str, int]
        Best partition found (highest exact Surprise across all restarts).
    """
    from pysurprise._core import louvain_surprise as _louvain_cpp

    # Collect and sort nodes so the index mapping is deterministic
    all_nodes: List[str] = sorted({n for e in edges for n in e})
    if not all_nodes:
        return {}

    node_to_idx: Dict[str, int] = {n: i for i, n in enumerate(all_nodes)}
    int_edges: List[Tuple[int, int]] = [
        (node_to_idx[u], node_to_idx[v]) for u, v in edges
    ]

    raw: List[int] = _louvain_cpp(int_edges, len(all_nodes), n_restarts, seed)

    # Re-map integer partition back to string node labels
    return {all_nodes[i]: raw[i] for i in range(len(all_nodes))}


# ---------------------------------------------------------------------------
# FAGSO  (Jiang et al. 2014)
# ---------------------------------------------------------------------------

def surprise_fagso(
    edges: Sequence[Edge],
    partition: Optional[Partition] = None,
) -> Partition:
    """Maximise Surprise via FAGSO (Fast Algorithm for Greedy Surprise Optimisation).

    Ranks every edge by Jaccard similarity then iterates in decreasing order,
    merging communities whenever the merge strictly increases exact Surprise.
    Uses an inter-community edge map for O(1) ΔM/Δp per candidate merge.

    Parameters
    ----------
    edges : list of (str, str)
    partition : ignored (kept for API symmetry)

    Returns
    -------
    dict[str, int]
    """
    from pysurprise._core import fagso_surprise as _fagso_cpp

    all_nodes: List[str] = sorted({n for e in edges for n in e})
    if not all_nodes:
        return {}
    node_to_idx: Dict[str, int] = {n: i for i, n in enumerate(all_nodes)}
    int_edges: List[Tuple[int, int]] = [
        (node_to_idx[u], node_to_idx[v]) for u, v in edges
    ]
    raw: List[int] = _fagso_cpp(int_edges, len(all_nodes))
    return {all_nodes[i]: raw[i] for i in range(len(all_nodes))}


# ---------------------------------------------------------------------------
# SP4  (Tang et al. 2019)
# ---------------------------------------------------------------------------

def surprise_sp4(
    edges: Sequence[Edge],
    partition: Optional[Partition] = None,
    *,
    n_restarts: int = 8,
    seed: int = 42,
) -> Partition:
    """Maximise Surprise via SP4 (Tang et al. 2019).

    Executes Pre_CS → Louvain with super-node aggregation → Post_CS.
    Multi-restart: *n_restarts* independent runs, best partition returned.

    Parameters
    ----------
    edges : list of (str, str)
    partition : ignored (kept for API symmetry)
    n_restarts : int, default 8
    seed : int, default 42

    Returns
    -------
    dict[str, int]
    """
    from pysurprise._core import sp4_surprise as _sp4_cpp

    all_nodes: List[str] = sorted({n for e in edges for n in e})
    if not all_nodes:
        return {}
    node_to_idx: Dict[str, int] = {n: i for i, n in enumerate(all_nodes)}
    int_edges: List[Tuple[int, int]] = [
        (node_to_idx[u], node_to_idx[v]) for u, v in edges
    ]
    raw: List[int] = _sp4_cpp(int_edges, len(all_nodes), n_restarts, seed)
    return {all_nodes[i]: raw[i] for i in range(len(all_nodes))}


# ---------------------------------------------------------------------------
# M² (Marin Squared) — hybrid asymptotic + exact Surprise maximizer
# ---------------------------------------------------------------------------

def surprise_m2(
    edges: Sequence[Edge],
    partition: Optional[Partition] = None,
    *,
    n_restarts: int = 4,
    seed: int = 42,
) -> Partition:
    """Maximise Surprise via M² (Marin Squared).

    A fast hybrid algorithm combining:
      - **FAGSO warm-start** — Jaccard-ranked merging for O(m·log m) initial
        partition (restart 0).
      - **Asymptotic Louvain** — node movements scored with KL-divergence
        approximation (O(1) per candidate).  Super-node aggregation compresses
        the graph after each convergence (no O(K²) merge pass).
      - **Exact refinement** — one final node pass with exact Surprise to
        correct any approximation errors.

    Multi-restart: *n_restarts* independent runs, best partition returned.

    Parameters
    ----------
    edges : list of (str, str)
    partition : ignored (kept for API symmetry)
    n_restarts : int, default 4
    seed : int, default 42

    Returns
    -------
    dict[str, int]
    """
    from pysurprise._core import m2_surprise as _m2_cpp

    all_nodes: List[str] = sorted({n for e in edges for n in e})
    if not all_nodes:
        return {}
    node_to_idx: Dict[str, int] = {n: i for i, n in enumerate(all_nodes)}
    int_edges: List[Tuple[int, int]] = [
        (node_to_idx[u], node_to_idx[v]) for u, v in edges
    ]
    raw: List[int] = _m2_cpp(int_edges, len(all_nodes), n_restarts, seed)
    return {all_nodes[i]: raw[i] for i in range(len(all_nodes))}
