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

//! Validation helper: run the FEMBA temporal encoder on an externally
//! supplied latent and dump the result, so a reference implementation can
//! compare against a byte-identical input.
//!
//! Usage:
//!   `validate_encoder --config config.json --weights CKPT.safetensors
//!   --hin h_in.safetensors --out h_out.safetensors`
//!
//! `h_in.safetensors` must contain an F32 tensor named `h_in` of shape
//! `[1, S, hidden]`. Writes `h_out` (same shape) to `--out`.

use std::borrow::Cow;
use std::path::Path;

use clap::Parser;
use safetensors::{Dtype, SafeTensors, View};

use lumamba::LuMambaEncoder;

#[derive(Parser, Debug)]
struct Args {
    #[arg(long)]
    config: String,
    #[arg(long)]
    weights: String,
    #[arg(long)]
    hin: String,
    #[arg(long)]
    out: String,
    /// Dump per-stage intermediates of mamba_blocks.0.mamba_fwd instead of the
    /// full encoder output.
    #[arg(long)]
    debug: bool,
    /// Dump bidirectional-wrapper stages of block 0.
    #[arg(long)]
    debug2: bool,
}

const DBG_STAGES: [&str; 9] =
    ["x_ssm", "x_conv", "x_act", "B", "C", "delta", "scan_y", "y_gated", "out"];
const DBG2_STAGES: [&str; 4] = ["flip_in", "fwd", "rev_out", "wrapper"];

struct Raw {
    data: Vec<u8>,
    shape: Vec<usize>,
}
impl View for Raw {
    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() }
}

fn main() -> anyhow::Result<()> {
    let args = Args::parse();

    let (mut model, _) =
        LuMambaEncoder::load(Path::new(&args.config), Path::new(&args.weights), rlx::Device::Cpu)?;

    let bytes = std::fs::read(&args.hin)?;
    let st = SafeTensors::deserialize(&bytes)?;
    let view = st.tensor("h_in")?;
    let shape = view.shape().to_vec();
    anyhow::ensure!(shape.len() == 3 && shape[0] == 1, "h_in must be [1, S, hidden], got {shape:?}");
    let s = shape[1];
    let h_in: Vec<f32> =
        view.data().chunks_exact(4).map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])).collect();

    let hidden = model.model_cfg.hidden_dim();
    let di = model.model_cfg.d_inner();
    let n = model.model_cfg.d_state;

    if args.debug {
        let stages = model.run_encoder_debug(&h_in, s)?;
        anyhow::ensure!(stages.len() == DBG_STAGES.len(), "expected {} stages", DBG_STAGES.len());
        // per-stage trailing dim, to write correct shapes
        let last_dim = |name: &str| -> usize {
            match name {
                "B" | "C" => n,
                "out" => hidden,
                _ => di,
            }
        };
        let mut tensors: Vec<(String, Raw)> = Vec::new();
        for (name, data) in DBG_STAGES.iter().zip(stages.iter()) {
            let ld = last_dim(name);
            tensors.push((
                name.to_string(),
                Raw { data: data.iter().flat_map(|f| f.to_le_bytes()).collect(), shape: vec![1, s, ld] },
            ));
        }
        std::fs::write(&args.out, safetensors::serialize(tensors, None)?)?;
        eprintln!("wrote {} debug stages → {}", stages.len(), args.out);
        return Ok(());
    }

    if args.debug2 {
        let stages = model.run_encoder_debug2(&h_in, s)?;
        let mut tensors: Vec<(String, Raw)> = Vec::new();
        for (name, data) in DBG2_STAGES.iter().zip(stages.iter()) {
            tensors.push((
                name.to_string(),
                Raw { data: data.iter().flat_map(|f| f.to_le_bytes()).collect(), shape: vec![1, s, hidden] },
            ));
        }
        std::fs::write(&args.out, safetensors::serialize(tensors, None)?)?;
        eprintln!("wrote {} debug2 stages → {}", stages.len(), args.out);
        return Ok(());
    }

    let h_out = model.run_encoder(&h_in, s)?;

    let out_bytes: Vec<u8> = h_out.iter().flat_map(|f| f.to_le_bytes()).collect();
    let raw = Raw { data: out_bytes, shape };
    std::fs::write(&args.out, safetensors::serialize([("h_out", raw)], None)?)?;
    eprintln!("wrote h_out ({} elems) → {}", h_out.len(), args.out);
    Ok(())
}