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

//! Downstream evaluation: score a fine-tuned LuMamba classifier on a labeled
//! EEG eval set and compare against the published number.
//!
//! Usage:
//!   eval --config <finetune_config.json> --weights <finetuned.safetensors> \
//!        --manifest <eval_set/manifest.json> [--device cpu]
//!
//! See `lumamba::eval` for the eval-set format, and `scripts/make_eval_set.py`
//! to build one from preprocessed signals + labels.

use std::path::Path;

use clap::Parser;

use lumamba::eval::{load_manifest, load_signals, run};
use lumamba::init_threads;
use lumamba::LuMambaEncoder;

#[derive(Parser, Debug)]
#[command(about = "Evaluate a fine-tuned LuMamba classifier on a labeled EEG eval set")]
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,

    /// Path to the eval-set manifest JSON. (`signals_path` inside it is
    /// resolved relative to the manifest's directory.)
    #[arg(long)]
    manifest: String,

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

fn main() -> anyhow::Result<()> {
    let args = Args::parse();
    let n_threads = init_threads(args.threads);
    let device = args.device;
    eprintln!("Device   : {device:?}  ({n_threads} threads)");

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

    let manifest = load_manifest(Path::new(&args.manifest))?;
    // Resolve signals_path relative to the manifest's directory.
    let base = Path::new(&args.manifest).parent().unwrap_or_else(|| Path::new("."));
    let signals_path = base.join(&manifest.signals_path);
    let signals = load_signals(&signals_path)?;
    eprintln!(
        "Eval set : {} task, {} epochs × {} ch × {} samples, metric={}",
        manifest.task, signals.n, signals.c, signals.t, manifest.metric
    );

    let report = run(&mut model, &manifest, &signals)?;
    println!("───────────────────────────────────────────");
    report.print();
    Ok(())
}