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

//! RLX-side epoch I/O. The EDF/FIF loaders depend on `exg` + `exg-luna` and
//! are gated behind the `edf` feature; `save_epochs` and the epoch/metadata
//! types are always available.

use anyhow::Result;

/// One preprocessed epoch ready for [`super::LuMambaEncoder::run_epoch`].
#[derive(Clone, Debug)]
pub struct RlxEpoch {
    /// EEG signal `[C, T]` row-major (one batch item).
    pub signal: Vec<f32>,
    /// Channel positions `[C, 3]` row-major.
    pub chan_pos: Vec<f32>,
    /// Vocabulary indices `[C]` when every channel name is known.
    pub channel_indices: Option<Vec<i32>>,
    /// Number of channels.
    pub n_channels: usize,
    /// Samples per epoch.
    pub n_samples: usize,
}

/// Metadata from a preprocessing run.
#[derive(Clone, Debug)]
pub struct PreprocInfo {
    /// Channel names in order.
    pub ch_names: Vec<String>,
    /// Number of channels.
    pub n_channels: usize,
    /// Number of epochs produced.
    pub n_epochs: usize,
    /// Target sampling rate (Hz) after resampling.
    pub target_sfreq: f32,
    /// Epoch duration in seconds.
    pub epoch_dur: f32,
}

#[cfg(feature = "edf")]
mod load {
    use std::path::Path;

    use anyhow::{Context, Result};
    use exg_luna::LunaPipelineConfig;

    use super::{PreprocInfo, RlxEpoch};
    use crate::channel_positions::bipolar_channel_xyz;
    use crate::channel_vocab;

    fn epochs_to_rlx(
        epochs: Vec<(ndarray::Array2<f32>, Vec<String>)>,
    ) -> Result<(Vec<RlxEpoch>, PreprocInfo)> {
        if epochs.is_empty() {
            anyhow::bail!("preprocessing produced zero epochs");
        }
        let ch_names = epochs[0].1.clone();
        let n_channels = ch_names.len();
        let n_samples = epochs[0].0.ncols();

        let positions: Vec<f32> = ch_names
            .iter()
            .flat_map(|name| bipolar_channel_xyz(name).unwrap_or([0.0, 0.0, 0.0]).to_vec())
            .collect();

        let channel_indices: Option<Vec<i32>> = {
            let indices: Vec<Option<usize>> =
                ch_names.iter().map(|n| channel_vocab::channel_index(n)).collect();
            if indices.iter().all(|i| i.is_some()) {
                Some(indices.iter().map(|i| i.unwrap() as i32).collect())
            } else {
                None
            }
        };

        let n_epochs = epochs.len();
        let mut out = Vec::with_capacity(n_epochs);
        for (epoch_data, _) in &epochs {
            out.push(RlxEpoch {
                signal: epoch_data.iter().copied().collect(),
                chan_pos: positions.clone(),
                channel_indices: channel_indices.clone(),
                n_channels,
                n_samples,
            });
        }
        Ok((
            out,
            PreprocInfo { ch_names, n_channels, n_epochs, target_sfreq: 256.0, epoch_dur: 5.0 },
        ))
    }

    /// Load and preprocess an EDF file into RLX epochs.
    pub fn load_edf(path: &Path) -> Result<(Vec<RlxEpoch>, PreprocInfo)> {
        load_edf_with_config(path, &LunaPipelineConfig::default())
    }

    /// Load + preprocess an EDF file with a custom pipeline config.
    pub fn load_edf_with_config(
        path: &Path,
        cfg: &LunaPipelineConfig,
    ) -> Result<(Vec<RlxEpoch>, PreprocInfo)> {
        let raw = exg::edf::open_raw_edf(path)
            .with_context(|| format!("opening EDF: {}", path.display()))?;
        let data = raw
            .read_all_data()
            .with_context(|| format!("reading EDF data: {}", path.display()))?;
        let ch_names: Vec<String> = raw.channel_names();
        let sfreq = raw.header.sample_rate;
        let epochs = exg_luna::preprocess_luna(data, &ch_names, sfreq, cfg)
            .with_context(|| "LUNA preprocessing failed")?;
        epochs_to_rlx(epochs)
    }

    /// Load and preprocess a FIF file into RLX epochs.
    pub fn load_fif(path: &Path) -> Result<(Vec<RlxEpoch>, PreprocInfo)> {
        load_fif_with_config(path, &LunaPipelineConfig::default())
    }

    /// Load + preprocess a FIF file with a custom pipeline config.
    pub fn load_fif_with_config(
        path: &Path,
        cfg: &LunaPipelineConfig,
    ) -> Result<(Vec<RlxEpoch>, PreprocInfo)> {
        let raw = exg::fiff::raw::open_raw(path)
            .with_context(|| format!("opening FIF: {}", path.display()))?;
        let data = raw
            .read_all_data()
            .with_context(|| format!("reading FIF data: {}", path.display()))?;
        let ch_names: Vec<String> = raw.info.chs.iter().map(|ch| ch.name.clone()).collect();
        let sfreq = raw.info.sfreq as f32;
        let data_f32 = data.mapv(|v| v as f32);
        let epochs = exg_luna::preprocess_luna(data_f32, &ch_names, sfreq, cfg)
            .with_context(|| "LUNA preprocessing failed")?;
        epochs_to_rlx(epochs)
    }
}

#[cfg(feature = "edf")]
pub use load::{load_edf, load_edf_with_config, load_fif, load_fif_with_config};

/// Write RLX epoch outputs to a safetensors file (`output_{i}`, `chan_pos_{i}`).
pub fn save_epochs(epochs: &[super::EpochEmbedding], path: &str) -> Result<()> {
    use safetensors::{Dtype, View};
    use std::borrow::Cow;

    struct RawTensor {
        data: Vec<u8>,
        shape: Vec<usize>,
    }
    impl View for RawTensor {
        fn dtype(&self) -> Dtype {
            Dtype::F32
        }
        fn shape(&self) -> &[usize] {
            &self.shape
        }
        fn data(&self) -> Cow<'_, [u8]> {
            Cow::Borrowed(&self.data)
        }
        fn data_len(&self) -> usize {
            self.data.len()
        }
    }

    let f32_bytes = |v: &[f32]| -> Vec<u8> { v.iter().flat_map(|f| f.to_le_bytes()).collect() };

    let mut keys = Vec::new();
    let mut tensors = Vec::new();
    for (i, ep) in epochs.iter().enumerate() {
        keys.push(format!("output_{i}"));
        tensors.push(RawTensor { data: f32_bytes(&ep.output), shape: ep.shape.clone() });
        keys.push(format!("chan_pos_{i}"));
        tensors.push(RawTensor { data: f32_bytes(&ep.chan_pos), shape: vec![ep.n_channels, 3] });
    }
    keys.push("n_samples".into());
    tensors.push(RawTensor { data: f32_bytes(&[epochs.len() as f32]), shape: vec![1] });

    let pairs: Vec<(&str, RawTensor)> = keys.iter().map(|s| s.as_str()).zip(tensors).collect();
    std::fs::write(path, safetensors::serialize(pairs, None)?)?;
    Ok(())
}