//! GPU vs CPU Performance Benchmarks
//! 
//! Comprehensive benchmarks comparing GPU and CPU performance for ML operations

use super::utils::*;
use candle_core::{Device, DType, Tensor};
use log::{info, warn};
use std::time::{Duration, Instant};

#[cfg(test)]
mod tests {
    use super::*;
    
    struct BenchmarkResult {
        operation: String,
        cpu_time: Duration,
        gpu_time: Option<Duration>,
        speedup: Option<f64>,
        data_size: String,
    }
    
    impl BenchmarkResult {
        fn new(operation: String, cpu_time: Duration, gpu_time: Option<Duration>, data_size: String) -> Self {
            let speedup = gpu_time.map(|gpu| cpu_time.as_nanos() as f64 / gpu.as_nanos() as f64);
            Self {
                operation,
                cpu_time,
                gpu_time,
                speedup,
                data_size,
            }
        }
        
        fn log_result(&self) {
            info!("📊 Benchmark: {}", self.operation);
            info!("   Data size: {}", self.data_size);
            info!("   CPU time: {:?}", self.cpu_time);
            
            if let Some(gpu_time) = self.gpu_time {
                info!("   GPU time: {:?}", gpu_time);
                if let Some(speedup) = self.speedup {
                    if speedup > 1.0 {
                        info!("   🚀 GPU speedup: {:.2}x", speedup);
                    } else {
                        info!("   ⚠️  GPU slower: {:.2}x", 1.0/speedup);
                    }
                }
            } else {
                info!("   ⚠️  GPU not available for comparison");
            }
        }
    }
    
    fn benchmark_operation<F>(name: &str, operation: F, warmup_iterations: usize, bench_iterations: usize) -> Duration
    where
        F: Fn() -> Result<(), Box<dyn std::error::Error>>,
    {
        // Warmup
        for _ in 0..warmup_iterations {
            let _ = operation();
        }
        
        // Benchmark
        let start = Instant::now();
        for _ in 0..bench_iterations {
            if let Err(e) = operation() {
                warn!("Benchmark operation '{}' failed: {}", name, e);
                break;
            }
        }
        let total_time = start.elapsed();
        
        // Return average time per operation
        total_time / bench_iterations as u32
    }
    
    #[test]
    fn test_matrix_multiplication_benchmark() {
        init_gpu_test_env();
        info!("🔢 Benchmarking matrix multiplication (GPU vs CPU)");
        
        let cpu_device = Device::Cpu;
        let gpu_device = Device::cuda_if_available(0);
        
        let sizes = vec![
            (256, 256, 256),
            (512, 512, 512),
            (1024, 1024, 1024),
            (2048, 1024, 512),
        ];
        
        let mut results = Vec::new();
        
        for (m, k, n) in sizes {
            info!("Testing matrix multiplication: {}x{} * {}x{}", m, k, k, n);
            
            // CPU benchmark
            let cpu_time = benchmark_operation(
                "CPU MatMul",
                || {
                    let a = Tensor::randn(0.0, 1.0, (m, k), &cpu_device)?;
                    let b = Tensor::randn(0.0, 1.0, (k, n), &cpu_device)?;
                    let _result = a.matmul(&b)?;
                    Ok(())
                },
                3, // warmup
                10, // benchmark iterations
            );
            
            // GPU benchmark (if available)
            let gpu_time = if let Ok(ref gpu_dev) = gpu_device {
                Some(benchmark_operation(
                    "GPU MatMul",
                    || {
                        let a = Tensor::randn(0.0, 1.0, (m, k), gpu_dev)?;
                        let b = Tensor::randn(0.0, 1.0, (k, n), gpu_dev)?;
                        let _result = a.matmul(&b)?;
                        Ok(())
                    },
                    3, // warmup
                    10, // benchmark iterations
                ))
            } else {
                None
            };
            
            let result = BenchmarkResult::new(
                "Matrix Multiplication".to_string(),
                cpu_time,
                gpu_time,
                format!("{}x{} * {}x{}", m, k, k, n),
            );
            
            result.log_result();
            results.push(result);
        }
        
        info!("✅ Matrix multiplication benchmark completed");
    }
    
    #[test]
    fn test_tensor_operations_benchmark() {
        init_gpu_test_env();
        info!("🧮 Benchmarking tensor operations (GPU vs CPU)");
        
        let cpu_device = Device::Cpu;
        let gpu_device = Device::cuda_if_available(0);
        
        let tensor_size = (1000, 1000);
        
        // Element-wise operations benchmark
        let operations = vec![
            ("Addition", |a: &Tensor, b: &Tensor| a.add(b)),
            ("Multiplication", |a: &Tensor, b: &Tensor| a.mul(b)),
            ("Subtraction", |a: &Tensor, b: &Tensor| a.sub(b)),
            ("Division", |a: &Tensor, b: &Tensor| a.div(b)),
        ];
        
        for (op_name, op_fn) in operations {
            info!("Benchmarking {}", op_name);
            
            // CPU benchmark
            let cpu_time = benchmark_operation(
                &format!("CPU {}", op_name),
                || {
                    let a = Tensor::randn(0.0, 1.0, tensor_size, &cpu_device)?;
                    let b = Tensor::randn(0.0, 1.0, tensor_size, &cpu_device)?;
                    let _result = op_fn(&a, &b)?;
                    Ok(())
                },
                5, // warmup
                20, // benchmark iterations
            );
            
            // GPU benchmark (if available)
            let gpu_time = if let Ok(ref gpu_dev) = gpu_device {
                Some(benchmark_operation(
                    &format!("GPU {}", op_name),
                    || {
                        let a = Tensor::randn(0.0, 1.0, tensor_size, gpu_dev)?;
                        let b = Tensor::randn(0.0, 1.0, tensor_size, gpu_dev)?;
                        let _result = op_fn(&a, &b)?;
                        Ok(())
                    },
                    5, // warmup
                    20, // benchmark iterations
                ))
            } else {
                None
            };
            
            let result = BenchmarkResult::new(
                op_name.to_string(),
                cpu_time,
                gpu_time,
                format!("{}x{}", tensor_size.0, tensor_size.1),
            );
            
            result.log_result();
        }
        
        info!("✅ Tensor operations benchmark completed");
    }
    
    #[test]
    fn test_convolution_benchmark() {
        init_gpu_test_env();
        info!("🔄 Benchmarking convolution operations (GPU vs CPU)");
        
        let cpu_device = Device::Cpu;
        let gpu_device = Device::cuda_if_available(0);
        
        // Test different convolution sizes
        let configs = vec![
            ("Small Conv", (1, 3, 32, 32), (16, 3, 3, 3)), // batch, channels, height, width
            ("Medium Conv", (8, 16, 64, 64), (32, 16, 5, 5)),
            ("Large Conv", (16, 32, 128, 128), (64, 32, 7, 7)),
        ];
        
        for (config_name, input_shape, kernel_shape) in configs {
            info!("Testing {}: input {:?}, kernel {:?}", config_name, input_shape, kernel_shape);
            
            // CPU benchmark
            let cpu_time = benchmark_operation(
                &format!("CPU {}", config_name),
                || {
                    let input = Tensor::randn(0.0, 1.0, input_shape, &cpu_device)?;
                    let kernel = Tensor::randn(0.0, 1.0, kernel_shape, &cpu_device)?;
                    
                    // Simple convolution operation (using available candle operations)
                    let _result = input.conv2d(&kernel, 1, 1, 1, 1)?;
                    Ok(())
                },
                2, // warmup
                5, // benchmark iterations (fewer for expensive operations)
            );
            
            // GPU benchmark (if available)
            let gpu_time = if let Ok(ref gpu_dev) = gpu_device {
                Some(benchmark_operation(
                    &format!("GPU {}", config_name),
                    || {
                        let input = Tensor::randn(0.0, 1.0, input_shape, gpu_dev)?;
                        let kernel = Tensor::randn(0.0, 1.0, kernel_shape, gpu_dev)?;
                        
                        let _result = input.conv2d(&kernel, 1, 1, 1, 1)?;
                        Ok(())
                    },
                    2, // warmup
                    5, // benchmark iterations
                ))
            } else {
                None
            };
            
            let result = BenchmarkResult::new(
                format!("Convolution {}", config_name),
                cpu_time,
                gpu_time,
                format!("input {:?}, kernel {:?}", input_shape, kernel_shape),
            );
            
            result.log_result();
        }
        
        info!("✅ Convolution benchmark completed");
    }
    
    #[test]
    fn test_memory_bandwidth_benchmark() {
        init_gpu_test_env();
        info!("💾 Benchmarking memory bandwidth (GPU vs CPU)");
        
        let cpu_device = Device::Cpu;
        let gpu_device = Device::cuda_if_available(0);
        
        let sizes = vec![
            (1024, 1024),      // 4MB
            (2048, 2048),      // 16MB
            (4096, 4096),      // 64MB
            (8192, 4096),      // 128MB
        ];
        
        for (rows, cols) in sizes {
            let size_mb = (rows * cols * 4) as f64 / (1024.0 * 1024.0); // f32 = 4 bytes
            info!("Testing memory operations with {:.1}MB tensors ({}x{})", size_mb, rows, cols);
            
            // CPU memory copy benchmark
            let cpu_time = benchmark_operation(
                "CPU Memory Copy",
                || {
                    let src = Tensor::randn(0.0, 1.0, (rows, cols), &cpu_device)?;
                    let _dst = src.copy()?; // Copy operation
                    Ok(())
                },
                3, // warmup
                10, // benchmark iterations
            );
            
            // GPU memory copy benchmark (if available)
            let gpu_time = if let Ok(ref gpu_dev) = gpu_device {
                Some(benchmark_operation(
                    "GPU Memory Copy",
                    || {
                        let src = Tensor::randn(0.0, 1.0, (rows, cols), gpu_dev)?;
                        let _dst = src.copy()?;
                        Ok(())
                    },
                    3, // warmup
                    10, // benchmark iterations
                ))
            } else {
                None
            };
            
            let result = BenchmarkResult::new(
                "Memory Copy".to_string(),
                cpu_time,
                gpu_time,
                format!("{:.1}MB ({}x{})", size_mb, rows, cols),
            );
            
            result.log_result();
            
            // Calculate bandwidth
            let data_size_bytes = (rows * cols * 4) as f64;
            let cpu_bandwidth_gbps = data_size_bytes / (cpu_time.as_secs_f64() * 1e9);
            info!("   CPU bandwidth: {:.2} GB/s", cpu_bandwidth_gbps);
            
            if let Some(gpu_time) = gpu_time {
                let gpu_bandwidth_gbps = data_size_bytes / (gpu_time.as_secs_f64() * 1e9);
                info!("   GPU bandwidth: {:.2} GB/s", gpu_bandwidth_gbps);
            }
        }
        
        info!("✅ Memory bandwidth benchmark completed");
    }
    
    #[test]
    fn test_ml_inference_latency_benchmark() {
        init_gpu_test_env();
        info!("🧠 Benchmarking ML inference latency (GPU vs CPU)");
        
        let cpu_device = Device::Cpu;
        let gpu_device = Device::cuda_if_available(0);
        
        // Simple neural network simulation
        let network_configs = vec![
            ("Small Network", vec![784, 128, 64, 10]),
            ("Medium Network", vec![1024, 512, 256, 128, 10]),
            ("Large Network", vec![2048, 1024, 512, 256, 128, 10]),
        ];
        
        for (config_name, layer_sizes) in network_configs {
            info!("Testing {} with layers: {:?}", config_name, layer_sizes);
            
            let batch_size = 32;
            
            // CPU inference benchmark
            let cpu_time = benchmark_operation(
                &format!("CPU {}", config_name),
                || {
                    let mut x = Tensor::randn(0.0, 1.0, (batch_size, layer_sizes[0]), &cpu_device)?;
                    
                    // Simulate forward pass through network
                    for i in 0..layer_sizes.len() - 1 {
                        let weight = Tensor::randn(0.0, 1.0, (layer_sizes[i], layer_sizes[i + 1]), &cpu_device)?;
                        x = x.matmul(&weight)?;
                        x = x.relu()?; // Activation function
                    }
                    
                    Ok(())
                },
                3, // warmup
                10, // benchmark iterations
            );
            
            // GPU inference benchmark (if available)
            let gpu_time = if let Ok(ref gpu_dev) = gpu_device {
                Some(benchmark_operation(
                    &format!("GPU {}", config_name),
                    || {
                        let mut x = Tensor::randn(0.0, 1.0, (batch_size, layer_sizes[0]), gpu_dev)?;
                        
                        for i in 0..layer_sizes.len() - 1 {
                            let weight = Tensor::randn(0.0, 1.0, (layer_sizes[i], layer_sizes[i + 1]), gpu_dev)?;
                            x = x.matmul(&weight)?;
                            x = x.relu()?;
                        }
                        
                        Ok(())
                    },
                    3, // warmup
                    10, // benchmark iterations
                ))
            } else {
                None
            };
            
            let result = BenchmarkResult::new(
                format!("ML Inference {}", config_name),
                cpu_time,
                gpu_time,
                format!("layers {:?}, batch_size {}", layer_sizes, batch_size),
            );
            
            result.log_result();
            
            // Check if we're meeting HFT latency requirements
            let cpu_latency_us = cpu_time.as_micros();
            info!("   CPU latency per batch: {}μs", cpu_latency_us);
            
            if let Some(gpu_time) = gpu_time {
                let gpu_latency_us = gpu_time.as_micros();
                info!("   GPU latency per batch: {}μs", gpu_latency_us);
                
                // Check if we're meeting the claimed sub-50μs requirements
                if gpu_latency_us < 50 {
                    info!("   ✅ GPU meets sub-50μs HFT requirement!");
                } else {
                    warn!("   ⚠️  GPU exceeds 50μs HFT requirement");
                }
            }
            
            if cpu_latency_us < 50 {
                info!("   ✅ CPU meets sub-50μs HFT requirement!");
            } else {
                warn!("   ⚠️  CPU exceeds 50μs HFT requirement");
            }
        }
        
        info!("✅ ML inference latency benchmark completed");
    }
    
    #[test]
    fn test_comprehensive_performance_summary() {
        init_gpu_test_env();
        info!("📈 Generating comprehensive performance summary");
        
        let cpu_device = Device::Cpu;
        let gpu_available = Device::cuda_if_available(0).is_ok();
        
        info!("=== GPU PERFORMANCE SUMMARY ===");
        info!("GPU Available: {}", gpu_available);
        info!("CPU Device: {:?}", cpu_device);
        
        if gpu_available {
            info!("GPU Device: {:?}", Device::cuda_if_available(0).unwrap());
            
            // Quick performance test
            let test_size = (1000, 1000);
            
            let cpu_start = Instant::now();
            let _cpu_tensor = Tensor::randn(0.0, 1.0, test_size, &cpu_device).unwrap();
            let cpu_time = cpu_start.elapsed();
            
            let gpu_device = Device::cuda_if_available(0).unwrap();
            let gpu_start = Instant::now();
            let _gpu_tensor = Tensor::randn(0.0, 1.0, test_size, &gpu_device).unwrap();
            let gpu_time = gpu_start.elapsed();
            
            info!("Tensor creation ({}x{}):", test_size.0, test_size.1);
            info!("  CPU: {:?}", cpu_time);
            info!("  GPU: {:?}", gpu_time);
            
            let speedup = cpu_time.as_nanos() as f64 / gpu_time.as_nanos() as f64;
            if speedup > 1.0 {
                info!("  🚀 GPU speedup: {:.2}x", speedup);
            } else {
                info!("  ⚠️  GPU slower: {:.2}x", 1.0/speedup);
            }
        } else {
            warn!("⚠️  GPU not available - CPU-only performance");
        }
        
        info!("=== PERFORMANCE RECOMMENDATIONS ===");
        if gpu_available {
            info!("✅ GPU acceleration available - recommend enabling for production");
            info!("✅ Use GPU for large matrix operations and ML inference");
            info!("✅ Consider GPU memory pooling for optimal performance");
        } else {
            info!("💻 CPU-only mode - consider GPU setup for better performance");
            info!("💻 Optimize CPU operations with SIMD and vectorization");
        }
        
        info!("✅ Comprehensive performance summary completed");
    }
}