#!/usr/bin/env python3
"""Benchmark GPU vs CPU backends for graph algorithms.

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

Compares:
    - cuGraph GPU (direct API)
    - 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.backends.dispatcher import AlgorithmDispatcher
from chainswarm_analyzers_baseline.graph.builder import build_money_flow_graph


def load_and_build_graph(parquet_path: str) -> nx.DiGraph:
    """Load parquet file and build NetworkX graph."""
    df = pd.read_parquet(parquet_path)
    flows = df.to_dict('records')
    return build_money_flow_graph(flows)


def benchmark_louvain(G: nx.DiGraph, gpu: GPUBackend) -> dict:
    """Benchmark Louvain community detection."""
    results = {}

    # GPU
    times = []
    for _ in range(3):
        t0 = time.perf_counter()
        communities_gpu = gpu.louvain()
        times.append(time.perf_counter() - t0)
    results['gpu_ms'] = min(times) * 1000
    results['gpu_communities'] = len(communities_gpu)

    # CPU
    G_und = G.to_undirected()
    # 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(3):
        t0 = time.perf_counter()
        communities_cpu = list(nx.community.louvain_communities(G_und, weight='weight', seed=42))
        times.append(time.perf_counter() - t0)
    results['cpu_ms'] = min(times) * 1000
    results['cpu_communities'] = len(communities_cpu)

    results['speedup'] = results['cpu_ms'] / results['gpu_ms']
    return results


def benchmark_pagerank(G: nx.DiGraph, gpu: GPUBackend) -> dict:
    """Benchmark PageRank."""
    results = {}

    # GPU
    times = []
    for _ in range(5):
        t0 = time.perf_counter()
        _ = gpu.pagerank()
        times.append(time.perf_counter() - t0)
    results['gpu_ms'] = min(times) * 1000

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

    results['speedup'] = results['cpu_ms'] / results['gpu_ms']
    return results


def benchmark_scc(G: nx.DiGraph, gpu: GPUBackend) -> dict:
    """Benchmark strongly connected components."""
    results = {}

    # CPU (expected to be faster)
    times = []
    for _ in range(5):
        t0 = time.perf_counter()
        sccs_cpu = list(nx.strongly_connected_components(G))
        times.append(time.perf_counter() - t0)
    results['cpu_ms'] = min(times) * 1000
    results['cpu_sccs'] = len(sccs_cpu)

    # Note: We skip GPU SCC benchmark because it's known to be 32x slower
    results['note'] = "GPU SCC skipped (known to be 32x slower)"

    return results


def benchmark_triangles(G: nx.DiGraph, gpu: GPUBackend) -> dict:
    """Benchmark triangle count."""
    results = {}

    # GPU
    times = []
    for _ in range(5):
        t0 = time.perf_counter()
        tri_gpu = gpu.triangle_count()
        times.append(time.perf_counter() - t0)
    results['gpu_ms'] = min(times) * 1000
    results['gpu_triangles'] = tri_gpu

    # CPU
    G_und = G.to_undirected()
    times = []
    for _ in range(5):
        t0 = time.perf_counter()
        triangles = nx.triangles(G_und)
        tri_cpu = sum(triangles.values()) // 3
        times.append(time.perf_counter() - t0)
    results['cpu_ms'] = min(times) * 1000
    results['cpu_triangles'] = tri_cpu

    results['speedup'] = results['cpu_ms'] / results['gpu_ms']
    return results


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

    # Check GPU availability
    if not gpu_available():
        logger.error("cuGraph not available. Install with: pip install cugraph-cu12")
        sys.exit(1)

    # Load data and build graph
    logger.info(f"Loading data from {args.input}")
    G = load_and_build_graph(args.input)
    logger.info(f"Loaded {G.number_of_edges():,} flows")
    n_nodes = G.number_of_nodes()
    n_edges = G.number_of_edges()
    logger.info(f"Graph: {n_nodes:,} nodes, {n_edges:,} edges")

    # Initialize GPU backend
    logger.info("Loading graph to GPU...")
    gpu = GPUBackend()
    gpu.load_from_networkx(G)

    # Run benchmarks
    print("\n" + "=" * 70)
    print(f"BENCHMARK: {n_nodes:,} nodes, {n_edges:,} edges")
    print("=" * 70)

    # Louvain
    print("\n=== Louvain Community Detection ===")
    louvain = benchmark_louvain(G, gpu)
    print(f"GPU:    {louvain['gpu_ms']:>8.1f}ms ({louvain['gpu_communities']} communities)")
    print(f"CPU:    {louvain['cpu_ms']:>8.1f}ms ({louvain['cpu_communities']} communities)")
    print(f"Speedup: {louvain['speedup']:.1f}x")

    # PageRank
    print("\n=== PageRank ===")
    pagerank = benchmark_pagerank(G, gpu)
    print(f"GPU:    {pagerank['gpu_ms']:>8.1f}ms")
    print(f"CPU:    {pagerank['cpu_ms']:>8.1f}ms")
    print(f"Speedup: {pagerank['speedup']:.1f}x")

    # SCC
    print("\n=== Strongly Connected Components ===")
    scc = benchmark_scc(G, gpu)
    print(f"CPU:    {scc['cpu_ms']:>8.1f}ms ({scc['cpu_sccs']} SCCs)")
    print(f"Note:   {scc['note']}")

    # Triangles
    print("\n=== Triangle Count ===")
    triangles = benchmark_triangles(G, gpu)
    print(f"GPU:    {triangles['gpu_ms']:>8.1f}ms ({triangles['gpu_triangles']} triangles)")
    print(f"CPU:    {triangles['cpu_ms']:>8.1f}ms ({triangles['cpu_triangles']} triangles)")
    print(f"Speedup: {triangles['speedup']:.1f}x")

    # Summary
    print("\n" + "=" * 70)
    print("SUMMARY")
    print("=" * 70)
    print(f"{'Algorithm':<25} {'GPU (ms)':<12} {'CPU (ms)':<12} {'Speedup':<10}")
    print("-" * 70)
    print(f"{'Louvain':<25} {louvain['gpu_ms']:<12.1f} {louvain['cpu_ms']:<12.1f} {louvain['speedup']:.1f}x")
    print(f"{'PageRank':<25} {pagerank['gpu_ms']:<12.1f} {pagerank['cpu_ms']:<12.1f} {pagerank['speedup']:.1f}x")
    print(f"{'Triangles':<25} {triangles['gpu_ms']:<12.1f} {triangles['cpu_ms']:<12.1f} {triangles['speedup']:.1f}x")
    print(f"{'SCC':<25} {'N/A':<12} {scc['cpu_ms']:<12.1f} {'CPU only':<10}")
    print("=" * 70)


if __name__ == "__main__":
    main()
