//! GPU Batch Inference Benchmarks
//!
//! This benchmark specifically tests batch inference performance to identify
//! why GPU speedup is only 1.05x instead of the target 10x.
//!
//! Key insights:
//! 1. Small models don't benefit from GPU (overhead dominates)
//! 2. Single inference has high CPU→GPU transfer overhead
//! 3. GPU shines with batch sizes ≥32
//! 4. FP16 precision doubles throughput

#![allow(unused_crate_dependencies)]

use candle_core::{Device, DType, Tensor};
use criterion::{black_box, criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
use std::time::Duration;

/// Generate input tensor on device (do NOT recreate inside benchmark loop!)
fn create_input_tensor(shape: &[usize], device: &Device) -> Tensor {
    Tensor::randn(0.0f32, 1.0f32, shape, device).expect("Failed to create tensor")
}

/// Simulate realistic neural network inference
fn simulate_forward_pass(input: &Tensor, weights: &Tensor) -> Tensor {
    // Matrix multiplication + activation
    let output = input.matmul(weights).expect("matmul failed");
    output.relu().expect("relu failed")
}

/// Test 1: Single vs Batch Inference (CPU)
fn bench_cpu_single_vs_batch(c: &mut Criterion) {
    let device = Device::Cpu;
    let mut group = c.benchmark_group("cpu_batch_comparison");
    group.measurement_time(Duration::from_secs(10));

    let batch_sizes = vec![1, 8, 16, 32, 64];
    let input_dim = 256;
    let output_dim = 128;

    for batch_size in batch_sizes {
        // Pre-create tensors OUTSIDE benchmark loop
        let input = create_input_tensor(&[batch_size, input_dim], &device);
        let weights = create_input_tensor(&[input_dim, output_dim], &device);

        group.bench_with_input(
            BenchmarkId::new("cpu", batch_size),
            &(input, weights),
            |b, (inp, w)| {
                b.iter(|| {
                    black_box(simulate_forward_pass(inp, w))
                })
            }
        );
    }

    group.finish();
}

/// Test 2: Single vs Batch Inference (GPU)
fn bench_gpu_single_vs_batch(c: &mut Criterion) {
    let gpu_device = match Device::new_cuda(0) {
        Ok(d) => d,
        Err(_) => {
            eprintln!("⚠️  GPU not available, skipping GPU batch benchmark");
            return;
        }
    };

    let mut group = c.benchmark_group("gpu_batch_comparison");
    group.measurement_time(Duration::from_secs(10));

    let batch_sizes = vec![1, 8, 16, 32, 64, 128];
    let input_dim = 256;
    let output_dim = 128;

    for batch_size in batch_sizes {
        // Pre-create tensors on GPU OUTSIDE benchmark loop
        let input = create_input_tensor(&[batch_size, input_dim], &gpu_device);
        let weights = create_input_tensor(&[input_dim, output_dim], &gpu_device);

        group.bench_with_input(
            BenchmarkId::new("gpu", batch_size),
            &(input, weights),
            |b, (inp, w)| {
                b.iter(|| {
                    black_box(simulate_forward_pass(inp, w))
                })
            }
        );
    }

    group.finish();
}

/// Test 3: Data Transfer Overhead
fn bench_cpu_to_gpu_transfer(c: &mut Criterion) {
    let gpu_device = match Device::new_cuda(0) {
        Ok(d) => d,
        Err(_) => return,
    };

    let cpu_device = Device::Cpu;
    let mut group = c.benchmark_group("cpu_to_gpu_transfer");
    group.measurement_time(Duration::from_secs(5));

    let sizes = vec![
        ("small", vec![1, 64]),
        ("medium", vec![32, 256]),
        ("large", vec![128, 512]),
    ];

    for (name, shape) in sizes {
        group.bench_function(name, |b| {
            b.iter_batched(
                || create_input_tensor(&shape, &cpu_device),
                |cpu_tensor| {
                    // Measure CPU→GPU transfer time
                    black_box(cpu_tensor.to_device(&gpu_device).expect("transfer failed"))
                },
                BatchSize::SmallInput
            )
        });
    }

    group.finish();
}

/// Test 4: GPU Utilization - Large Model
fn bench_gpu_large_model(c: &mut Criterion) {
    let gpu_device = match Device::new_cuda(0) {
        Ok(d) => d,
        Err(_) => return,
    };

    let mut group = c.benchmark_group("gpu_large_model");
    group.measurement_time(Duration::from_secs(15));

    // Large model that should benefit from GPU
    let batch_size = 64;
    let layers = vec![
        (512, 1024),
        (1024, 2048),
        (2048, 1024),
        (1024, 256),
    ];

    // Pre-create all tensors on GPU
    let input = create_input_tensor(&[batch_size, layers[0].0], &gpu_device);
    let weights: Vec<Tensor> = layers.iter()
        .map(|(in_dim, out_dim)| create_input_tensor(&[*in_dim, *out_dim], &gpu_device))
        .collect();

    group.bench_function("4_layer_network", |b| {
        b.iter(|| {
            let mut current = input.clone();
            for weight in &weights {
                current = black_box(simulate_forward_pass(&current, weight));
            }
            black_box(current)
        })
    });

    group.finish();
}

/// Test 5: FP16 vs FP32 (GPU only)
fn bench_gpu_precision(c: &mut Criterion) {
    let gpu_device = match Device::new_cuda(0) {
        Ok(d) => d,
        Err(_) => return,
    };

    let mut group = c.benchmark_group("gpu_precision");
    group.measurement_time(Duration::from_secs(10));

    let batch_size = 32;
    let input_dim = 512;
    let output_dim = 256;

    // FP32
    let input_fp32 = create_input_tensor(&[batch_size, input_dim], &gpu_device);
    let weights_fp32 = create_input_tensor(&[input_dim, output_dim], &gpu_device);

    group.bench_function("fp32", |b| {
        b.iter(|| {
            black_box(simulate_forward_pass(&input_fp32, &weights_fp32))
        })
    });

    // FP16
    let input_fp16 = input_fp32.to_dtype(DType::F16).expect("FP16 conversion failed");
    let weights_fp16 = weights_fp32.to_dtype(DType::F16).expect("FP16 conversion failed");

    group.bench_function("fp16", |b| {
        b.iter(|| {
            black_box(simulate_forward_pass(&input_fp16, &weights_fp16))
        })
    });

    group.finish();
}

/// Test 6: Cold Start Penalty (includes model creation)
fn bench_cold_start_overhead(c: &mut Criterion) {
    let gpu_device = match Device::new_cuda(0) {
        Ok(d) => d,
        Err(_) => return,
    };

    let mut group = c.benchmark_group("cold_start");
    group.measurement_time(Duration::from_secs(10));
    group.sample_size(10);

    let input_dim = 256;
    let output_dim = 128;

    group.bench_function("with_tensor_creation", |b| {
        b.iter(|| {
            // This includes tensor creation overhead (simulates cold start)
            let input = create_input_tensor(&[1, input_dim], &gpu_device);
            let weights = create_input_tensor(&[input_dim, output_dim], &gpu_device);
            black_box(simulate_forward_pass(&input, &weights))
        })
    });

    // Pre-create tensors
    let input = create_input_tensor(&[1, input_dim], &gpu_device);
    let weights = create_input_tensor(&[input_dim, output_dim], &gpu_device);

    group.bench_function("warm_cache", |b| {
        b.iter(|| {
            black_box(simulate_forward_pass(&input, &weights))
        })
    });

    group.finish();
}

criterion_group! {
    name = gpu_optimization_benchmarks;
    config = Criterion::default()
        .measurement_time(Duration::from_secs(10))
        .warm_up_time(Duration::from_secs(2));
    targets =
        bench_cpu_single_vs_batch,
        bench_gpu_single_vs_batch,
        bench_cpu_to_gpu_transfer,
        bench_gpu_large_model,
        bench_gpu_precision,
        bench_cold_start_overhead
}

criterion_main!(gpu_optimization_benchmarks);