/// Example: load SleepFM model and generate embeddings for synthetic data.
///
/// Usage:
///   cargo run --example embed -- \
///       --config data/config.json \
///       --weights data/model_base.safetensors

use std::path::Path;
use burn::backend::NdArray as B;
use burn::prelude::*;
use burn::tensor::TensorData;
use sleepfm::SleepFmEncoder;

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 device = burn::backend::ndarray::NdArrayDevice::Cpu;

    // Load model
    let (encoder, ms) = SleepFmEncoder::<B>::load(
        Path::new(config_path),
        Path::new(weights_path),
        device.clone(),
    )?;
    eprintln!("{}", encoder.describe());
    eprintln!("Loaded in {ms:.1}ms");

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

    let t_samples = encoder.model_cfg.samples_per_chunk();

    for (name, n_channels) in &modalities {
        // Create random input
        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,
        );

        let result = encoder.run(signal, mask)?;
        println!("{name:>4}: pooled [{:>3}]  patches {:?}",
            result.pooled.len(), result.patches_shape);
    }

    Ok(())
}
