//! Direct ML Inference Benchmarks
//!
//! Simple, direct timing measurements for ML model inference validation.
//! Tests GPU acceleration and validates <1ms p99 target.

#![allow(unused_crate_dependencies)]

use candle_core::{Device, Tensor};
use candle_nn::ops::softmax;
use std::time::{Duration, Instant};

// Statistics helper
struct BenchStats {
    samples: Vec<Duration>,
}

impl BenchStats {
    fn new() -> Self {
        Self { samples: Vec::new() }
    }

    fn add(&mut self, duration: Duration) {
        self.samples.push(duration);
    }

    fn percentile(&mut self, p: f64) -> Duration {
        self.samples.sort();
        let index = ((p / 100.0) * self.samples.len() as f64) as usize;
        self.samples[index.min(self.samples.len() - 1)]
    }

    fn mean(&self) -> Duration {
        let sum: Duration = self.samples.iter().sum();
        sum / self.samples.len() as u32
    }

    fn min(&self) -> Duration {
        *self.samples.iter().min().unwrap()
    }

    fn max(&self) -> Duration {
        *self.samples.iter().max().unwrap()
    }
}

fn main() {
    println!("{}", "=".repeat(80));
    println!("ML INFERENCE PERFORMANCE BENCHMARKS");
    println!("Wave 131 Phase 2 - Agent 204");
    println!("{}", "=".repeat(80));
    println!();

    // Check GPU availability
    let gpu_available = Device::cuda_if_available(0).is_ok();
    println!("GPU Available: {}", gpu_available);
    if gpu_available {
        println!("GPU Device: NVIDIA RTX 3050 Ti");
    }
    println!();

    let iterations = 1000;

    // =========================================================================
    // MAMBA-2 BENCHMARKS
    // =========================================================================
    println!("---[ MAMBA-2 State Space Model ]---");
    bench_mamba2(iterations);
    println!();

    // =========================================================================
    // DQN BENCHMARKS
    // =========================================================================
    println!("---[ DQN Deep Q-Network ]---");
    bench_dqn(iterations);
    println!();

    // =========================================================================
    // PPO BENCHMARKS
    // =========================================================================
    println!("---[ PPO Proximal Policy Optimization ]---");
    bench_ppo(iterations);
    println!();

    // =========================================================================
    // TFT BENCHMARKS
    // =========================================================================
    println!("---[ TFT Temporal Fusion Transformer ]---");
    bench_tft(iterations);
    println!();

    // =========================================================================
    // BATCH INFERENCE
    // =========================================================================
    println!("---[ Batch Inference (100 samples) ]---");
    bench_batch(100);
    println!();

    // =========================================================================
    // COLD START
    // =========================================================================
    println!("---[ Cold Start (Load + Inference) ]---");
    bench_cold_start();
    println!();

    // =========================================================================
    // SUMMARY
    // =========================================================================
    println!("{}", "=".repeat(80));
    println!("VALIDATION SUMMARY");
    println!("{}", "=".repeat(80));
    println!("Target: <1ms p99 inference (warm cache)");
    println!("Target: <10s cold start");
    println!("Target: <50ms batch (100 samples)");
    println!("Target: >10x GPU speedup");
    println!();
}

fn bench_mamba2(iterations: usize) {
    let cpu_device = Device::Cpu;
    let shape = vec![1, 256, 512]; // batch, seq_len, d_model

    // CPU timing
    let mut cpu_stats = BenchStats::new();
    for _ in 0..iterations {
        let input = Tensor::randn(0.0f32, 1.0f32, shape.as_slice(), &cpu_device).unwrap();

        let start = Instant::now();
        let _output = input.matmul(&input.t().unwrap()).unwrap();
        cpu_stats.add(start.elapsed());
    }

    println!("CPU - Shape: {:?}", shape);
    println!("  Mean:   {:?}", cpu_stats.mean());
    println!("  P50:    {:?}", cpu_stats.percentile(50.0));
    println!("  P95:    {:?}", cpu_stats.percentile(95.0));
    println!("  P99:    {:?}", cpu_stats.percentile(99.0));
    println!("  Min/Max: {:?} / {:?}", cpu_stats.min(), cpu_stats.max());

    // GPU timing
    if let Ok(gpu_device) = Device::cuda_if_available(0) {
        let mut gpu_stats = BenchStats::new();
        for _ in 0..iterations {
            let input = Tensor::randn(0.0f32, 1.0f32, shape.as_slice(), &gpu_device).unwrap();

            let start = Instant::now();
            let _output = input.matmul(&input.t().unwrap()).unwrap();
            gpu_stats.add(start.elapsed());
        }

        println!("GPU - Shape: {:?}", shape);
        println!("  Mean:   {:?}", gpu_stats.mean());
        println!("  P50:    {:?}", gpu_stats.percentile(50.0));
        println!("  P95:    {:?}", gpu_stats.percentile(95.0));
        println!("  P99:    {:?}", gpu_stats.percentile(99.0));
        println!("  Min/Max: {:?} / {:?}", gpu_stats.min(), gpu_stats.max());

        let speedup = cpu_stats.mean().as_nanos() as f64 / gpu_stats.mean().as_nanos() as f64;
        println!("  Speedup: {:.2}x", speedup);
    }
}

fn bench_dqn(iterations: usize) {
    let cpu_device = Device::Cpu;
    let state_dim = 128;
    let action_dim = 16;

    // CPU timing
    let mut cpu_stats = BenchStats::new();
    for _ in 0..iterations {
        let input = Tensor::randn(0.0f32, 1.0f32, &[1, state_dim], &cpu_device).unwrap();

        let start = Instant::now();
        // 3-layer MLP
        let h1 = input.matmul(&Tensor::randn(0.0f32, 1.0f32, &[state_dim, 256], &cpu_device).unwrap()).unwrap();
        let h1_relu = h1.relu().unwrap();
        let h2 = h1_relu.matmul(&Tensor::randn(0.0f32, 1.0f32, &[256, 128], &cpu_device).unwrap()).unwrap();
        let h2_relu = h2.relu().unwrap();
        let _output = h2_relu.matmul(&Tensor::randn(0.0f32, 1.0f32, &[128, action_dim], &cpu_device).unwrap()).unwrap();
        cpu_stats.add(start.elapsed());
    }

    println!("CPU - State: {}, Actions: {}", state_dim, action_dim);
    println!("  Mean:   {:?}", cpu_stats.mean());
    println!("  P50:    {:?}", cpu_stats.percentile(50.0));
    println!("  P95:    {:?}", cpu_stats.percentile(95.0));
    println!("  P99:    {:?}", cpu_stats.percentile(99.0));
    println!("  Min/Max: {:?} / {:?}", cpu_stats.min(), cpu_stats.max());

    // GPU timing
    if let Ok(gpu_device) = Device::cuda_if_available(0) {
        let mut gpu_stats = BenchStats::new();
        for _ in 0..iterations {
            let input = Tensor::randn(0.0f32, 1.0f32, &[1, state_dim], &gpu_device).unwrap();

            let start = Instant::now();
            let h1 = input.matmul(&Tensor::randn(0.0f32, 1.0f32, &[state_dim, 256], &gpu_device).unwrap()).unwrap();
            let h1_relu = h1.relu().unwrap();
            let h2 = h1_relu.matmul(&Tensor::randn(0.0f32, 1.0f32, &[256, 128], &gpu_device).unwrap()).unwrap();
            let h2_relu = h2.relu().unwrap();
            let _output = h2_relu.matmul(&Tensor::randn(0.0f32, 1.0f32, &[128, action_dim], &gpu_device).unwrap()).unwrap();
            gpu_stats.add(start.elapsed());
        }

        println!("GPU - State: {}, Actions: {}", state_dim, action_dim);
        println!("  Mean:   {:?}", gpu_stats.mean());
        println!("  P50:    {:?}", gpu_stats.percentile(50.0));
        println!("  P95:    {:?}", gpu_stats.percentile(95.0));
        println!("  P99:    {:?}", gpu_stats.percentile(99.0));
        println!("  Min/Max: {:?} / {:?}", gpu_stats.min(), gpu_stats.max());

        let speedup = cpu_stats.mean().as_nanos() as f64 / gpu_stats.mean().as_nanos() as f64;
        println!("  Speedup: {:.2}x", speedup);
    }
}

fn bench_ppo(iterations: usize) {
    let cpu_device = Device::Cpu;
    let state_dim = 64;

    // CPU timing
    let mut cpu_stats = BenchStats::new();
    for _ in 0..iterations {
        let input = Tensor::randn(0.0f32, 1.0f32, &[1, state_dim], &cpu_device).unwrap();

        let start = Instant::now();
        let h1 = input.matmul(&Tensor::randn(0.0f32, 1.0f32, &[state_dim, 64], &cpu_device).unwrap()).unwrap();
        let h1_tanh = h1.tanh().unwrap();
        let _mean = h1_tanh.matmul(&Tensor::randn(0.0f32, 1.0f32, &[64, state_dim / 2], &cpu_device).unwrap()).unwrap();
        cpu_stats.add(start.elapsed());
    }

    println!("CPU - State: {}", state_dim);
    println!("  Mean:   {:?}", cpu_stats.mean());
    println!("  P50:    {:?}", cpu_stats.percentile(50.0));
    println!("  P95:    {:?}", cpu_stats.percentile(95.0));
    println!("  P99:    {:?}", cpu_stats.percentile(99.0));
    println!("  Min/Max: {:?} / {:?}", cpu_stats.min(), cpu_stats.max());

    // GPU timing
    if let Ok(gpu_device) = Device::cuda_if_available(0) {
        let mut gpu_stats = BenchStats::new();
        for _ in 0..iterations {
            let input = Tensor::randn(0.0f32, 1.0f32, &[1, state_dim], &gpu_device).unwrap();

            let start = Instant::now();
            let h1 = input.matmul(&Tensor::randn(0.0f32, 1.0f32, &[state_dim, 64], &gpu_device).unwrap()).unwrap();
            let h1_tanh = h1.tanh().unwrap();
            let _mean = h1_tanh.matmul(&Tensor::randn(0.0f32, 1.0f32, &[64, state_dim / 2], &gpu_device).unwrap()).unwrap();
            gpu_stats.add(start.elapsed());
        }

        println!("GPU - State: {}", state_dim);
        println!("  Mean:   {:?}", gpu_stats.mean());
        println!("  P50:    {:?}", gpu_stats.percentile(50.0));
        println!("  P95:    {:?}", gpu_stats.percentile(95.0));
        println!("  P99:    {:?}", gpu_stats.percentile(99.0));
        println!("  Min/Max: {:?} / {:?}", gpu_stats.min(), gpu_stats.max());

        let speedup = cpu_stats.mean().as_nanos() as f64 / gpu_stats.mean().as_nanos() as f64;
        println!("  Speedup: {:.2}x", speedup);
    }
}

fn bench_tft(iterations: usize) {
    let cpu_device = Device::Cpu;
    let shape = vec![1, 64, 128]; // batch, seq_len, features

    // CPU timing
    let mut cpu_stats = BenchStats::new();
    for _ in 0..iterations {
        let input = Tensor::randn(0.0f32, 1.0f32, shape.as_slice(), &cpu_device).unwrap();
        let features = shape[2];

        let start = Instant::now();
        let qkv = input.matmul(&Tensor::randn(0.0f32, 1.0f32, &[features, features * 3], &cpu_device).unwrap()).unwrap();
        let attention = qkv.matmul(&qkv.t().unwrap()).unwrap();
        let _output = softmax(&attention, 1).unwrap().matmul(&input).unwrap();
        cpu_stats.add(start.elapsed());
    }

    println!("CPU - Shape: {:?}", shape);
    println!("  Mean:   {:?}", cpu_stats.mean());
    println!("  P50:    {:?}", cpu_stats.percentile(50.0));
    println!("  P95:    {:?}", cpu_stats.percentile(95.0));
    println!("  P99:    {:?}", cpu_stats.percentile(99.0));
    println!("  Min/Max: {:?} / {:?}", cpu_stats.min(), cpu_stats.max());

    // GPU timing
    if let Ok(gpu_device) = Device::cuda_if_available(0) {
        let mut gpu_stats = BenchStats::new();
        for _ in 0..iterations {
            let input = Tensor::randn(0.0f32, 1.0f32, shape.as_slice(), &gpu_device).unwrap();
            let features = shape[2];

            let start = Instant::now();
            let qkv = input.matmul(&Tensor::randn(0.0f32, 1.0f32, &[features, features * 3], &gpu_device).unwrap()).unwrap();
            let attention = qkv.matmul(&qkv.t().unwrap()).unwrap();
            let _output = softmax(&attention, 1).unwrap().matmul(&input).unwrap();
            gpu_stats.add(start.elapsed());
        }

        println!("GPU - Shape: {:?}", shape);
        println!("  Mean:   {:?}", gpu_stats.mean());
        println!("  P50:    {:?}", gpu_stats.percentile(50.0));
        println!("  P95:    {:?}", gpu_stats.percentile(95.0));
        println!("  P99:    {:?}", gpu_stats.percentile(99.0));
        println!("  Min/Max: {:?} / {:?}", gpu_stats.min(), gpu_stats.max());

        let speedup = cpu_stats.mean().as_nanos() as f64 / gpu_stats.mean().as_nanos() as f64;
        println!("  Speedup: {:.2}x", speedup);
    }
}

fn bench_batch(batch_size: usize) {
    let cpu_device = Device::Cpu;
    let shape = vec![1, 128];

    // CPU timing
    let start = Instant::now();
    for _ in 0..batch_size {
        let input = Tensor::randn(0.0f32, 1.0f32, shape.as_slice(), &cpu_device).unwrap();
        let _output = input.matmul(&Tensor::randn(0.0f32, 1.0f32, &[128, 64], &cpu_device).unwrap()).unwrap();
    }
    let cpu_time = start.elapsed();

    println!("CPU - Batch Size: {}", batch_size);
    println!("  Total Time: {:?}", cpu_time);
    println!("  Per Sample: {:?}", cpu_time / batch_size as u32);

    // GPU timing
    if let Ok(gpu_device) = Device::cuda_if_available(0) {
        let start = Instant::now();
        for _ in 0..batch_size {
            let input = Tensor::randn(0.0f32, 1.0f32, shape.as_slice(), &gpu_device).unwrap();
            let _output = input.matmul(&Tensor::randn(0.0f32, 1.0f32, &[128, 64], &gpu_device).unwrap()).unwrap();
        }
        let gpu_time = start.elapsed();

        println!("GPU - Batch Size: {}", batch_size);
        println!("  Total Time: {:?}", gpu_time);
        println!("  Per Sample: {:?}", gpu_time / batch_size as u32);

        let speedup = cpu_time.as_nanos() as f64 / gpu_time.as_nanos() as f64;
        println!("  Speedup: {:.2}x", speedup);
    }
}

fn bench_cold_start() {
    let device = Device::cuda_if_available(0).unwrap_or(Device::Cpu);

    let start = Instant::now();
    // Simulate loading large model weights
    let _weights = Tensor::randn(0.0f32, 1.0f32, &[1000, 1000], &device).unwrap();

    // First inference
    let input = Tensor::randn(0.0f32, 1.0f32, &[1, 1000], &device).unwrap();
    let _output = input.matmul(&_weights).unwrap();
    let cold_start_time = start.elapsed();

    println!("Device: {:?}", device);
    println!("  Cold Start Time: {:?}", cold_start_time);
}