"""
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
# ---------------------------------------------------------------------------

def _find_repo_root() -> 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
        # Also check two levels up from pysurprise/src/pysurprise/
        if (parent / "pysurprise" / "src" / "pysurprise").is_dir():
            return parent
    # Fallback: try common layout
    candidate = p.parent.parent.parent.parent  # pysurprise/src/pysurprise -> repo
    if (candidate / "src" / "Surprise").is_dir():
        return candidate
    raise FileNotFoundError(
        "Cannot locate SurpriseMe binaries. Make sure the repo root "
        "contains src/Surprise/, src/Jerarca/, etc."
    )


_REPO: Optional[Path] = None


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


def _bin(relpath: str) -> str:
    """Return absolute path to a compiled binary, raising if missing."""
    p = _repo() / relpath
    if not p.exists():
        raise FileNotFoundError(
            f"Binary not found: {p}\n"
            f"Run 'make' in {_repo()} to compile."
        )
    return str(p)


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 _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("src/CPM/bin/slicer")
    community_bin = _bin("src/CPM/bin/community")
    compute_s = _bin("src/Surprise/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

        while step > 0.001:
            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:
                            break
                gamma += step
            step /= 2
            count = 0

        return best_part


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

    Uses the Python ``infomap`` package (pip install infomap).
    """
    try:
        import infomap as im
    except ImportError:
        raise ImportError(
            "The 'infomap' Python package is required. "
            "Install it with: pip install infomap"
        )

    node2idx, idx2node = _node_index_maps(edges)
    network = im.Infomap(silent=True, num_trials=attempts)
    for u, v in edges:
        network.add_link(node2idx[u], node2idx[v])
    network.run()

    part: Partition = {}
    for node_id in network.tree:
        if node_id.is_leaf:
            part[idx2node[node_id.node_id]] = node_id.module_id
    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("src/RB/bin/slicer")
    community_bin = _bin("src/RB/bin/community")
    compute_s = _bin("src/Surprise/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

        while step > 0.01:
            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:
                            break
                gamma += step
            step /= 2
            count = 0

        return best_part


def rn(edges: Sequence[Edge], *, timeout: int = 120) -> Partition:
    """RN — Ronhovde-Nussinov (2010).

    Resolution-limit-free Potts model.  Uses igraph's Leiden algorithm with
    the CPM objective function (same Potts-model family) and scans the
    resolution parameter to maximise Surprise.

    Requires the ``igraph`` Python package (``pip install igraph``).
    """
    try:
        import igraph as ig
    except ImportError:
        raise ImportError(
            "The 'igraph' Python package is required for RN. "
            "Install it with: pip install igraph"
        )

    from pysurprise._core import (
        compute_surprise as _compute_surprise,
        surprise_from_partition_dict as _sfpd,
    )

    node2idx, idx2node = _node_index_maps(edges)
    n_nodes = len(node2idx)

    g = ig.Graph(n=n_nodes, edges=[(node2idx[u], node2idx[v]) for u, v in edges],
                 directed=False)

    best_part: Partition = {}
    best_s = -1.0

    # Scan resolution parameter (gamma) — coarse then fine
    for gamma in [0.001, 0.005, 0.01, 0.02, 0.05, 0.1, 0.15, 0.2, 0.3,
                  0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0, 1.2, 1.5, 2.0,
                  3.0, 5.0, 8.0, 10.0, 15.0, 20.0]:
        try:
            membership = g.community_leiden(
                objective_function="CPM",
                resolution=gamma,
                n_iterations=-1,
            )
            str_edges = [(str(u), str(v)) for u, v in edges]
            part_dict = {idx2node[i]: int(c)
                         for i, c in enumerate(membership.membership)}
            s = _sfpd(str_edges, {str(k): v for k, v in part_dict.items()})
            if s > best_s:
                best_s = s
                best_part = part_dict
        except Exception:
            continue

    return best_part


def rnsc(edges: Sequence[Edge], *, timeout: int = 120) -> Partition:
    """RNSC — Restricted Neighbourhood Search Clustering (King et al. 2004)."""
    rnscconvert = _bin("src/RNSC/rnscconvert")
    rnsc_bin = _bin("src/RNSC/rnsc")
    rnscfilter = _bin("src/RNSC/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("src/Jerarca/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("src/Jerarca/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


# ---------------------------------------------------------------------------
# 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
