"""Shootout: closeness & harmonic centrality — igraph vs networkit vs approx vs multiprocessing.

Usage:
    source .venv/bin/activate
    python3 benchmarks/shootout_closeness.py
"""

from __future__ import annotations

import multiprocessing as mp
import time
from functools import partial

import numpy as np


def _time(fn, warmup=1, iterations=5):
    """Run fn, return (p50_ms, result)."""
    for _ in range(warmup):
        fn()
    times = []
    for _ in range(iterations):
        t0 = time.perf_counter()
        result = fn()
        times.append((time.perf_counter() - t0) * 1000)
    return float(np.percentile(times, 50)), result


def _load_graph():
    """Load the benchmark graph from AGE into igraph + networkit."""
    from benchmarks.config import AGE_GRAPH
    from benchmarks.adapters.age_adapter import AgeAdapter
    from benchmarks.datagen.synthetic import generate
    from benchmarks.config import SCALES

    # Generate + load 10k data
    scale = SCALES["10k"]
    data = generate(scale)
    adapter = AgeAdapter()
    adapter.connect()
    adapter.setup_schema(scale.vec_dims)
    adapter.load_nodes(data.nodes)
    adapter.load_edges(data.edges)
    driver = adapter.db

    # igraph
    from piggie._integrations.igraph_integration import to_igraph_sql
    G_ig, id_map = to_igraph_sql(driver, AGE_GRAPH, directed=True)
    G_ig_u = G_ig.as_undirected()

    # networkit — build from igraph edge list
    import networkit as nk
    G_nk = nk.Graph(G_ig_u.vcount(), weighted=False, directed=False)
    for e in G_ig_u.get_edgelist():
        if not G_nk.hasEdge(e[0], e[1]):
            G_nk.addEdge(e[0], e[1])

    driver.close()
    return G_ig_u, G_nk, id_map


# ---------------------------------------------------------------------------
# Approach 1: igraph (single-threaded C)
# ---------------------------------------------------------------------------

def igraph_closeness(G):
    return G.closeness(mode="all")


def igraph_harmonic(G):
    return G.harmonic_centrality(mode="all")


# ---------------------------------------------------------------------------
# Approach 2: networkit (multi-threaded C++)
# ---------------------------------------------------------------------------

def networkit_closeness(G_nk):
    import networkit as nk
    algo = nk.centrality.Closeness(G_nk, True, nk.centrality.ClosenessVariant.Generalized)
    algo.run()
    return algo.scores()


def networkit_harmonic(G_nk):
    import networkit as nk
    algo = nk.centrality.HarmonicCloseness(G_nk)
    algo.run()
    return algo.scores()


# ---------------------------------------------------------------------------
# Approach 3: Approximate sampling (igraph BFS from k sources)
# ---------------------------------------------------------------------------

def approx_closeness(G, k=500, seed=42):
    """Sample k source vertices, compute shortest paths, estimate closeness."""
    import random
    rng = random.Random(seed)
    n = G.vcount()
    sources = rng.sample(range(n), min(k, n))

    # For each source, get distances to all other vertices
    # closeness(v) ≈ (k-1) / sum(d(s,v) for s in sources if d(s,v) < inf)
    # This is an approximation using the sampled sources
    total_dist = [0.0] * n
    reachable = [0] * n
    for s in sources:
        dists = G.distances(s, mode="all")[0]
        for v, d in enumerate(dists):
            if d > 0 and d < float("inf"):
                total_dist[v] += d
                reachable[v] += 1

    scores = []
    for v in range(n):
        if reachable[v] > 0 and total_dist[v] > 0:
            scores.append(reachable[v] / total_dist[v])
        else:
            scores.append(0.0)
    return scores


def approx_harmonic(G, k=500, seed=42):
    """Sample k sources, estimate harmonic centrality."""
    import random
    rng = random.Random(seed)
    n = G.vcount()
    sources = rng.sample(range(n), min(k, n))
    scale = n / len(sources)

    inv_dist_sum = [0.0] * n
    for s in sources:
        dists = G.distances(s, mode="all")[0]
        for v, d in enumerate(dists):
            if d > 0 and d < float("inf"):
                inv_dist_sum[v] += 1.0 / d

    return [x * scale for x in inv_dist_sum]


# ---------------------------------------------------------------------------
# Approach 4: multiprocessing with igraph
# ---------------------------------------------------------------------------

def _closeness_chunk(edges, n, vertex_range, directed):
    """Worker: compute closeness for a chunk of vertices."""
    import igraph as ig
    G = ig.Graph(n=n, edges=edges, directed=directed)
    results = {}
    for v in vertex_range:
        dists = G.distances(v, mode="all")[0]
        reachable = sum(1 for d in dists if 0 < d < float("inf"))
        total = sum(d for d in dists if 0 < d < float("inf"))
        results[v] = reachable / total if total > 0 else 0.0
    return results


def mp_closeness(G, workers=None):
    """Compute closeness using multiprocessing."""
    if workers is None:
        workers = min(mp.cpu_count(), 8)
    n = G.vcount()
    edges = G.get_edgelist()
    chunks = [list(range(i, n, workers)) for i in range(workers)]

    with mp.Pool(workers) as pool:
        fn = partial(_closeness_chunk, edges, n, directed=False)
        chunk_results = pool.map(fn, chunks)

    scores = [0.0] * n
    for chunk in chunk_results:
        for v, s in chunk.items():
            scores[v] = s
    return scores


def _harmonic_chunk(edges, n, vertex_range, directed):
    """Worker: compute harmonic centrality for a chunk of vertices."""
    import igraph as ig
    G = ig.Graph(n=n, edges=edges, directed=directed)
    results = {}
    for v in vertex_range:
        dists = G.distances(v, mode="all")[0]
        results[v] = sum(1.0 / d for d in dists if 0 < d < float("inf"))
    return results


def mp_harmonic(G, workers=None):
    """Compute harmonic centrality using multiprocessing."""
    if workers is None:
        workers = min(mp.cpu_count(), 8)
    n = G.vcount()
    edges = G.get_edgelist()
    chunks = [list(range(i, n, workers)) for i in range(workers)]

    with mp.Pool(workers) as pool:
        fn = partial(_harmonic_chunk, edges, n, directed=False)
        chunk_results = pool.map(fn, chunks)

    scores = [0.0] * n
    for chunk in chunk_results:
        for v, s in chunk.items():
            scores[v] = s
    return scores


# ---------------------------------------------------------------------------
# Also test greedy modularity alternatives
# ---------------------------------------------------------------------------

def igraph_fastgreedy(G):
    d = G.community_fastgreedy()
    return d.as_clustering()


def networkit_plm(G_nk):
    """Parallel Louvain Method — networkit's fast community detection."""
    import networkit as nk
    algo = nk.community.PLM(G_nk)
    algo.run()
    return algo.getPartition()


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    print("Loading graph from AGE...", flush=True)
    G_ig_u, G_nk, id_map = _load_graph()
    print(f"  igraph:    {G_ig_u.vcount():,} nodes, {G_ig_u.ecount():,} edges")
    print(f"  networkit: {G_nk.numberOfNodes():,} nodes, {G_nk.numberOfEdges():,} edges")
    print()

    gds_closeness = 1810  # ms, from R1 scorecard
    gds_harmonic = 298    # ms
    gds_greedy_mod = 147  # ms

    # --- Closeness ---
    print("=== CLOSENESS CENTRALITY ===")
    print()

    t, _ = _time(lambda: igraph_closeness(G_ig_u), warmup=1, iterations=3)
    print(f"  igraph (single-thread):     {t:>8.1f}ms  {'WIN' if t < gds_closeness else 'LOSE'} vs GDS {gds_closeness}ms")

    t, _ = _time(lambda: networkit_closeness(G_nk), warmup=1, iterations=3)
    print(f"  networkit (multi-thread):   {t:>8.1f}ms  {'WIN' if t < gds_closeness else 'LOSE'} vs GDS {gds_closeness}ms")

    for k in (200, 500, 1000):
        t, _ = _time(lambda k=k: approx_closeness(G_ig_u, k=k), warmup=1, iterations=3)
        print(f"  approx igraph k={k:<4d}:      {t:>8.1f}ms  {'WIN' if t < gds_closeness else 'LOSE'} vs GDS {gds_closeness}ms")

    t, _ = _time(lambda: mp_closeness(G_ig_u), warmup=0, iterations=2)
    print(f"  multiprocessing (8-core):   {t:>8.1f}ms  {'WIN' if t < gds_closeness else 'LOSE'} vs GDS {gds_closeness}ms")

    print()

    # --- Harmonic ---
    print("=== HARMONIC CENTRALITY ===")
    print()

    t, _ = _time(lambda: igraph_harmonic(G_ig_u), warmup=1, iterations=3)
    print(f"  igraph (single-thread):     {t:>8.1f}ms  {'WIN' if t < gds_harmonic else 'LOSE'} vs GDS {gds_harmonic}ms")

    t, _ = _time(lambda: networkit_harmonic(G_nk), warmup=1, iterations=3)
    print(f"  networkit (multi-thread):   {t:>8.1f}ms  {'WIN' if t < gds_harmonic else 'LOSE'} vs GDS {gds_harmonic}ms")

    for k in (200, 500, 1000):
        t, _ = _time(lambda k=k: approx_harmonic(G_ig_u, k=k), warmup=1, iterations=3)
        print(f"  approx igraph k={k:<4d}:      {t:>8.1f}ms  {'WIN' if t < gds_harmonic else 'LOSE'} vs GDS {gds_harmonic}ms")

    t, _ = _time(lambda: mp_harmonic(G_ig_u), warmup=0, iterations=2)
    print(f"  multiprocessing (8-core):   {t:>8.1f}ms  {'WIN' if t < gds_harmonic else 'LOSE'} vs GDS {gds_harmonic}ms")

    print()

    # --- Greedy Modularity ---
    print("=== GREEDY MODULARITY ===")
    print()

    t, _ = _time(lambda: igraph_fastgreedy(G_ig_u), warmup=1, iterations=3)
    print(f"  igraph fastgreedy:          {t:>8.1f}ms  {'WIN' if t < gds_greedy_mod else 'LOSE'} vs GDS {gds_greedy_mod}ms")

    t, _ = _time(lambda: networkit_plm(G_nk), warmup=1, iterations=3)
    print(f"  networkit PLM:              {t:>8.1f}ms  {'WIN' if t < gds_greedy_mod else 'LOSE'} vs GDS {gds_greedy_mod}ms")
