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 harness: load a labeled EEG eval set, run the
//! fine-tuned LuMamba classifier per epoch, optionally aggregate epoch
//! predictions to the recording level, and score with the paper's metrics.
//!
//! ## Eval-set format
//!
//! A JSON **manifest** plus a **signals** safetensors:
//!
//! ```jsonc
//! // manifest.json
//! {
//!   "task": "TUAB",
//!   "num_classes": 2,
//!   "metric": "balanced_accuracy",   // | "auroc" | "aupr"
//!   "positive_class": 1,              // for binary auroc / aupr
//!   "aggregate": "mean",              // "mean" (recording-level) | "none"
//!   "signals_path": "signals.safetensors",
//!   "sample_rate": 256.0,
//!   "labels": [0, 1, 1, 0, ...],          // length N (per epoch)
//!   "recording_ids": [0, 0, 1, 1, ...],   // optional, length N
//!   "expected": 0.8099                    // optional published number
//! }
//! ```
//!
//! `signals.safetensors` holds `signals` `[N, C, T]` (f32), `chan_pos`
//! `[C, 3]` (f32), and optionally `channel_idx` `[C]` (i64). Build one with
//! `scripts/make_eval_set.py`.

use std::path::Path;

use anyhow::{Context, Result};
use safetensors::SafeTensors;

use crate::metrics;
use crate::rlx::LuMambaEncoder;

/// Eval-set description (the manifest JSON).
#[derive(Debug, Clone, serde::Deserialize)]
pub struct EvalManifest {
    /// Task name, e.g. `"TUAB"` (informational; used in the report).
    pub task: String,
    /// Number of classes the classifier outputs.
    pub num_classes: usize,
    /// Metric to compute: `"balanced_accuracy"` | `"auroc"` | `"aupr"` | `"accuracy"`.
    pub metric: String,
    /// Positive class index for binary `auroc` / `aupr` (default 1).
    #[serde(default = "default_pos")]
    pub positive_class: usize,
    /// `"mean"` (recording-level mean of probabilities) | `"none"` (epoch-level).
    #[serde(default = "default_agg")]
    pub aggregate: String,
    /// Path to the signals safetensors, relative to the manifest's directory.
    pub signals_path: String,
    /// Sampling rate in Hz (informational).
    #[serde(default = "default_sr")]
    pub sample_rate: f32,
    /// Per-epoch integer labels, length `N`.
    pub labels: Vec<i64>,
    /// Optional per-epoch recording id, length `N` (for recording-level aggregation).
    #[serde(default)]
    pub recording_ids: Option<Vec<i64>>,
    /// Optional published metric value to compare against.
    #[serde(default)]
    pub expected: Option<f64>,
}

fn default_pos() -> usize { 1 }
fn default_agg() -> String { "none".to_string() }
fn default_sr() -> f32 { 256.0 }

/// EEG signals + channel geometry, parsed from the signals safetensors.
pub struct Signals {
    /// Epoch signals `[N, C, T]`, row-major.
    pub data: Vec<f32>,
    /// Number of epochs.
    pub n: usize,
    /// Number of channels.
    pub c: usize,
    /// Samples per epoch.
    pub t: usize,
    /// Channel positions `[C, 3]`, row-major.
    pub chan_pos: Vec<f32>,
    /// Optional channel vocabulary indices `[C]`.
    pub channel_idx: Option<Vec<i32>>,
}

/// Load the signals safetensors (`signals`, `chan_pos`, optional `channel_idx`).
pub fn load_signals(path: &Path) -> Result<Signals> {
    let bytes = std::fs::read(path).with_context(|| format!("reading {}", path.display()))?;
    let st = SafeTensors::deserialize(&bytes)?;

    let sig = st.tensor("signals").context("missing `signals` tensor")?;
    let shape = sig.shape().to_vec();
    anyhow::ensure!(shape.len() == 3, "`signals` must be [N, C, T], got {shape:?}");
    let (n, c, t) = (shape[0], shape[1], shape[2]);
    let data: Vec<f32> =
        sig.data().chunks_exact(4).map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])).collect();

    let cp = st.tensor("chan_pos").context("missing `chan_pos` tensor")?;
    anyhow::ensure!(cp.shape() == [c, 3], "`chan_pos` must be [C, 3]");
    let chan_pos: Vec<f32> =
        cp.data().chunks_exact(4).map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])).collect();

    let channel_idx = st.tensor("channel_idx").ok().map(|v| {
        v.data().chunks_exact(8).map(|b| {
            i64::from_le_bytes([b[0], b[1], b[2], b[3], b[4], b[5], b[6], b[7]]) as i32
        }).collect::<Vec<i32>>()
    });

    Ok(Signals { data, n, c, t, chan_pos, channel_idx })
}

/// Result of an evaluation run.
#[derive(Debug, Clone)]
pub struct EvalReport {
    /// Task name (from the manifest).
    pub task: String,
    /// Metric name that was computed.
    pub metric: String,
    /// Computed metric value.
    pub value: f64,
    /// Published value to compare against, if provided.
    pub expected: Option<f64>,
    /// Number of input epochs.
    pub n_epochs: usize,
    /// Number of scored units after aggregation (== `n_epochs` if epoch-level).
    pub n_units: usize,
    /// Confusion matrix `conf[true][pred]` for the accuracy-family metrics.
    pub confusion: Option<Vec<Vec<u64>>>,
}

impl EvalReport {
    /// Print the report to stdout (value, expected, Δ).
    pub fn print(&self) {
        println!("task        : {}", self.task);
        println!("epochs      : {}", self.n_epochs);
        println!("scored units: {} ({})", self.n_units,
            if self.n_units == self.n_epochs { "epoch-level" } else { "recording-level (mean prob)" });
        if let Some(cm) = &self.confusion {
            println!("confusion   : {cm:?}");
        }
        println!("{:<18}= {:.4}", self.metric, self.value);
        if let Some(e) = self.expected {
            println!("{:<18}= {:.4}  (published)", "expected", e);
            println!("{:<18}= {:+.4}", "Δ vs paper", self.value - e);
        }
    }
}

/// Mean-aggregate epoch probabilities to the recording level. Returns
/// `(unit_probs, unit_labels)`.
fn aggregate_mean(
    probs: &[Vec<f32>],
    labels: &[usize],
    rec_ids: &[i64],
) -> (Vec<Vec<f32>>, Vec<usize>) {
    use std::collections::BTreeMap;
    let k = probs.first().map(|p| p.len()).unwrap_or(0);
    let mut acc: BTreeMap<i64, (Vec<f64>, usize, usize)> = BTreeMap::new(); // id -> (sum, count, label)
    for ((p, &lab), &rid) in probs.iter().zip(labels).zip(rec_ids) {
        let e = acc.entry(rid).or_insert((vec![0.0; k], 0, lab));
        for (a, &v) in e.0.iter_mut().zip(p) {
            *a += v as f64;
        }
        e.1 += 1;
    }
    let mut up = Vec::new();
    let mut ul = Vec::new();
    for (_, (sum, count, lab)) in acc {
        up.push(sum.iter().map(|&s| (s / count as f64) as f32).collect());
        ul.push(lab);
    }
    (up, ul)
}

/// Run an evaluation: classify every epoch, aggregate, score.
pub fn run(model: &mut LuMambaEncoder, manifest: &EvalManifest, signals: &Signals) -> Result<EvalReport> {
    anyhow::ensure!(
        signals.n == manifest.labels.len(),
        "labels ({}) must match #epochs ({})",
        manifest.labels.len(),
        signals.n
    );
    let k = manifest.num_classes;
    let (c, t) = (signals.c, signals.t);

    let mut probs: Vec<Vec<f32>> = Vec::with_capacity(signals.n);
    for i in 0..signals.n {
        let sig = &signals.data[i * c * t..(i + 1) * c * t];
        let logits = model.classify_epoch(sig, &signals.chan_pos, signals.channel_idx.as_deref(), c, t)?;
        anyhow::ensure!(logits.len() == k, "classifier produced {} logits, expected {k}", logits.len());
        probs.push(metrics::softmax(&logits));
    }
    let labels: Vec<usize> = manifest.labels.iter().map(|&l| l as usize).collect();
    score(&probs, &labels, manifest)
}

/// Aggregate (if requested) and compute the manifest's metric from per-epoch
/// class probabilities. Pure — unit-tested without a model.
pub fn score(probs: &[Vec<f32>], labels: &[usize], manifest: &EvalManifest) -> Result<EvalReport> {
    let k = manifest.num_classes;
    let n_epochs = labels.len();

    let (probs, labels) = match (manifest.aggregate.as_str(), &manifest.recording_ids) {
        ("mean", Some(rids)) => {
            anyhow::ensure!(rids.len() == n_epochs, "recording_ids length mismatch");
            aggregate_mean(probs, labels, rids)
        }
        ("mean", None) => anyhow::bail!("aggregate=\"mean\" requires recording_ids in the manifest"),
        _ => (probs.to_vec(), labels.to_vec()),
    };

    let n_units = labels.len();
    let pos = manifest.positive_class;
    let (value, confusion) = match manifest.metric.as_str() {
        "balanced_accuracy" => {
            let preds: Vec<usize> = probs.iter().map(|p| metrics::argmax(p)).collect();
            (metrics::balanced_accuracy(&preds, &labels, k), Some(metrics::confusion(&preds, &labels, k)))
        }
        "accuracy" => {
            let preds: Vec<usize> = probs.iter().map(|p| metrics::argmax(p)).collect();
            (metrics::accuracy(&preds, &labels), Some(metrics::confusion(&preds, &labels, k)))
        }
        "auroc" => {
            if k == 2 {
                let scores: Vec<f64> = probs.iter().map(|p| p[pos] as f64).collect();
                let bin: Vec<u8> = labels.iter().map(|&l| (l == pos) as u8).collect();
                (metrics::auroc_binary(&scores, &bin), None)
            } else {
                (metrics::auroc_macro_ovr(&probs, &labels, k), None)
            }
        }
        "aupr" => {
            let scores: Vec<f64> = probs.iter().map(|p| p[pos] as f64).collect();
            let bin: Vec<u8> = labels.iter().map(|&l| (l == pos) as u8).collect();
            (metrics::average_precision(&scores, &bin), None)
        }
        other => anyhow::bail!("unknown metric: {other}"),
    };

    Ok(EvalReport {
        task: manifest.task.clone(),
        metric: manifest.metric.clone(),
        value,
        expected: manifest.expected,
        n_epochs,
        n_units,
        confusion,
    })
}

/// Load a manifest JSON from disk.
pub fn load_manifest(path: &Path) -> Result<EvalManifest> {
    let s = std::fs::read_to_string(path).with_context(|| format!("reading {}", path.display()))?;
    Ok(serde_json::from_str(&s)?)
}

#[cfg(test)]
mod tests {
    use super::*;

    fn manifest(metric: &str, aggregate: &str, rec: Option<Vec<i64>>) -> EvalManifest {
        EvalManifest {
            task: "T".into(),
            num_classes: 2,
            metric: metric.into(),
            positive_class: 1,
            aggregate: aggregate.into(),
            signals_path: String::new(),
            sample_rate: 256.0,
            labels: vec![],
            recording_ids: rec,
            expected: None,
        }
    }

    #[test]
    fn score_recording_level_balanced_accuracy() {
        // 4 epochs → 2 recordings. Per-epoch noisy, but recording-mean is correct.
        let probs = vec![
            vec![0.8f32, 0.2],
            vec![0.6, 0.4], // rec 0 → mean [0.7,0.3] → class 0
            vec![0.3, 0.7],
            vec![0.1, 0.9], // rec 1 → mean [0.2,0.8] → class 1
        ];
        let labels = vec![0usize, 0, 1, 1];
        let m = manifest("balanced_accuracy", "mean", Some(vec![0, 0, 1, 1]));
        let r = score(&probs, &labels, &m).unwrap();
        assert_eq!(r.n_units, 2);
        assert_eq!(r.n_epochs, 4);
        assert!((r.value - 1.0).abs() < 1e-9, "{}", r.value);
    }

    #[test]
    fn score_epoch_level_auroc() {
        let probs = vec![
            vec![0.9f32, 0.1],
            vec![0.6, 0.4],
            vec![0.35, 0.65],
            vec![0.2, 0.8],
        ];
        let labels = vec![0usize, 0, 1, 1];
        let m = manifest("auroc", "none", None);
        let r = score(&probs, &labels, &m).unwrap();
        // pos scores {0.65,0.8} both exceed neg {0.1,0.4} → AUROC 1.0
        assert!((r.value - 1.0).abs() < 1e-9, "{}", r.value);
    }
}