"""
Benchmark: surprise_refine vs surprise_anneal vs surprise_louvain.

Graphs tested
-------------
1. Karate Club          (n=34,  classic, known ground truth)
2. Ring of Cliques      (n=varies, clear planted structure)
3. LFR small  mu=0.1   (n=100, easy)
4. LFR small  mu=0.3   (n=100, medium)
5. LFR medium mu=0.1   (n=250, easy)
6. LFR medium mu=0.3   (n=250, medium)
7. Random Erdos-Renyi   (n=100, no community structure — sanity check)

Metrics reported per graph × algorithm
---------------------------------------
- Surprise value (higher = better)
- NMI with ground truth (where available)
- Wall-clock time (seconds)
"""
from __future__ import annotations

import time
from typing import Dict, List, Optional, Tuple

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


# ── Helpers ───────────────────────────────────────────────────────────────────

def nmi(p1: Partition, p2: Partition) -> float:
    """Normalised Mutual Information between two partitions (same node set)."""
    import math
    nodes = list(p1.keys())
    n = len(nodes)
    if n == 0:
        return 0.0

    def entropy(labels: list) -> float:
        from collections import Counter
        c = Counter(labels)
        return -sum((v / n) * math.log(v / n) for v in c.values() if v > 0)

    l1 = [p1[nd] for nd in nodes]
    l2 = [p2[nd] for nd in nodes]

    from collections import Counter
    joint = Counter(zip(l1, l2))
    h1 = entropy(l1)
    h2 = entropy(l2)
    mi = sum((cnt / n) * math.log((cnt / n) / ((l1c / n) * (l2c / n)))
             for (a, b), cnt in joint.items()
             if (l1c := Counter(l1)[a]) and (l2c := Counter(l2)[b]))
    denom = (h1 + h2) / 2.0
    return mi / denom if denom > 1e-12 else 1.0


def run_algo(name: str, fn, edges: List[Edge], **kw) -> Tuple[Partition, float, float]:
    """Run an optimiser, return (partition, surprise, elapsed_seconds)."""
    import pysurprise as pys
    t0 = time.perf_counter()
    part = fn(edges, **kw)
    elapsed = time.perf_counter() - t0
    s = pys.surprise(edges, part) if part else 0.0
    return part, s, elapsed


def print_row(algo: str, s: float, elapsed: float, nmi_val: Optional[float]) -> None:
    nmi_str = f"{nmi_val:.4f}" if nmi_val is not None else "  n/a "
    print(f"  {algo:<20}  S={s:10.4f}  NMI={nmi_str}  t={elapsed:.3f}s")


# ── Graph builders ────────────────────────────────────────────────────────────

def karate_club() -> Tuple[List[Edge], Optional[Partition]]:
    """Zachary's Karate Club with two-faction ground truth."""
    raw_edges = [
        (1,2),(1,3),(1,4),(1,5),(1,6),(1,7),(1,8),(1,9),(1,11),(1,12),(1,13),
        (1,14),(1,18),(1,20),(1,22),(1,32),(2,3),(2,4),(2,8),(2,14),(2,18),
        (2,20),(2,22),(2,31),(3,4),(3,8),(3,9),(3,10),(3,14),(3,28),(3,29),
        (3,33),(4,8),(4,13),(4,14),(5,7),(5,11),(6,7),(6,11),(6,17),(7,17),
        (9,31),(9,33),(9,34),(10,34),(14,34),(15,33),(15,34),(16,33),(16,34),
        (19,33),(19,34),(20,34),(21,33),(21,34),(23,33),(23,34),(24,26),(24,28),
        (24,30),(24,33),(24,34),(25,26),(25,28),(25,32),(26,32),(27,30),(27,34),
        (28,34),(29,32),(29,34),(30,33),(30,34),(31,33),(31,34),(32,33),(32,34),
        (33,34),
    ]
    edges = [(str(u), str(v)) for u, v in raw_edges]
    # Two factions: 1-based node IDs
    faction1 = {1,2,3,4,5,6,7,8,10,11,12,13,17,19,21}
    gt: Partition = {}
    for u, v in raw_edges:
        gt[str(u)] = 0 if u in faction1 else 1
        gt[str(v)] = 0 if v in faction1 else 1
    return edges, gt


def lfr_graph(n: int, mu: float, seed: int = 0) -> Tuple[List[Edge], Partition]:
    try:
        import networkx as nx
        G = nx.generators.community.LFR_benchmark_graph(
            n=n, tau1=2.5, tau2=1.5, mu=mu,
            min_degree=5, max_degree=20,
            min_community=10, max_community=50,
            seed=seed,
        )
        edges = [(str(u), str(v)) for u, v in G.edges()]
        gt: Partition = {}
        for node in G.nodes():
            comm_set = G.nodes[node]["community"]
            gt[str(node)] = min(comm_set)
        # Remap gt to 0-based ints
        unique = sorted(set(gt.values()))
        remap = {old: new for new, old in enumerate(unique)}
        gt = {k: remap[v] for k, v in gt.items()}
        return edges, gt
    except Exception as exc:
        print(f"  [LFR n={n} mu={mu}] failed: {exc}")
        return [], {}


def erdos_renyi(n: int = 100, p: float = 0.05, seed: int = 0) -> List[Edge]:
    import random
    rng = random.Random(seed)
    nodes = [str(i) for i in range(n)]
    edges = [(nodes[i], nodes[j])
             for i in range(n) for j in range(i + 1, n)
             if rng.random() < p]
    return edges


# ── Main benchmark ────────────────────────────────────────────────────────────

def benchmark():
    from pysurprise.algorithms import (
        surprise_refine, surprise_anneal,
        surprise_louvain, surprise_fagso, surprise_sp4,
    )

    graphs = []

    # 1. Karate club
    e, gt = karate_club()
    graphs.append(("Karate Club (n=34)", e, gt))

    # 2. Aldecoa RC (512 nodes, 16 communities, broken-stick sizes, J=0.75)
    from pysurprise.benchmarks.rc_benchmark import rc_benchmark, pielou_index, PAPER_PRESETS
    e, gt = rc_benchmark(**PAPER_PRESETS["2013_scirep"], degradation=0.0)
    _sz = sorted([sum(1 for v in gt.values() if v == c)
                  for c in set(gt.values())], reverse=True)
    _pi = pielou_index(_sz)
    graphs.append((f"Aldecoa RC (n=512, K=16, J={_pi:.2f})", e, gt))

    # 3-6. LFR benchmarks
    for n, mu in [(100, 0.1), (100, 0.3), (250, 0.1), (250, 0.3)]:
        e, gt = lfr_graph(n, mu, seed=42)
        if e:
            graphs.append((f"LFR n={n} mu={mu}", e, gt))

    # 7. Erdos-Renyi (no structure)
    e = erdos_renyi(n=80, p=0.06, seed=7)
    graphs.append(("Erdos-Renyi n=80 p=0.06", e, None))

    algos = [
        ("refine",  lambda edges: surprise_refine(edges, {n: i for i, n in
            enumerate(sorted({x for e in edges for x in e}))})),
        ("anneal",  lambda edges: surprise_anneal(edges, steps=50_000, seed=42)),
        ("fagso",   lambda edges: surprise_fagso(edges)),
        ("louvain", lambda edges: surprise_louvain(edges, n_restarts=8, seed=42)),
        ("sp4",     lambda edges: surprise_sp4(edges, n_restarts=8, seed=42)),
    ]

    for graph_name, edges, ground_truth in graphs:
        if not edges:
            continue
        print(f"\n{'─'*60}")
        print(f"Graph: {graph_name}  ({len(edges)} edges)")
        print(f"{'─'*60}")

        best_s = {}
        for algo_name, algo_fn in algos:
            part, s, t = run_algo(algo_name, algo_fn, edges)
            nmi_v = nmi(part, ground_truth) if ground_truth and part else None
            print_row(algo_name, s, t, nmi_v)
            best_s[algo_name] = s

        winner = max(best_s, key=best_s.get)
        print(f"  → Best Surprise: {winner}  ({best_s[winner]:.4f})")

    print(f"\n{'═'*60}")
    print("Done.")


if __name__ == "__main__":
    benchmark()
