#[cfg(feature = "cuda")]
use std::time::Instant;

#[cfg(feature = "cuda")]
use narsil_engine::backend::mlp::deterministic_input;
#[cfg(feature = "cuda")]
use narsil_engine::{BurnMlpModel, MlpConfig};

#[cfg(feature = "cuda")]
fn parse_arg(args: &[String], name: &str, default: usize) -> narsil_engine::Result<usize> {
    let Some(index) = args.iter().position(|arg| arg == name) else {
        return Ok(default);
    };
    let value = args.get(index + 1).ok_or_else(|| {
        narsil_engine::NarsilError::InvalidRequest(format!("missing value for {name}"))
    })?;
    value.parse().map_err(|err| {
        narsil_engine::NarsilError::InvalidRequest(format!("invalid {name} value '{value}': {err}"))
    })
}

#[cfg(feature = "cuda")]
fn main() -> narsil_engine::Result<()> {
    let args = std::env::args().collect::<Vec<_>>();
    let batch = parse_arg(&args, "--batch", 1)?;
    let input_dim = parse_arg(&args, "--input", 1024)?;
    let hidden_dim = parse_arg(&args, "--hidden", 2048)?;
    let output_dim = parse_arg(&args, "--output", 1000)?;
    let warmup = parse_arg(&args, "--warmup", 10)?;
    let iters = parse_arg(&args, "--iters", 100)?;
    let device_index = parse_arg(&args, "--device", 0)?;

    let config = MlpConfig::new(input_dim, hidden_dim, output_dim)?;
    let device = burn_cuda::CudaDevice::new(device_index);
    let model = BurnMlpModel::<burn_cuda::Cuda<f32, i32>>::new(config, device)?;
    let input = deterministic_input(batch, input_dim);

    for _ in 0..warmup {
        let output = model.infer_batch_f32(input.clone(), batch)?;
        std::hint::black_box(output);
    }

    let started = Instant::now();
    let mut checksum = 0.0_f32;
    for _ in 0..iters {
        let output = model.infer_batch_f32(input.clone(), batch)?;
        checksum += output.iter().take(16).sum::<f32>();
        std::hint::black_box(&output);
    }
    let elapsed = started.elapsed();
    let avg_ms = elapsed.as_secs_f64() * 1000.0 / iters as f64;
    let rps = iters as f64 / elapsed.as_secs_f64();

    println!("backend=burn-cuda");
    println!("device={device_index}");
    println!("batch={batch}");
    println!("input_dim={input_dim}");
    println!("hidden_dim={hidden_dim}");
    println!("output_dim={output_dim}");
    println!("warmup={warmup}");
    println!("iters={iters}");
    println!("elapsed_s={:.6}", elapsed.as_secs_f64());
    println!("avg_ms={avg_ms:.3}");
    println!("rps={rps:.1}");
    println!("checksum={checksum:.6}");

    Ok(())
}

#[cfg(not(feature = "cuda"))]
fn main() {
    eprintln!("This example requires: cargo run --release --features cuda --example burn_cuda_mlp");
    std::process::exit(1);
}
