lumamba 0.0.1

// lumamba-rs — LuMamba EEG foundation-model inference on the RLX runtime.
// Copyright (C) 2026 Nataliya Kosmyna.
// SPDX-License-Identifier: GPL-3.0-only

//! LuMamba EEG inference — thin CLI over [`lumamba::LuMambaEncoder`].

use std::time::Instant;

use clap::Parser;

use lumamba::init_threads;
use lumamba::rlx::{save_epochs, LuMambaEncoder, PreprocInfo, RlxEpoch};

#[derive(Parser, Debug)]
#[command(about = "LuMamba EEG foundation-model inference (RLX runtime)")]
struct Args {
    /// Backend: cpu, metal, mlx, gpu, cuda, rocm, tpu.
    #[arg(long, default_value = "cpu", value_parser = lumamba::parse_device)]
    device: rlx::Device,

    /// Weights path, or `hf://<owner>/<repo>/<file>` (needs `--features hf-download`).
    #[arg(long, env = "LUMAMBA_WEIGHTS")]
    weights: String,

    /// Config path, or `hf://<owner>/<repo>/<file>`.
    #[arg(long, env = "LUMAMBA_CONFIG")]
    config: String,

    /// Raw EEG recording (.edf or .fif). Requires the `edf` feature.
    /// Omit to run on a synthetic epoch.
    #[arg(long)]
    input: Option<String>,

    #[arg(long)]
    output: String,

    /// Emit the encoder latent (foundation embedding) instead of the
    /// reconstruction / classification head output.
    #[arg(long)]
    encode: bool,

    /// Repeat inference N times; reports cold (first) vs warm (cached-graph)
    /// per-epoch time. Use to measure steady-state GPU throughput.
    #[arg(long, default_value_t = 1)]
    repeat: usize,

    #[arg(long, env = "RAYON_NUM_THREADS")]
    threads: Option<usize>,

    #[arg(long, short = 'v')]
    verbose: bool,
}

fn main() -> anyhow::Result<()> {
    let args = Args::parse();
    let n_threads = init_threads(args.threads);
    let device = args.device;
    let t0 = Instant::now();

    eprintln!("Device   : {:?}  ({n_threads} threads)", device);

    let config_path = lumamba::hf::resolve(&args.config)?;
    let weights_path = lumamba::hf::resolve(&args.weights)?;
    let (mut model, ms_load) = LuMambaEncoder::load(&config_path, &weights_path, device)?;
    eprintln!("Model    : {}  ({ms_load:.0} ms)", model.describe());

    let t_pp = Instant::now();
    let (epochs, info) = load_epochs(args.input.as_deref())?;
    let ms_preproc = t_pp.elapsed().as_secs_f64() * 1000.0;
    eprintln!(
        "Input    : {} epochs × {} ch × {} samples",
        epochs.len(),
        info.n_channels,
        epochs.first().map(|e| e.n_samples).unwrap_or(0),
    );
    eprintln!("Preproc  : {ms_preproc:.1} ms");

    // Run `repeat` passes; the first pass pays one-time graph compilation
    // (cold), later passes reuse the cached compiled graph (warm steady-state).
    let runs = args.repeat.max(1);
    let mut outputs = Vec::with_capacity(epochs.len());
    let mut times = Vec::with_capacity(runs);
    for _ in 0..runs {
        outputs.clear();
        let t = Instant::now();
        for ep in &epochs {
            if args.encode {
                let (latent, shape) = model.encode(
                    &ep.signal,
                    &ep.chan_pos,
                    ep.channel_indices.as_deref(),
                    ep.n_channels,
                    ep.n_samples,
                )?;
                outputs.push(lumamba::rlx::EpochEmbedding {
                    output: latent,
                    shape,
                    chan_pos: ep.chan_pos.clone(),
                    n_channels: ep.n_channels,
                });
            } else {
                outputs.push(model.run_rlx_epoch(ep)?);
            }
        }
        times.push(t.elapsed().as_secs_f64() * 1000.0);
    }
    let n = outputs.len().max(1);
    let cold = times[0];
    if runs > 1 {
        let warm: f64 = times[1..].iter().sum::<f64>() / (runs - 1) as f64;
        eprintln!(
            "Infer    : cold {cold:.0} ms  |  warm {:.1} ms/epoch  ({n} epochs × {} warm runs)",
            warm / n as f64,
            runs - 1,
        );
    } else {
        eprintln!("Infer    : {cold:.0} ms  ({n} epochs; {:.0} ms/epoch)", cold / n as f64);
    }

    if args.verbose {
        for (i, ep) in outputs.iter().enumerate() {
            let n = ep.output.len().max(1);
            let mean: f64 = ep.output.iter().map(|&v| v as f64).sum::<f64>() / n as f64;
            let max = ep.output.iter().cloned().fold(f32::MIN, f32::max);
            let min = ep.output.iter().cloned().fold(f32::MAX, f32::min);
            eprintln!("  [ep {}] shape={:?}  mean={mean:+.5}  min={min:+.4} max={max:+.4}", i + 1, ep.shape);
        }
    }

    save_epochs(&outputs, &args.output)?;
    eprintln!(
        "Output   → {}  ({:.0} ms total)",
        args.output,
        t0.elapsed().as_secs_f64() * 1000.0
    );
    Ok(())
}

fn load_epochs(input: Option<&str>) -> anyhow::Result<(Vec<RlxEpoch>, PreprocInfo)> {
    match input {
        Some(path) => load_recording(path),
        None => synthetic_epoch(),
    }
}

#[cfg(feature = "edf")]
fn load_recording(path: &str) -> anyhow::Result<(Vec<RlxEpoch>, PreprocInfo)> {
    use lumamba::rlx::{load_edf, load_fif};
    let p = std::path::Path::new(path);
    if p.extension().and_then(|s| s.to_str()) == Some("edf") {
        load_edf(p)
    } else {
        load_fif(p)
    }
}

#[cfg(not(feature = "edf"))]
fn load_recording(_path: &str) -> anyhow::Result<(Vec<RlxEpoch>, PreprocInfo)> {
    anyhow::bail!("EDF/FIF input requires building with `--features edf`")
}

/// One synthetic epoch on a 22-channel TUEG bipolar montage (128 patches).
fn synthetic_epoch() -> anyhow::Result<(Vec<RlxEpoch>, PreprocInfo)> {
    use lumamba::channel_positions;

    let n_channels = 22usize;
    let n_samples = 5120usize; // 128 patches × patch_size 40
    let names = [
        "FP1-F7", "F7-T3", "T3-T5", "T5-O1", "FP2-F8", "F8-T4", "T4-T6", "T6-O2", "T3-C3", "C3-CZ",
        "CZ-C4", "C4-T4", "FP1-F3", "F3-C3", "C3-P3", "P3-O1", "FP2-F4", "F4-C4", "C4-P4", "P4-O2",
        "A1-T3", "T4-A2",
    ];
    let positions: Vec<f32> = names
        .iter()
        .flat_map(|n| channel_positions::bipolar_channel_xyz(n).unwrap_or([0.0, 0.0, 0.0]).to_vec())
        .collect();
    let channel_indices: Vec<i32> =
        names.iter().map(|n| lumamba::channel_index(n).unwrap_or(0) as i32).collect();

    // Mild deterministic sinusoid so the smoke test isn't all-zeros.
    let mut signal = vec![0f32; n_channels * n_samples];
    for ch in 0..n_channels {
        for ti in 0..n_samples {
            let f = 0.05 + 0.01 * ch as f32;
            signal[ch * n_samples + ti] = (f * ti as f32).sin();
        }
    }

    Ok((
        vec![RlxEpoch {
            signal,
            chan_pos: positions,
            channel_indices: Some(channel_indices),
            n_channels,
            n_samples,
        }],
        PreprocInfo {
            ch_names: names.iter().map(|s| s.to_string()).collect(),
            n_channels,
            n_epochs: 1,
            target_sfreq: 256.0,
            epoch_dur: 20.0,
        },
    ))
}