/// SleepFM PSG inference — thin CLI.
///
/// Build — CPU (default):
///   cargo build --release
///
/// Build — GPU:
///   cargo build --release --no-default-features --features wgpu
///
/// Usage:
///   infer --weights <st> --config <json>

use std::path::Path;
use std::time::Instant;
use clap::Parser;
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 / MSL shaders)";
    #[cfg(feature = "vulkan")]
    pub const NAME: &str = "GPU (wgpu — Vulkan / SPIR-V shaders)";
    #[cfg(not(any(feature = "metal", feature = "vulkan")))]
    pub const NAME: &str = "GPU (wgpu — WGSL shaders)";
}

#[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 + Apple 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};

// ── CLI ───────────────────────────────────────────────────────────────────────
#[derive(Parser, Debug)]
#[command(about = "SleepFM PSG model inference (Burn 0.20.1)")]
struct Args {
    /// Safetensors weights file.
    #[arg(long)]
    weights: String,

    /// Model config JSON file.
    #[arg(long)]
    config: String,

    /// Number of channels (simulated input).
    #[arg(long, default_value = "10")]
    channels: usize,
}

fn main() -> anyhow::Result<()> {
    let args = Args::parse();
    eprintln!("Backend: {}", backend::NAME);
    let device = device();

    // Load model
    let (encoder, ms_load) = SleepFmEncoder::<B>::load(
        Path::new(&args.config),
        Path::new(&args.weights),
        device.clone(),
    )?;
    eprintln!("{}", encoder.describe());
    eprintln!("Model loaded in {ms_load:.1}ms");

    // Generate synthetic input
    let cfg = &encoder.model_cfg;
    let t_samples = cfg.samples_per_chunk();
    let n_channels = args.channels;

    eprintln!("Running inference: {} channels × {} samples ({} patches)",
        n_channels, t_samples, cfg.num_patches());

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

    let t = Instant::now();
    let result = encoder.run(signal, mask)?;
    let ms_infer = t.elapsed().as_secs_f64() * 1000.0;

    eprintln!("Inference done in {ms_infer:.1}ms");
    eprintln!("Pooled embedding: {} dims", result.pooled.len());
    eprintln!("Patch embeddings: {:?}", result.patches_shape);
    eprintln!("First 5 pooled values: {:?}", &result.pooled[..5.min(result.pooled.len())]);

    Ok(())
}
