torsh-graph 0.1.3

//! Performance benchmarks for torsh-graph vs PyTorch Geometric
//!
//! This module provides comprehensive performance testing and benchmarking
//! for graph neural network operations, comparing scalability and throughput.

use std::time::Instant;
use torsh_graph::{
    conv::{AggregationType, GCNConv, GINConv, GraphTransformer, MPNNConv, SAGEConv},
    scirs2_integration::generation,
    GraphData, GraphLayer,
};
use torsh_tensor::creation::randn;

/// Benchmark configuration for systematic testing
#[derive(Debug, Clone)]
struct BenchmarkConfig {
    name: String,
    num_nodes: usize,
    num_features: usize,
    out_features: usize,
    edge_probability: f64,
    num_iterations: usize,
}

impl BenchmarkConfig {
    fn small() -> Self {
        Self {
            name: "Small".to_string(),
            num_nodes: 100,
            num_features: 16,
            out_features: 32,
            edge_probability: 0.1,
            num_iterations: 100,
        }
    }

    fn medium() -> Self {
        Self {
            name: "Medium".to_string(),
            num_nodes: 500,
            num_features: 64,
            out_features: 128,
            edge_probability: 0.05,
            num_iterations: 50,
        }
    }

    fn large() -> Self {
        Self {
            name: "Large".to_string(),
            num_nodes: 1000,
            num_features: 128,
            out_features: 256,
            edge_probability: 0.03,
            num_iterations: 20,
        }
    }

    #[allow(dead_code)]
    fn xlarge() -> Self {
        Self {
            name: "XLarge".to_string(),
            num_nodes: 2000,
            num_features: 256,
            out_features: 512,
            edge_probability: 0.02,
            num_iterations: 10,
        }
    }
}

/// Generate benchmark graph according to configuration
fn create_benchmark_graph(config: &BenchmarkConfig) -> GraphData {
    if config.num_features == 16 {
        generation::erdos_renyi(config.num_nodes, config.edge_probability)
    } else {
        // For larger feature sizes, create custom graph
        let x = randn(&[config.num_nodes, config.num_features]).unwrap();
        let base_graph = generation::erdos_renyi(config.num_nodes, config.edge_probability);

        GraphData {
            x,
            edge_index: base_graph.edge_index,
            edge_attr: base_graph.edge_attr,
            batch: base_graph.batch,
            num_nodes: base_graph.num_nodes,
            num_edges: base_graph.num_edges,
        }
    }
}

/// Benchmark a graph layer with given configuration
fn benchmark_layer<L: GraphLayer>(
    layer: &L,
    config: &BenchmarkConfig,
    layer_name: &str,
) -> (f64, f64, usize) {
    let graph = create_benchmark_graph(config);

    // Warmup runs
    for _ in 0..5 {
        let _ = layer.forward(&graph);
    }

    // Actual benchmark
    let start = Instant::now();
    for _ in 0..config.num_iterations {
        let _output = layer.forward(&graph);
    }
    let duration = start.elapsed();

    let total_ms = duration.as_millis() as f64;
    let avg_ms = total_ms / config.num_iterations as f64;
    let throughput = config.num_nodes as f64 / avg_ms * 1000.0; // nodes/second

    println!(
        "📊 {} | {} | Avg: {:.2}ms | Total: {:.2}ms | Throughput: {:.0} nodes/sec",
        config.name, layer_name, avg_ms, total_ms, throughput
    );

    (avg_ms, throughput, graph.num_edges)
}

/// Comprehensive performance test comparing all layers
///
/// This test is ignored by default due to long execution time (>300s)
#[test]
#[ignore = "timeout"]
fn test_comprehensive_layer_performance() {
    println!("\n🚀 ToRSh Graph Neural Networks - Performance Benchmarks");
    println!("{}", "=".repeat(80));

    let configs = vec![
        BenchmarkConfig::small(),
        BenchmarkConfig::medium(),
        BenchmarkConfig::large(),
    ];

    for config in configs {
        println!(
            "\n📈 Configuration: {} ({} nodes, {} features)",
            config.name, config.num_nodes, config.num_features
        );
        println!("{}", "-".repeat(60));

        // Test GCN
        let gcn = GCNConv::new(config.num_features, config.out_features, true);
        let (gcn_time, _gcn_throughput, num_edges) = benchmark_layer(&gcn, &config, "GCN");

        // Test SAGE
        let sage = SAGEConv::new(config.num_features, config.out_features, true);
        let (sage_time, _sage_throughput, _) = benchmark_layer(&sage, &config, "SAGE");

        // Test GIN
        let gin = GINConv::new(config.num_features, config.out_features, 0.0, false, true);
        let (gin_time, _gin_throughput, _) = benchmark_layer(&gin, &config, "GIN");

        // Test MPNN
        let mpnn = MPNNConv::new(
            config.num_features,
            config.out_features,
            0,
            config.out_features * 2,
            config.out_features * 2,
            AggregationType::Mean,
            true,
        );
        let (mpnn_time, _mpnn_throughput, _) = benchmark_layer(&mpnn, &config, "MPNN");

        // Test GraphTransformer (with smaller head count for performance)
        let heads = if config.num_features >= 64 { 8 } else { 4 };
        let transformer = GraphTransformer::new(
            config.num_features,
            config.out_features,
            config.num_features / heads, // head_dim
            heads,
            0.0,
            true,
        );
        let (transformer_time, _transformer_throughput, _) =
            benchmark_layer(&transformer, &config, "Transformer");

        // Summary statistics
        println!("\n📋 Summary for {}:", config.name);
        println!("   Graph: {} nodes, {} edges", config.num_nodes, num_edges);

        let times = vec![
            ("GCN", gcn_time),
            ("SAGE", sage_time),
            ("GIN", gin_time),
            ("MPNN", mpnn_time),
            ("Transformer", transformer_time),
        ];

        let fastest = times
            .iter()
            .min_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
            .unwrap();
        let slowest = times
            .iter()
            .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
            .unwrap();

        println!("   ⚡ Fastest: {} ({:.2}ms)", fastest.0, fastest.1);
        println!("   🐌 Slowest: {} ({:.2}ms)", slowest.0, slowest.1);
        println!("   📊 Speedup: {:.1}x", slowest.1 / fastest.1);
    }
}

/// Memory efficiency benchmark
#[test]
fn test_memory_scalability() {
    println!("\n🧠 Memory Scalability Benchmark");
    println!("{}", "=".repeat(50));

    let node_counts = vec![100, 500, 1000, 2000];
    let base_config = BenchmarkConfig::small();

    for num_nodes in node_counts {
        let config = BenchmarkConfig {
            num_nodes,
            num_iterations: 10,
            ..base_config.clone()
        };

        println!("\n📊 Testing {} nodes", num_nodes);

        // Test different layers for memory efficiency
        let gcn = GCNConv::new(16, 32, true);
        let graph = create_benchmark_graph(&config);

        let start = Instant::now();
        for _ in 0..config.num_iterations {
            let _output = gcn.forward(&graph);
        }
        let duration = start.elapsed();

        let avg_time = duration.as_millis() as f64 / config.num_iterations as f64;
        let time_per_node = avg_time / num_nodes as f64;

        println!("   ⏱️  Avg time: {:.2}ms", avg_time);
        println!("   🔧 Time/node: {:.4}ms", time_per_node);
        println!("   📈 Edges: {}", graph.num_edges);

        // Validate linear scaling expectation (adjusted for realistic performance in test builds)
        // Relaxed threshold to account for debug builds, varied hardware, and parallel test execution
        if num_nodes >= 1000 {
            assert!(
                time_per_node < 5.0,
                "Time per node should remain reasonable: {:.4}ms",
                time_per_node
            );
        }
    }
}

/// Edge density impact benchmark
#[test]
fn test_edge_density_performance() {
    println!("\n🌐 Edge Density Performance Impact");
    println!("{}", "=".repeat(50));

    let densities = vec![0.01, 0.05, 0.1, 0.2];
    let base_nodes = 500;

    for density in densities {
        println!("\n📊 Edge density: {:.1}%", density * 100.0);

        let graph = generation::erdos_renyi(base_nodes, density);
        let gcn = GCNConv::new(16, 32, true);

        let start = Instant::now();
        for _ in 0..20 {
            let _output = gcn.forward(&graph);
        }
        let duration = start.elapsed();

        let avg_time = duration.as_millis() as f64 / 20.0;
        let time_per_edge = avg_time / graph.num_edges as f64;

        println!("   📈 Edges: {}", graph.num_edges);
        println!("   ⏱️  Avg time: {:.2}ms", avg_time);
        println!("   🔗 Time/edge: {:.4}ms", time_per_edge);

        // Basic performance assertions
        assert!(
            avg_time < 1000.0,
            "Performance degradation too severe at {:.1}% density",
            density * 100.0
        );
    }
}

/// Deep network performance test
#[test]
fn test_deep_network_performance() {
    println!("\n🏗️  Deep Network Performance");
    println!("{}", "=".repeat(40));

    let graph = create_benchmark_graph(&BenchmarkConfig::medium());
    let layer_counts = vec![1, 2, 4, 8];

    for num_layers in layer_counts {
        println!("\n📊 Testing {} layer(s)", num_layers);

        // Create layers
        let mut layers: Vec<Box<dyn GraphLayer>> = Vec::new();

        // First layer
        layers.push(Box::new(GCNConv::new(64, 32, true)));

        // Hidden layers
        for _ in 1..num_layers {
            layers.push(Box::new(GCNConv::new(32, 32, true)));
        }

        // Benchmark deep forward pass
        let start = Instant::now();
        for _ in 0..10 {
            let mut current_graph = graph.clone();

            for layer in &layers {
                current_graph = layer.forward(&current_graph);
            }

            // Validate final output
            assert_eq!(current_graph.num_nodes, graph.num_nodes);
            assert_eq!(current_graph.x.shape().dims()[1], 32);
        }
        let duration = start.elapsed();

        let avg_time = duration.as_millis() as f64 / 10.0;
        let time_per_layer = avg_time / num_layers as f64;

        println!("   ⏱️  Total time: {:.2}ms", avg_time);
        println!("   🏗️  Time/layer: {:.2}ms", time_per_layer);

        // Validate reasonable scaling (relaxed for test builds and varied hardware)
        if num_layers >= 4 {
            assert!(
                time_per_layer < 200.0,
                "Per-layer time should be reasonable: {:.2}ms",
                time_per_layer
            );
        }
    }
}

/// Batch processing performance simulation
#[test]
fn test_batch_processing_simulation() {
    println!("\n📦 Batch Processing Simulation");
    println!("{}", "=".repeat(40));

    let batch_sizes = vec![1, 4, 8, 16];
    let gcn = GCNConv::new(16, 32, true);

    for batch_size in batch_sizes {
        println!("\n📊 Batch size: {}", batch_size);

        // Create multiple graphs
        let graphs: Vec<GraphData> = (0..batch_size)
            .map(|_| generation::erdos_renyi(100, 0.1))
            .collect();

        // Time sequential processing (current approach)
        let start = Instant::now();
        for _ in 0..10 {
            for graph in &graphs {
                let _output = gcn.forward(graph);
            }
        }
        let duration = start.elapsed();

        let total_time = duration.as_millis() as f64 / 10.0;
        let time_per_graph = total_time / batch_size as f64;

        println!("   ⏱️  Total time: {:.2}ms", total_time);
        println!("   📊 Time/graph: {:.2}ms", time_per_graph);
        println!(
            "   🚀 Throughput: {:.0} graphs/sec",
            1000.0 / time_per_graph
        );

        // Basic efficiency check
        assert!(
            time_per_graph < 50.0,
            "Per-graph time should be reasonable: {:.2}ms",
            time_per_graph
        );
    }
}

/// Comparative analysis with theoretical PyTorch Geometric performance
#[test]
fn test_pytorch_geometric_comparison_analysis() {
    println!("\n🆚 Theoretical PyTorch Geometric Comparison");
    println!("{}", "=".repeat(60));
    println!("Note: This test provides baseline performance analysis for future comparison");

    let config = BenchmarkConfig::medium();
    let _graph = create_benchmark_graph(&config);

    // Benchmark our implementations
    let gcn = GCNConv::new(config.num_features, config.out_features, true);
    let (torsh_time, torsh_throughput, _) = benchmark_layer(&gcn, &config, "ToRSh-GCN");

    println!("\n📊 Performance Analysis:");
    println!("   🦀 ToRSh Performance:");
    println!("      - Average time: {:.2}ms", torsh_time);
    println!("      - Throughput: {:.0} nodes/sec", torsh_throughput);
    println!("      - Memory: Rust zero-cost abstractions");
    println!("      - Backend: Pure Rust with SciRS2 optimization");

    println!("\n   🐍 Expected PyTorch Geometric Performance:");
    println!(
        "      - Estimated time: ~{:.2}ms (GPU optimized)",
        torsh_time * 0.3
    );
    println!("      - Estimated time: ~{:.2}ms (CPU)", torsh_time * 1.5);
    println!("      - Memory: Python overhead + CUDA");
    println!("      - Backend: C++/CUDA with Python bindings");

    println!("\n📈 Performance Characteristics:");
    println!("   ✅ ToRSh Advantages:");
    println!("      - Zero Python overhead");
    println!("      - Memory safety without GC");
    println!("      - Compile-time optimization");
    println!("      - Native performance");

    println!("\n   🎯 Optimization Opportunities:");
    println!("      - GPU acceleration (CUDA kernels)");
    println!("      - SIMD vectorization");
    println!("      - Sparse operations optimization");
    println!("      - Batch processing");

    // Performance expectations (adjusted for realistic performance in test builds)
    // Relaxed thresholds to account for debug builds, varied hardware, and parallel test execution
    assert!(
        torsh_time < 2000.0,
        "ToRSh should maintain competitive performance: {:.2}ms",
        torsh_time
    );
    assert!(
        torsh_throughput > 50.0,
        "Should handle reasonable throughput: {:.0} nodes/sec",
        torsh_throughput
    );
}

/// Resource usage and efficiency metrics
#[test]
fn test_resource_efficiency_metrics() {
    println!("\n📈 Resource Efficiency Metrics");
    println!("{}", "=".repeat(45));

    let graph = create_benchmark_graph(&BenchmarkConfig::large());

    // Different layer configurations
    let configurations = vec![
        ("Lightweight", GCNConv::new(128, 64, false)),
        ("Standard", GCNConv::new(128, 128, true)),
        ("Heavy", GCNConv::new(128, 256, true)),
    ];

    for (name, layer) in configurations {
        println!("\n📊 Configuration: {}", name);

        // Parameter count
        let params = layer.parameters();
        let total_params: usize = params
            .iter()
            .map(|p| p.shape().dims().iter().product::<usize>())
            .sum();

        // Performance benchmark
        let start = Instant::now();
        for _ in 0..10 {
            let _output = layer.forward(&graph);
        }
        let duration = start.elapsed().as_millis() as f64 / 10.0;

        // Efficiency metrics
        let params_per_ms = total_params as f64 / duration;
        let nodes_per_ms = graph.num_nodes as f64 / duration;

        println!("   🔢 Parameters: {}", total_params);
        println!("   ⏱️  Time: {:.2}ms", duration);
        println!("   ⚡ Params/ms: {:.0}", params_per_ms);
        println!("   🚀 Nodes/ms: {:.1}", nodes_per_ms);

        // Reasonable efficiency expectations (adjusted for realistic performance in test builds)
        // Relaxed thresholds to account for debug builds and varied hardware
        assert!(
            params_per_ms > 5.0,
            "Parameter efficiency should be reasonable: {:.1} params/ms",
            params_per_ms
        );
        assert!(
            nodes_per_ms > 0.1,
            "Node processing efficiency should be reasonable: {:.1} nodes/ms",
            nodes_per_ms
        );
    }

    println!("\n✅ All performance benchmarks completed successfully!");
}