#!/usr/bin/env python3
"""Benchmark Rust+petgraph vs cuGraph GPU vs NetworkX CPU.

Usage:
    python scripts/benchmark_all.py --input data/money_flows.parquet

Compares:
    - Rust (petgraph) CPU
    - cuGraph GPU
    - NetworkX CPU
"""

import argparse
import time
import sys
from pathlib import Path

import networkx as nx
import pandas as pd
from loguru import logger

# Add src to path
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))

from chainswarm_analyzers_baseline.backends.gpu import GPUBackend, gpu_available
from chainswarm_analyzers_baseline.graph.builder import build_money_flow_graph


def load_graph(parquet_path: str):
    """Load parquet and build graphs for all backends."""
    df = pd.read_parquet(parquet_path)
    flows = df.to_dict('records')

    # NetworkX graph
    G = build_money_flow_graph(flows)

    # Extract edge data for Rust
    sources = df['from_address'].tolist()
    targets = df['to_address'].tolist()
    weights = df['amount_usd_sum'].tolist()

    return G, sources, targets, weights


def benchmark_rust(sources, targets, weights, n_runs=3):
    """Benchmark Rust+petgraph."""
    try:
        from chainswarm_analyzers_baseline._rust import RustGraph
    except ImportError:
        logger.warning("Rust backend not available. Build with: maturin develop --release")
        return None

    results = {}

    # Create graph
    t0 = time.perf_counter()
    g = RustGraph()
    g.add_edges(sources, targets, weights)
    results['load_ms'] = (time.perf_counter() - t0) * 1000

    n_nodes, n_edges = g.stats()
    logger.info(f"Rust graph: {n_nodes:,} nodes, {n_edges:,} edges")

    # SCC
    times = []
    for _ in range(n_runs):
        t0 = time.perf_counter()
        sccs = g.strongly_connected_components()
        times.append(time.perf_counter() - t0)
    results['scc_ms'] = min(times) * 1000
    results['scc_count'] = len(sccs)

    # PageRank
    times = []
    for _ in range(n_runs):
        t0 = time.perf_counter()
        pr = g.pagerank(0.85, 100, 1e-6)
        times.append(time.perf_counter() - t0)
    results['pagerank_ms'] = min(times) * 1000

    # Triangle Count
    times = []
    for _ in range(n_runs):
        t0 = time.perf_counter()
        tri = g.triangle_count()
        times.append(time.perf_counter() - t0)
    results['triangles_ms'] = min(times) * 1000
    results['triangles_count'] = tri

    # Louvain
    times = []
    for _ in range(n_runs):
        t0 = time.perf_counter()
        communities = g.louvain(1.0)
        times.append(time.perf_counter() - t0)
    results['louvain_ms'] = min(times) * 1000
    results['louvain_communities'] = len(communities)

    return results


def benchmark_gpu(G: nx.DiGraph, n_runs=3):
    """Benchmark cuGraph GPU."""
    if not gpu_available():
        logger.warning("GPU not available")
        return None

    results = {}

    # Load to GPU
    t0 = time.perf_counter()
    gpu = GPUBackend()
    gpu.load_from_networkx(G)
    results['load_ms'] = (time.perf_counter() - t0) * 1000

    # Louvain
    times = []
    for _ in range(n_runs):
        t0 = time.perf_counter()
        communities = gpu.louvain()
        times.append(time.perf_counter() - t0)
    results['louvain_ms'] = min(times) * 1000
    results['louvain_communities'] = len(communities)

    # PageRank
    times = []
    for _ in range(n_runs):
        t0 = time.perf_counter()
        pr = gpu.pagerank()
        times.append(time.perf_counter() - t0)
    results['pagerank_ms'] = min(times) * 1000

    # Triangle Count
    times = []
    for _ in range(n_runs):
        t0 = time.perf_counter()
        tri = gpu.triangle_count()
        times.append(time.perf_counter() - t0)
    results['triangles_ms'] = min(times) * 1000
    results['triangles_count'] = tri

    # Note: SCC skipped for GPU (known to be 32x slower)
    results['scc_ms'] = None
    results['scc_note'] = "Skipped (GPU SCC is 32x slower)"

    return results


def benchmark_cpu(G: nx.DiGraph, n_runs=3):
    """Benchmark NetworkX CPU."""
    results = {}

    # SCC
    times = []
    for _ in range(n_runs):
        t0 = time.perf_counter()
        sccs = list(nx.strongly_connected_components(G))
        times.append(time.perf_counter() - t0)
    results['scc_ms'] = min(times) * 1000
    results['scc_count'] = len(sccs)

    # PageRank
    times = []
    for _ in range(n_runs):
        t0 = time.perf_counter()
        pr = nx.pagerank(G, weight='weight')
        times.append(time.perf_counter() - t0)
    results['pagerank_ms'] = min(times) * 1000

    # Triangle Count
    G_und = G.to_undirected()
    times = []
    for _ in range(n_runs):
        t0 = time.perf_counter()
        triangles = nx.triangles(G_und)
        tri = sum(triangles.values()) // 3
        times.append(time.perf_counter() - t0)
    results['triangles_ms'] = min(times) * 1000
    results['triangles_count'] = tri

    # Louvain
    # Ensure weights
    for u, v, data in G_und.edges(data=True):
        if data.get('weight', 0) == 0:
            data['weight'] = max(data.get('tx_count', 1), 1)

    times = []
    for _ in range(n_runs):
        t0 = time.perf_counter()
        communities = list(nx.community.louvain_communities(G_und, weight='weight', seed=42))
        times.append(time.perf_counter() - t0)
    results['louvain_ms'] = min(times) * 1000
    results['louvain_communities'] = len(communities)

    return results


def print_results(rust, gpu, cpu, n_nodes, n_edges):
    """Print comparison table."""
    print("\n" + "=" * 90)
    print(f"BENCHMARK RESULTS: {n_nodes:,} nodes, {n_edges:,} edges")
    print("=" * 90)

    algorithms = ['louvain', 'pagerank', 'triangles', 'scc']

    print(f"\n{'Algorithm':<15} {'Rust (ms)':<15} {'GPU (ms)':<15} {'CPU (ms)':<15} {'Best':<15}")
    print("-" * 90)

    for algo in algorithms:
        key = f'{algo}_ms'

        rust_ms = rust.get(key) if rust else None
        gpu_ms = gpu.get(key) if gpu else None
        cpu_ms = cpu.get(key) if cpu else None

        rust_str = f"{rust_ms:.1f}" if rust_ms else "N/A"
        gpu_str = f"{gpu_ms:.1f}" if gpu_ms else "N/A"
        cpu_str = f"{cpu_ms:.1f}" if cpu_ms else "N/A"

        # Determine best
        times = []
        if rust_ms: times.append(('Rust', rust_ms))
        if gpu_ms: times.append(('GPU', gpu_ms))
        if cpu_ms: times.append(('CPU', cpu_ms))

        if times:
            best_name, best_time = min(times, key=lambda x: x[1])

            # Calculate speedup vs slowest
            if len(times) > 1:
                slowest = max(t for _, t in times)
                speedup = slowest / best_time
                best_str = f"{best_name} ({speedup:.1f}x)"
            else:
                best_str = best_name
        else:
            best_str = "N/A"

        print(f"{algo.upper():<15} {rust_str:<15} {gpu_str:<15} {cpu_str:<15} {best_str:<15}")

    # Print result counts for verification
    print("\n" + "-" * 90)
    print("VERIFICATION (counts should match across backends):")
    print("-" * 90)

    if rust:
        print(f"Rust:  Louvain={rust.get('louvain_communities', 'N/A')} communities, "
              f"Triangles={rust.get('triangles_count', 'N/A')}, "
              f"SCCs={rust.get('scc_count', 'N/A')}")
    if gpu:
        print(f"GPU:   Louvain={gpu.get('louvain_communities', 'N/A')} communities, "
              f"Triangles={gpu.get('triangles_count', 'N/A')}, "
              f"SCCs={gpu.get('scc_note', 'N/A')}")
    if cpu:
        print(f"CPU:   Louvain={cpu.get('louvain_communities', 'N/A')} communities, "
              f"Triangles={cpu.get('triangles_count', 'N/A')}, "
              f"SCCs={cpu.get('scc_count', 'N/A')}")

    print("=" * 90)


def main():
    parser = argparse.ArgumentParser(description="Benchmark Rust vs GPU vs CPU graph algorithms")
    parser.add_argument("--input", required=True, help="Path to money_flows.parquet")
    args = parser.parse_args()

    # Load data
    logger.info(f"Loading data from {args.input}")
    G, sources, targets, weights = load_graph(args.input)
    n_nodes = G.number_of_nodes()
    n_edges = G.number_of_edges()
    logger.info(f"Graph: {n_nodes:,} nodes, {n_edges:,} edges")

    # Run benchmarks
    print("\nRunning Rust benchmark...")
    rust_results = benchmark_rust(sources, targets, weights)

    print("Running GPU benchmark...")
    gpu_results = benchmark_gpu(G)

    print("Running CPU benchmark...")
    cpu_results = benchmark_cpu(G)

    # Print comparison
    print_results(rust_results, gpu_results, cpu_results, n_nodes, n_edges)


if __name__ == "__main__":
    main()
