//! GPU index building benchmark framework
//!
//! Provides reproducible benchmarks for comparing GPU vs CPU index building
//! performance across different dataset sizes, dimensions, and index types.
//!
//! ## Benchmark Suites
//!
//! - **KMeans**: GPU vs CPU KMeans clustering (varying K, N, D)
//! - **Quantization**: GPU vs CPU PQ/SQ training
//! - **CAGRA**: GPU graph build vs CPU k-NN graph
//! - **Pipeline**: End-to-end IVF index build (GPU vs CPU)
//! - **Scaling**: How throughput scales with dataset size
//!
//! ## Usage
//!
//! ```ignore
//! use aifastdb_core::gpu::bench::*;
//!
//! let suite = BenchSuite::default();
//! let report = suite.run_all();
//! println!("{}", report.to_markdown());
//! ```

use std::time::{Duration, Instant};

use super::cagra::CagraConfig;
use super::config::GpuIndexConfig;
use super::kmeans::{GpuKMeansConfig, GpuKMeansTrainer};
use super::memory::GpuMemoryPool;
use super::pipeline::IndexBuildPipeline;
use super::quantization::{GpuPqTrainer, GpuSqTrainer, PqTrainConfig, SqTrainConfig};

// ============================================================================
// Benchmark Configuration
// ============================================================================

/// Configuration for the benchmark suite.
#[derive(Debug, Clone)]
pub struct BenchConfig {
    /// Dataset sizes to benchmark (number of vectors)
    pub dataset_sizes: Vec<usize>,
    /// Dimensions to benchmark
    pub dimensions: Vec<usize>,
    /// Number of KMeans clusters
    pub kmeans_clusters: Vec<usize>,
    /// Number of warmup runs before timing
    pub warmup_runs: usize,
    /// Number of timed runs (results are averaged)
    pub timed_runs: usize,
    /// Whether to include GPU benchmarks (auto-detected if not set)
    pub include_gpu: Option<bool>,
    /// Whether to print progress during benchmarking
    pub verbose: bool,
}

impl Default for BenchConfig {
    fn default() -> Self {
        Self {
            dataset_sizes: vec![1_000, 10_000, 50_000, 100_000],
            dimensions: vec![128, 768],
            kmeans_clusters: vec![16, 64, 256],
            warmup_runs: 1,
            timed_runs: 3,
            include_gpu: None,
            verbose: true,
        }
    }
}

impl BenchConfig {
    /// Quick benchmark (small datasets, single run).
    pub fn quick() -> Self {
        Self {
            dataset_sizes: vec![1_000, 10_000],
            dimensions: vec![128],
            kmeans_clusters: vec![16],
            warmup_runs: 0,
            timed_runs: 1,
            include_gpu: None,
            verbose: true,
        }
    }

    /// Full benchmark (large datasets, multiple runs).
    pub fn full() -> Self {
        Self {
            dataset_sizes: vec![1_000, 10_000, 50_000, 100_000, 500_000],
            dimensions: vec![128, 384, 768, 1536],
            kmeans_clusters: vec![16, 64, 256, 1024],
            warmup_runs: 2,
            timed_runs: 5,
            include_gpu: None,
            verbose: true,
        }
    }
}

// ============================================================================
// Benchmark Result Types
// ============================================================================

/// A single benchmark measurement.
#[derive(Debug, Clone)]
pub struct BenchMeasurement {
    /// Benchmark name
    pub name: String,
    /// Dataset size (number of vectors)
    pub n_vectors: usize,
    /// Vector dimension
    pub dimension: usize,
    /// Additional parameters (e.g., "k=256")
    pub params: String,
    /// CPU execution time
    pub cpu_time: Duration,
    /// GPU execution time (None if GPU not available)
    pub gpu_time: Option<Duration>,
    /// Speedup ratio (CPU time / GPU time). >1 means GPU is faster.
    pub speedup: Option<f64>,
    /// Throughput in vectors/second (CPU)
    pub cpu_throughput: f64,
    /// Throughput in vectors/second (GPU)
    pub gpu_throughput: Option<f64>,
}

impl BenchMeasurement {
    fn new(
        name: &str,
        n_vectors: usize,
        dimension: usize,
        params: &str,
        cpu_time: Duration,
        gpu_time: Option<Duration>,
    ) -> Self {
        let cpu_throughput = n_vectors as f64 / cpu_time.as_secs_f64();
        let gpu_throughput = gpu_time.map(|t| n_vectors as f64 / t.as_secs_f64());
        let speedup = gpu_time.map(|gt| cpu_time.as_secs_f64() / gt.as_secs_f64());

        Self {
            name: name.into(),
            n_vectors,
            dimension,
            params: params.into(),
            cpu_time,
            gpu_time,
            speedup,
            cpu_throughput,
            gpu_throughput,
        }
    }
}

/// Full benchmark report.
#[derive(Debug, Clone)]
pub struct BenchReport {
    /// Individual measurements
    pub measurements: Vec<BenchMeasurement>,
    /// GPU device name (if available)
    pub gpu_device: Option<String>,
    /// Total benchmark time
    pub total_time: Duration,
    /// System info
    pub system_info: String,
}

impl BenchReport {
    /// Render report as a Markdown table.
    pub fn to_markdown(&self) -> String {
        let mut s = String::new();

        s.push_str("# GPU Index Build Benchmark Report\n\n");

        s.push_str(&format!("**System**: {}\n", self.system_info));
        if let Some(ref gpu) = self.gpu_device {
            s.push_str(&format!("**GPU**: {}\n", gpu));
        } else {
            s.push_str("**GPU**: Not available (CPU-only results)\n");
        }
        s.push_str(&format!(
            "**Total Time**: {:.1}s\n\n",
            self.total_time.as_secs_f64()
        ));

        // Table header
        s.push_str("| Benchmark | N | Dim | Params | CPU (ms) | GPU (ms) | Speedup | CPU vecs/s | GPU vecs/s |\n");
        s.push_str("|-----------|---|-----|--------|----------|----------|---------|------------|------------|\n");

        for m in &self.measurements {
            let gpu_ms = m
                .gpu_time
                .map(|t| format!("{:.1}", t.as_secs_f64() * 1000.0))
                .unwrap_or_else(|| "N/A".into());
            let speedup = m
                .speedup
                .map(|s| format!("{:.1}x", s))
                .unwrap_or_else(|| "N/A".into());
            let gpu_tp = m
                .gpu_throughput
                .map(|t| format!("{:.0}", t))
                .unwrap_or_else(|| "N/A".into());

            s.push_str(&format!(
                "| {} | {} | {} | {} | {:.1} | {} | {} | {:.0} | {} |\n",
                m.name,
                m.n_vectors,
                m.dimension,
                m.params,
                m.cpu_time.as_secs_f64() * 1000.0,
                gpu_ms,
                speedup,
                m.cpu_throughput,
                gpu_tp,
            ));
        }

        s.push_str("\n");

        // Summary
        let gpu_measurements: Vec<_> = self.measurements.iter().filter(|m| m.speedup.is_some()).collect();
        if !gpu_measurements.is_empty() {
            let avg_speedup: f64 =
                gpu_measurements.iter().map(|m| m.speedup.unwrap()).sum::<f64>()
                    / gpu_measurements.len() as f64;
            let max_speedup = gpu_measurements
                .iter()
                .map(|m| m.speedup.unwrap())
                .fold(0.0f64, f64::max);

            s.push_str(&format!(
                "**Average GPU Speedup**: {:.1}x | **Max Speedup**: {:.1}x\n",
                avg_speedup, max_speedup
            ));
        }

        s
    }

    /// Render report as CSV.
    pub fn to_csv(&self) -> String {
        let mut s = String::new();
        s.push_str("benchmark,n_vectors,dimension,params,cpu_ms,gpu_ms,speedup,cpu_vecs_per_s,gpu_vecs_per_s\n");

        for m in &self.measurements {
            s.push_str(&format!(
                "{},{},{},{},{:.2},{},{},{:.0},{}\n",
                m.name,
                m.n_vectors,
                m.dimension,
                m.params,
                m.cpu_time.as_secs_f64() * 1000.0,
                m.gpu_time
                    .map(|t| format!("{:.2}", t.as_secs_f64() * 1000.0))
                    .unwrap_or_default(),
                m.speedup.map(|s| format!("{:.2}", s)).unwrap_or_default(),
                m.cpu_throughput,
                m.gpu_throughput
                    .map(|t| format!("{:.0}", t))
                    .unwrap_or_default(),
            ));
        }

        s
    }
}

// ============================================================================
// Benchmark Suite
// ============================================================================

/// Benchmark suite runner.
pub struct BenchSuite {
    config: BenchConfig,
    gpu_available: bool,
}

impl BenchSuite {
    /// Create a new benchmark suite.
    pub fn new(config: BenchConfig) -> Self {
        let gpu_available = config.include_gpu.unwrap_or_else(|| {
            super::gpu_available().is_available()
        });

        Self {
            config,
            gpu_available,
        }
    }

    /// Run all benchmarks and return a report.
    pub fn run_all(&self) -> BenchReport {
        let start = Instant::now();
        let mut measurements = Vec::new();

        if self.config.verbose {
            println!("=== GPU Index Build Benchmark Suite ===");
            println!(
                "GPU: {}",
                if self.gpu_available {
                    "Available"
                } else {
                    "Not available (CPU-only)"
                }
            );
            println!();
        }

        // 1. KMeans benchmarks
        measurements.extend(self.bench_kmeans());

        // 2. SQ Training benchmarks
        measurements.extend(self.bench_sq());

        // 3. PQ Training benchmarks
        measurements.extend(self.bench_pq());

        // 4. CAGRA graph build benchmarks
        measurements.extend(self.bench_cagra());

        // 5. End-to-end IVF pipeline benchmarks
        measurements.extend(self.bench_pipeline());

        let gpu_device = if self.gpu_available {
            super::gpu_available()
                .best_device()
                .map(|d| d.name.clone())
        } else {
            None
        };

        BenchReport {
            measurements,
            gpu_device,
            total_time: start.elapsed(),
            system_info: format!(
                "CPU cores: {}, OS: {}",
                num_cpus::get(),
                std::env::consts::OS
            ),
        }
    }

    // ========================================================================
    // Individual benchmark suites
    // ========================================================================

    fn bench_kmeans(&self) -> Vec<BenchMeasurement> {
        let mut results = Vec::new();

        for &n in &self.config.dataset_sizes {
            for &dim in &self.config.dimensions {
                for &k in &self.config.kmeans_clusters {
                    if k >= n {
                        continue; // Skip invalid combinations
                    }

                    let data = generate_random_data(n, dim);
                    let params = format!("k={}", k);

                    if self.config.verbose {
                        print!("KMeans: N={}, D={}, k={} ... ", n, dim, k);
                    }

                    let cpu_time = self.time_cpu_kmeans(&data, n, dim, k);
                    let gpu_time = if self.gpu_available {
                        Some(self.time_gpu_kmeans(&data, n, dim, k))
                    } else {
                        None
                    };

                    let m = BenchMeasurement::new("KMeans", n, dim, &params, cpu_time, gpu_time);

                    if self.config.verbose {
                        println!(
                            "CPU={:.1}ms{}",
                            cpu_time.as_secs_f64() * 1000.0,
                            gpu_time
                                .map(|t| format!(
                                    ", GPU={:.1}ms, speedup={:.1}x",
                                    t.as_secs_f64() * 1000.0,
                                    m.speedup.unwrap_or(0.0)
                                ))
                                .unwrap_or_default()
                        );
                    }

                    results.push(m);
                }
            }
        }

        results
    }

    fn bench_sq(&self) -> Vec<BenchMeasurement> {
        let mut results = Vec::new();

        for &n in &self.config.dataset_sizes {
            for &dim in &self.config.dimensions {
                let data = generate_random_data(n, dim);

                if self.config.verbose {
                    print!("SQ Train: N={}, D={} ... ", n, dim);
                }

                let cpu_time = self.time_cpu_sq(&data, n, dim);

                let m = BenchMeasurement::new("SQ Train", n, dim, "", cpu_time, None);

                if self.config.verbose {
                    println!("CPU={:.1}ms", cpu_time.as_secs_f64() * 1000.0);
                }

                results.push(m);
            }
        }

        results
    }

    fn bench_pq(&self) -> Vec<BenchMeasurement> {
        let mut results = Vec::new();

        for &n in &self.config.dataset_sizes {
            for &dim in &self.config.dimensions {
                if dim < 16 {
                    continue;
                }
                let n_sub = (dim / 8).max(2).min(dim / 2);
                let data = generate_random_data(n, dim);
                let params = format!("sub={}", n_sub);

                if self.config.verbose {
                    print!("PQ Train: N={}, D={}, sub={} ... ", n, dim, n_sub);
                }

                let cpu_time = self.time_cpu_pq(&data, n, dim, n_sub);

                let m = BenchMeasurement::new("PQ Train", n, dim, &params, cpu_time, None);

                if self.config.verbose {
                    println!("CPU={:.1}ms", cpu_time.as_secs_f64() * 1000.0);
                }

                results.push(m);
            }
        }

        results
    }

    fn bench_cagra(&self) -> Vec<BenchMeasurement> {
        let mut results = Vec::new();

        // CAGRA is O(N²) for brute-force k-NN, so use smaller datasets
        let cagra_sizes: Vec<usize> = self
            .config
            .dataset_sizes
            .iter()
            .copied()
            .filter(|&n| n <= 10_000)
            .collect();

        for &n in &cagra_sizes {
            for &dim in &self.config.dimensions {
                let data = generate_random_data(n, dim);
                let degree = 16;
                let params = format!("deg={}", degree);

                if self.config.verbose {
                    print!("CAGRA Build: N={}, D={}, deg={} ... ", n, dim, degree);
                }

                let cpu_time = self.time_cpu_cagra(&data, n, dim, degree);

                let m = BenchMeasurement::new("CAGRA Build", n, dim, &params, cpu_time, None);

                if self.config.verbose {
                    println!("CPU={:.1}ms", cpu_time.as_secs_f64() * 1000.0);
                }

                results.push(m);
            }
        }

        results
    }

    fn bench_pipeline(&self) -> Vec<BenchMeasurement> {
        let mut results = Vec::new();

        for &n in &self.config.dataset_sizes {
            for &dim in &self.config.dimensions {
                let data = generate_random_data(n, dim);
                let n_partitions = (n as f64).sqrt() as usize;
                let n_partitions = n_partitions.max(4).min(1024);
                let params = format!("parts={}", n_partitions);

                if self.config.verbose {
                    print!(
                        "IVF Pipeline: N={}, D={}, parts={} ... ",
                        n, dim, n_partitions
                    );
                }

                let cpu_time = self.time_cpu_pipeline(&data, n, dim, n_partitions);

                let m =
                    BenchMeasurement::new("IVF Pipeline", n, dim, &params, cpu_time, None);

                if self.config.verbose {
                    println!("CPU={:.1}ms", cpu_time.as_secs_f64() * 1000.0);
                }

                results.push(m);
            }
        }

        results
    }

    // ========================================================================
    // Timing helpers
    // ========================================================================

    fn time_cpu_kmeans(&self, data: &[f32], n: usize, dim: usize, k: usize) -> Duration {
        self.time_fn(|| {
            let config = GpuKMeansConfig {
                n_clusters: k,
                max_iterations: 10,
                ..Default::default()
            };
            let trainer = GpuKMeansTrainer::new(config);
            // Pass None for memory_pool to force CPU path
            let _ = trainer.train(data, n, dim, None);
        })
    }

    fn time_gpu_kmeans(&self, data: &[f32], n: usize, dim: usize, k: usize) -> Duration {
        let pool = self.create_memory_pool();
        self.time_fn(|| {
            let config = GpuKMeansConfig {
                n_clusters: k,
                max_iterations: 10,
                ..Default::default()
            };
            let trainer = GpuKMeansTrainer::new(config);
            let _ = trainer.train(data, n, dim, pool.as_ref());
        })
    }

    fn time_cpu_sq(&self, data: &[f32], n: usize, dim: usize) -> Duration {
        self.time_fn(|| {
            let config = SqTrainConfig::default();
            let trainer = GpuSqTrainer::new(config);
            let _ = trainer.train(data, n, dim, None);
        })
    }

    fn time_cpu_pq(&self, data: &[f32], n: usize, dim: usize, n_sub: usize) -> Duration {
        self.time_fn(|| {
            let config = PqTrainConfig {
                num_sub_vectors: n_sub,
                ..Default::default()
            };
            let trainer = GpuPqTrainer::new(config);
            let _ = trainer.train(data, n, dim, None);
        })
    }

    fn time_cpu_cagra(&self, data: &[f32], n: usize, dim: usize, degree: usize) -> Duration {
        self.time_fn(|| {
            let config = CagraConfig {
                graph_degree: degree,
                ..Default::default()
            };
            let _ = super::cagra::CagraIndex::build(data, n, dim, &config, None);
        })
    }

    fn time_cpu_pipeline(
        &self,
        data: &[f32],
        n: usize,
        dim: usize,
        n_partitions: usize,
    ) -> Duration {
        self.time_fn(|| {
            let pipeline = IndexBuildPipeline::new(GpuIndexConfig::default());
            let _ = pipeline.build_ivf(data, n, dim, n_partitions);
        })
    }

    /// Run a function multiple times and return the average duration.
    fn time_fn<F: Fn()>(&self, f: F) -> Duration {
        // Warmup
        for _ in 0..self.config.warmup_runs {
            f();
        }

        // Timed runs
        let runs = self.config.timed_runs.max(1);
        let mut total = Duration::ZERO;

        for _ in 0..runs {
            let start = Instant::now();
            f();
            total += start.elapsed();
        }

        total / runs as u32
    }

    fn create_memory_pool(&self) -> Option<GpuMemoryPool> {
        super::gpu_available()
            .best_device()
            .map(|info| GpuMemoryPool::new(info.clone(), None))
    }
}

// ============================================================================
// Data generation
// ============================================================================

/// Generate pseudo-random float data for benchmarking.
///
/// Uses a simple LCG to be deterministic and fast (no dependency on `rand`).
fn generate_random_data(n_vectors: usize, dimension: usize) -> Vec<f32> {
    let total = n_vectors * dimension;
    let mut data = Vec::with_capacity(total);
    let mut state: u64 = 42;

    for _ in 0..total {
        state = state.wrapping_mul(6364136223846793005).wrapping_add(1);
        let val = ((state >> 33) as f32) / (u32::MAX as f32) * 2.0 - 1.0;
        data.push(val);
    }

    data
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_generate_random_data() {
        let data = generate_random_data(100, 4);
        assert_eq!(data.len(), 400);

        // Check values are in reasonable range
        for &v in &data {
            assert!(v >= -1.5 && v <= 1.5, "Value out of range: {}", v);
        }

        // Check determinism
        let data2 = generate_random_data(100, 4);
        assert_eq!(data, data2);
    }

    #[test]
    fn test_bench_measurement() {
        let m = BenchMeasurement::new(
            "test",
            10000,
            128,
            "k=16",
            Duration::from_millis(100),
            Some(Duration::from_millis(10)),
        );
        assert_eq!(m.name, "test");
        assert_eq!(m.n_vectors, 10000);
        assert!((m.speedup.unwrap() - 10.0).abs() < 0.1);
        assert!(m.cpu_throughput > 0.0);
        assert!(m.gpu_throughput.unwrap() > 0.0);
    }

    #[test]
    fn test_bench_measurement_no_gpu() {
        let m = BenchMeasurement::new(
            "test",
            10000,
            128,
            "k=16",
            Duration::from_millis(100),
            None,
        );
        assert!(m.speedup.is_none());
        assert!(m.gpu_throughput.is_none());
    }

    #[test]
    fn test_bench_report_markdown() {
        let report = BenchReport {
            measurements: vec![
                BenchMeasurement::new(
                    "KMeans",
                    10000,
                    128,
                    "k=16",
                    Duration::from_millis(50),
                    Some(Duration::from_millis(5)),
                ),
                BenchMeasurement::new(
                    "KMeans",
                    50000,
                    128,
                    "k=16",
                    Duration::from_millis(200),
                    None,
                ),
            ],
            gpu_device: Some("Test GPU".into()),
            total_time: Duration::from_secs(1),
            system_info: "Test System".into(),
        };

        let md = report.to_markdown();
        assert!(md.contains("GPU Index Build Benchmark Report"));
        assert!(md.contains("Test GPU"));
        assert!(md.contains("KMeans"));
        assert!(md.contains("10.0x")); // speedup
        assert!(md.contains("N/A")); // no GPU for second entry
    }

    #[test]
    fn test_bench_report_csv() {
        let report = BenchReport {
            measurements: vec![BenchMeasurement::new(
                "KMeans",
                10000,
                128,
                "k=16",
                Duration::from_millis(50),
                None,
            )],
            gpu_device: None,
            total_time: Duration::from_secs(1),
            system_info: "Test".into(),
        };

        let csv = report.to_csv();
        assert!(csv.contains("benchmark,n_vectors"));
        assert!(csv.contains("KMeans,10000,128"));
    }

    #[test]
    fn test_quick_bench_config() {
        let config = BenchConfig::quick();
        assert_eq!(config.dataset_sizes.len(), 2);
        assert_eq!(config.timed_runs, 1);
    }

    #[test]
    fn test_bench_suite_quick() {
        // Run a minimal benchmark to verify the framework works
        let config = BenchConfig {
            dataset_sizes: vec![200],
            dimensions: vec![32],
            kmeans_clusters: vec![4],
            warmup_runs: 0,
            timed_runs: 1,
            include_gpu: Some(false), // Force CPU-only for test
            verbose: false,
        };

        let suite = BenchSuite::new(config);
        let report = suite.run_all();

        assert!(!report.measurements.is_empty());

        // Should have at least KMeans, SQ, PQ, CAGRA, Pipeline measurements
        let names: Vec<&str> = report.measurements.iter().map(|m| m.name.as_str()).collect();
        assert!(names.contains(&"KMeans"), "Missing KMeans benchmark");
        assert!(names.contains(&"SQ Train"), "Missing SQ benchmark");
        assert!(names.contains(&"PQ Train"), "Missing PQ benchmark, got: {:?}", names);
        assert!(names.contains(&"CAGRA Build"), "Missing CAGRA benchmark");
        assert!(
            names.contains(&"IVF Pipeline"),
            "Missing Pipeline benchmark"
        );

        // All CPU times should be > 0
        for m in &report.measurements {
            assert!(
                m.cpu_time > Duration::ZERO,
                "{}: CPU time should be > 0",
                m.name
            );
            assert!(
                m.cpu_throughput > 0.0,
                "{}: throughput should be > 0",
                m.name
            );
            // No GPU in this test
            assert!(m.gpu_time.is_none());
        }
    }
}
