/// Benchmark SleepFM inference across modalities and output JSON results.
///
/// Usage:
///   cargo run --release --example benchmark -- \
///       --config data/config.json \
///       --weights data/model_base.safetensors \
///       --output data/benchmark_rust.json

use std::path::Path;
use std::time::Instant;
use burn::prelude::*;
use burn::tensor::TensorData;
use sleepfm::SleepFmEncoder;

// ── Backend ───────────────────────────────────────────────────────────────────
#[cfg(all(feature = "wgpu", not(feature = "ndarray")))]
mod backend {
    pub use burn::backend::{Wgpu as B, wgpu::WgpuDevice as Device};
    pub fn device() -> Device { Device::DefaultDevice }
    #[cfg(feature = "metal")]
    pub const NAME: &str = "GPU (wgpu — Metal)";
    #[cfg(feature = "vulkan")]
    pub const NAME: &str = "GPU (wgpu — Vulkan)";
    #[cfg(not(any(feature = "metal", feature = "vulkan")))]
    pub const NAME: &str = "GPU (wgpu — WGSL)";
}

#[cfg(feature = "ndarray")]
mod backend {
    pub use burn::backend::NdArray as B;
    pub type Device = burn::backend::ndarray::NdArrayDevice;
    pub fn device() -> Device { Device::Cpu }
    #[cfg(feature = "blas-accelerate")]
    pub const NAME: &str = "CPU (NdArray + Accelerate)";
    #[cfg(feature = "openblas-system")]
    pub const NAME: &str = "CPU (NdArray + OpenBLAS)";
    #[cfg(not(any(feature = "blas-accelerate", feature = "openblas-system")))]
    pub const NAME: &str = "CPU (NdArray + Rayon)";
}

use backend::{B, device};

fn main() -> anyhow::Result<()> {
    let args: Vec<String> = std::env::args().collect();
    let config_path = args.iter().position(|a| a == "--config").map(|i| &args[i + 1])
        .expect("--config <path> required");
    let weights_path = args.iter().position(|a| a == "--weights").map(|i| &args[i + 1])
        .expect("--weights <path> required");
    let output_path = args.iter().position(|a| a == "--output").map(|i| args[i + 1].as_str())
        .unwrap_or("data/benchmark_rust.json");

    eprintln!("Backend: {}", backend::NAME);
    let device = device();
    let (encoder, ms_load) = SleepFmEncoder::<B>::load(
        Path::new(config_path), Path::new(weights_path), device.clone(),
    )?;
    eprintln!("{}", encoder.describe());
    eprintln!("Model loaded in {ms_load:.1}ms\n");

    let warmup = 3;
    let runs = 10;
    let t_samples = encoder.model_cfg.samples_per_chunk();

    let modalities = vec![
        ("EKG", 2usize),
        ("EMG", 4),
        ("RESP", 7),
        ("BAS", 10),
    ];

    let mut all_results = serde_json::Map::new();

    for (name, n_channels) in &modalities {
        let signal = Tensor::<B, 3>::zeros([1, *n_channels, t_samples], &device);
        let mask_data: Vec<bool> = vec![false; *n_channels];
        let mask = Tensor::<B, 2, Bool>::from_data(
            TensorData::new(mask_data, vec![1, *n_channels]), &device,
        );

        // Warmup
        for _ in 0..warmup {
            let _ = encoder.run(signal.clone(), mask.clone())?;
        }

        // Timed runs
        let mut times = Vec::with_capacity(runs);
        for _ in 0..runs {
            let t = Instant::now();
            let _ = encoder.run(signal.clone(), mask.clone())?;
            times.push(t.elapsed().as_secs_f64() * 1000.0);
        }

        let mean: f64 = times.iter().sum::<f64>() / times.len() as f64;
        let variance: f64 = times.iter().map(|t| (t - mean).powi(2)).sum::<f64>() / times.len() as f64;
        let std = variance.sqrt();

        eprintln!("{name:>4} ({n_channels:>2}ch): {mean:.1} ± {std:.1} ms");

        let mut entry = serde_json::Map::new();
        entry.insert("n_channels".into(), serde_json::Value::from(*n_channels));
        entry.insert("mean_ms".into(), serde_json::Value::from(mean));
        entry.insert("std_ms".into(), serde_json::Value::from(std));
        entry.insert("times_ms".into(), serde_json::Value::from(times));
        all_results.insert(name.to_string(), serde_json::Value::Object(entry));
    }

    let json = serde_json::to_string_pretty(&serde_json::Value::Object(all_results))?;
    std::fs::write(output_path, &json)?;
    eprintln!("\nSaved to {output_path}");

    Ok(())
}
