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

//! Safetensors → flat parameter map for the LuMamba RLX graph.
//!
//! Linear weights are transposed from PyTorch `[out, in]` to RLX `[in, out]`
//! layout (so `g.mm(x, w)` matches `x @ wᵀ`). Fused `nn.MultiheadAttention`
//! `in_proj` tensors are split into separate Q/K/V. The Mamba `A_log`
//! parameter is materialised as `A = -exp(A_log)` at load time, and the
//! depthwise conv weight `[d_inner, 1, k]` is reshaped to grouped-`conv2d`
//! NCHW layout `[d_inner, 1, 1, k]`.

use std::collections::HashMap;

use half::bf16;
use safetensors::SafeTensors;

use crate::config::ModelConfig;

/// A named parameter tensor: row-major `f32` data plus its shape.
#[derive(Clone, Debug)]
pub struct ParamBuf {
    /// Row-major tensor data.
    pub data:  Vec<f32>,
    /// Tensor shape.
    pub shape: Vec<usize>,
}

/// Map from parameter name to [`ParamBuf`].
pub type ParamMap = HashMap<String, ParamBuf>;

/// Load a safetensors file into a name→[`ParamBuf`] map (BF16/F32 → f32,
/// stripping any leading `model.` prefix from keys).
pub fn load_safetensors(path: &str) -> anyhow::Result<HashMap<String, ParamBuf>> {
    let bytes = std::fs::read(path)?;
    let st = SafeTensors::deserialize(&bytes)?;
    let mut out = HashMap::with_capacity(st.len());
    for (raw_key, view) in st.tensors() {
        let key = raw_key
            .strip_prefix("model.")
            .unwrap_or(raw_key.as_str())
            .to_string();
        let shape: Vec<usize> = view.shape().to_vec();
        let data = match view.dtype() {
            safetensors::Dtype::BF16 => view
                .data()
                .chunks_exact(2)
                .map(|b| bf16::from_le_bytes([b[0], b[1]]).to_f32())
                .collect(),
            safetensors::Dtype::F32 => view
                .data()
                .chunks_exact(4)
                .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
                .collect(),
            other => anyhow::bail!("unsupported dtype {:?} for key {}", other, key),
        };
        out.insert(key, ParamBuf { data, shape });
    }
    Ok(out)
}

fn transpose(data: &[f32], rows: usize, cols: usize) -> Vec<f32> {
    let mut out = vec![0f32; data.len()];
    for r in 0..rows {
        for c in 0..cols {
            out[c * rows + r] = data[r * cols + c];
        }
    }
    out
}

/// Take a PyTorch Linear weight `[out, in]` and transpose to RLX `[in, out]`.
fn take_linear_w(raw: &mut HashMap<String, ParamBuf>, key: &str) -> anyhow::Result<ParamBuf> {
    let p = raw
        .remove(key)
        .ok_or_else(|| anyhow::anyhow!("missing weight key: {key}"))?;
    anyhow::ensure!(
        p.shape.len() == 2,
        "Linear weight {key} must be 2-D, got {:?}",
        p.shape
    );
    let (out_d, in_d) = (p.shape[0], p.shape[1]);
    let data = transpose(&p.data, out_d, in_d);
    Ok(ParamBuf {
        data,
        shape: vec![in_d, out_d],
    })
}

fn take(raw: &mut HashMap<String, ParamBuf>, key: &str) -> anyhow::Result<ParamBuf> {
    raw.remove(key)
        .ok_or_else(|| anyhow::anyhow!("missing weight key: {key}"))
}

/// Split a fused `in_proj` weight `[3*dim, dim]` → three `[dim, dim]` RLX layouts.
fn split_in_proj(p: ParamBuf, dim: usize) -> anyhow::Result<(ParamBuf, ParamBuf, ParamBuf)> {
    anyhow::ensure!(
        p.shape == vec![dim, 3 * dim],
        "in_proj shape mismatch: {:?}",
        p.shape
    );
    let mut wq = vec![0f32; dim * dim];
    let mut wk = vec![0f32; dim * dim];
    let mut wv = vec![0f32; dim * dim];
    for i in 0..dim {
        for j in 0..dim {
            let src = i * 3 * dim + j;
            wq[i * dim + j] = p.data[src];
            wk[i * dim + j] = p.data[src + dim];
            wv[i * dim + j] = p.data[src + 2 * dim];
        }
    }
    Ok((
        ParamBuf { data: wq, shape: vec![dim, dim] },
        ParamBuf { data: wk, shape: vec![dim, dim] },
        ParamBuf { data: wv, shape: vec![dim, dim] },
    ))
}

fn split_in_proj_bias(p: ParamBuf, dim: usize) -> anyhow::Result<(ParamBuf, ParamBuf, ParamBuf)> {
    anyhow::ensure!(p.shape == vec![3 * dim], "in_proj bias shape mismatch");
    Ok((
        ParamBuf { data: p.data[0..dim].to_vec(), shape: vec![dim] },
        ParamBuf { data: p.data[dim..2 * dim].to_vec(), shape: vec![dim] },
        ParamBuf { data: p.data[2 * dim..3 * dim].to_vec(), shape: vec![dim] },
    ))
}

fn insert_fused_mha(
    p: &mut ParamMap,
    prefix: &str,
    raw: &mut HashMap<String, ParamBuf>,
    dim: usize,
) -> anyhow::Result<()> {
    let w = take_linear_w(raw, &format!("{prefix}.in_proj_weight"))?;
    let (wq, wk, wv) = split_in_proj(w, dim)?;
    p.insert(format!("{prefix}.wq.weight"), wq);
    p.insert(format!("{prefix}.wk.weight"), wk);
    p.insert(format!("{prefix}.wv.weight"), wv);
    if let Ok(b) = take(raw, &format!("{prefix}.in_proj_bias")) {
        let (bq, bk, bv) = split_in_proj_bias(b, dim)?;
        p.insert(format!("{prefix}.wq.bias"), bq);
        p.insert(format!("{prefix}.wk.bias"), bk);
        p.insert(format!("{prefix}.wv.bias"), bv);
    }
    p.insert(
        format!("{prefix}.wo.weight"),
        take_linear_w(raw, &format!("{prefix}.out_proj.weight"))?,
    );
    if let Ok(b) = take(raw, &format!("{prefix}.out_proj.bias")) {
        p.insert(format!("{prefix}.wo.bias"), b);
    }
    Ok(())
}

/// One direction (`mamba_fwd` / `mamba_rev`) of a bidirectional Mamba block.
fn insert_mamba_dir(
    p: &mut ParamMap,
    prefix: &str,
    raw: &mut HashMap<String, ParamBuf>,
) -> anyhow::Result<()> {
    // in_proj: [2*d_inner, d_model] → [d_model, 2*d_inner] (no bias in mamba_ssm).
    p.insert(
        format!("{prefix}.in_proj.weight"),
        take_linear_w(raw, &format!("{prefix}.in_proj.weight"))?,
    );

    // Depthwise causal conv1d: [d_inner, 1, k] → grouped conv2d NCHW [d_inner, 1, 1, k].
    let mut conv = take(raw, &format!("{prefix}.conv1d.weight"))?;
    anyhow::ensure!(
        conv.shape.len() == 3 && conv.shape[1] == 1,
        "conv1d.weight expected [d_inner,1,k], got {:?}",
        conv.shape
    );
    conv.shape = vec![conv.shape[0], 1, 1, conv.shape[2]];
    p.insert(format!("{prefix}.conv1d.weight"), conv);
    p.insert(
        format!("{prefix}.conv1d.bias"),
        take(raw, &format!("{prefix}.conv1d.bias"))?,
    );

    // x_proj: [dt_rank + 2N, d_inner] → [d_inner, dt_rank+2N] (no bias).
    p.insert(
        format!("{prefix}.x_proj.weight"),
        take_linear_w(raw, &format!("{prefix}.x_proj.weight"))?,
    );

    // dt_proj: [d_inner, dt_rank] → [dt_rank, d_inner] (+ bias).
    p.insert(
        format!("{prefix}.dt_proj.weight"),
        take_linear_w(raw, &format!("{prefix}.dt_proj.weight"))?,
    );
    p.insert(
        format!("{prefix}.dt_proj.bias"),
        take(raw, &format!("{prefix}.dt_proj.bias"))?,
    );

    // A = -exp(A_log)  [d_inner, N].  RLX selective_scan does NOT negate A.
    let alog = take(raw, &format!("{prefix}.A_log"))?;
    let a = ParamBuf {
        data: alog.data.iter().map(|v| -v.exp()).collect(),
        shape: alog.shape.clone(),
    };
    p.insert(format!("{prefix}.A"), a);

    // D skip term [d_inner]; applied as y += D ⊙ x outside the scan op.
    p.insert(format!("{prefix}.D"), take(raw, &format!("{prefix}.D"))?);

    // out_proj: [d_model, d_inner] → [d_inner, d_model] (no bias).
    p.insert(
        format!("{prefix}.out_proj.weight"),
        take_linear_w(raw, &format!("{prefix}.out_proj.weight"))?,
    );
    Ok(())
}

/// Build the parameter map consumed by [`super::graph::build_forward_graph`].
pub fn build_forward_params(
    raw: &mut HashMap<String, ParamBuf>,
    cfg: &ModelConfig,
) -> anyhow::Result<ParamMap> {
    let d = cfg.embed_dim;
    let mut p = ParamMap::new();

    // ── Cross-attention channel unification ─────────────────────────────────
    p.insert(
        "cross_attn.query_embed".into(),
        take(raw, "cross_attn.query_embed")?,
    );
    for norm in ["queries_norm", "keys_norm", "values_norm"] {
        p.insert(
            format!("cross_attn.{norm}.weight"),
            take(raw, &format!("cross_attn.{norm}.weight"))?,
        );
        p.insert(
            format!("cross_attn.{norm}.bias"),
            take(raw, &format!("cross_attn.{norm}.bias"))?,
        );
    }
    insert_fused_mha(&mut p, "cross_attn.cross_attention", raw, d)?;
    p.insert("cross_attn.ffn.fc1.weight".into(), take_linear_w(raw, "cross_attn.ffn.fc1.weight")?);
    p.insert("cross_attn.ffn.fc1.bias".into(),   take(raw, "cross_attn.ffn.fc1.bias")?);
    p.insert("cross_attn.ffn.norm.weight".into(), take(raw, "cross_attn.ffn.norm.weight")?);
    p.insert("cross_attn.ffn.norm.bias".into(),   take(raw, "cross_attn.ffn.norm.bias")?);
    p.insert("cross_attn.ffn.fc2.weight".into(), take_linear_w(raw, "cross_attn.ffn.fc2.weight")?);
    p.insert("cross_attn.ffn.fc2.bias".into(),   take(raw, "cross_attn.ffn.fc2.bias")?);
    for i in 0..3 {
        let q = format!("cross_attn.query_self_attn.layers.{i}");
        p.insert(format!("{q}.norm1.weight"), take(raw, &format!("{q}.norm1.weight"))?);
        p.insert(format!("{q}.norm1.bias"),   take(raw, &format!("{q}.norm1.bias"))?);
        insert_fused_mha(&mut p, &format!("{q}.self_attn"), raw, d)?;
        p.insert(format!("{q}.norm2.weight"), take(raw, &format!("{q}.norm2.weight"))?);
        p.insert(format!("{q}.norm2.bias"),   take(raw, &format!("{q}.norm2.bias"))?);
        p.insert(format!("{q}.linear1.weight"), take_linear_w(raw, &format!("{q}.linear1.weight"))?);
        p.insert(format!("{q}.linear1.bias"),   take(raw, &format!("{q}.linear1.bias"))?);
        p.insert(format!("{q}.linear2.weight"), take_linear_w(raw, &format!("{q}.linear2.weight"))?);
        p.insert(format!("{q}.linear2.bias"),   take(raw, &format!("{q}.linear2.bias"))?);
    }

    // ── FEMBA bidirectional-Mamba encoder ───────────────────────────────────
    for i in 0..cfg.num_blocks {
        for dir in ["mamba_fwd", "mamba_rev"] {
            insert_mamba_dir(&mut p, &format!("mamba_blocks.{i}.{dir}"), raw)?;
        }
        p.insert(
            format!("norm_layers.{i}.weight"),
            take(raw, &format!("norm_layers.{i}.weight"))?,
        );
        p.insert(
            format!("norm_layers.{i}.bias"),
            take(raw, &format!("norm_layers.{i}.bias"))?,
        );
    }
    // NB: LuMamba deletes LUNA's trailing `self.norm` — no encoder-output norm.

    // ── LUNA-style classifier head (fine-tuned, classifier_option=None) ──────
    // Loaded by key presence so a fine-tuned checkpoint works regardless of the
    // `num_classes` in the JSON config. The graph only references these when
    // `num_classes > 0` (see build_forward_graph). The linear-classifier head
    // is loaded host-side in build_prepare_params instead.
    if raw.contains_key("classifier.decoder_attn.in_proj_weight") {
        if let Ok(agg) = take(raw, "classifier.learned_agg") {
            p.insert("classifier.learned_agg".into(), agg);
        }
        insert_fused_mha(&mut p, "classifier.decoder_attn", raw, cfg.hidden_dim())?;
        p.insert("classifier.decoder_ffn.fc1.weight".into(), take_linear_w(raw, "classifier.decoder_ffn.fc1.weight")?);
        p.insert("classifier.decoder_ffn.fc1.bias".into(),   take(raw, "classifier.decoder_ffn.fc1.bias")?);
        p.insert("classifier.decoder_ffn.fc2.weight".into(), take_linear_w(raw, "classifier.decoder_ffn.fc2.weight")?);
        p.insert("classifier.decoder_ffn.fc2.bias".into(),   take(raw, "classifier.decoder_ffn.fc2.bias")?);
    }

    // ── Reconstruction head (present in pretrained checkpoints) ──────────────
    if !raw.contains_key("decoder_head.decoder_pred.layers.0.norm1.weight") {
        return Ok(p);
    }
    let dp = "decoder_head.decoder_pred.layers.0";
    p.insert(format!("{dp}.norm1.weight"), take(raw, &format!("{dp}.norm1.weight"))?);
    p.insert(format!("{dp}.norm1.bias"),   take(raw, &format!("{dp}.norm1.bias"))?);
    insert_fused_mha(&mut p, &format!("{dp}.self_attn"), raw, d)?;
    p.insert(format!("{dp}.norm2.weight"), take(raw, &format!("{dp}.norm2.weight"))?);
    p.insert(format!("{dp}.norm2.bias"),   take(raw, &format!("{dp}.norm2.bias"))?);
    insert_fused_mha(&mut p, &format!("{dp}.multihead_attn"), raw, d)?;
    p.insert(format!("{dp}.norm3.weight"), take(raw, &format!("{dp}.norm3.weight"))?);
    p.insert(format!("{dp}.norm3.bias"),   take(raw, &format!("{dp}.norm3.bias"))?);
    p.insert(format!("{dp}.linear1.weight"), take_linear_w(raw, &format!("{dp}.linear1.weight"))?);
    p.insert(format!("{dp}.linear1.bias"),   take(raw, &format!("{dp}.linear1.bias"))?);
    p.insert(format!("{dp}.linear2.weight"), take_linear_w(raw, &format!("{dp}.linear2.weight"))?);
    p.insert(format!("{dp}.linear2.bias"),   take(raw, &format!("{dp}.linear2.bias"))?);
    p.insert("decoder_head.norm.weight".into(), take(raw, "decoder_head.norm.weight")?);
    p.insert("decoder_head.norm.bias".into(),   take(raw, "decoder_head.norm.bias")?);
    p.insert("decoder_head.decoder_linear.fc1.weight".into(), take_linear_w(raw, "decoder_head.decoder_linear.fc1.weight")?);
    p.insert("decoder_head.decoder_linear.fc1.bias".into(),   take(raw, "decoder_head.decoder_linear.fc1.bias")?);
    p.insert("decoder_head.decoder_linear.fc2.weight".into(), take_linear_w(raw, "decoder_head.decoder_linear.fc2.weight")?);
    p.insert("decoder_head.decoder_linear.fc2.bias".into(),   take(raw, "decoder_head.decoder_linear.fc2.bias")?);

    Ok(p)
}

/// Parameters for the CPU token-preparation path (patch / freq / channel MLP).
pub fn build_prepare_params(raw: &mut HashMap<String, ParamBuf>) -> anyhow::Result<ParamMap> {
    let mut p = ParamMap::new();
    // proj_in layout: conv0, gn1, (gelu), conv3, gn4, (gelu), conv6, gn7, (gelu).
    for (conv_idx, gn_idx, conv, gn) in [(0, 1, "conv1", 1), (3, 4, "conv2", 2), (6, 7, "conv3", 3)] {
        p.insert(
            format!("patch_embed.{conv}.weight"),
            take(raw, &format!("patch_embed.proj_in.{conv_idx}.weight"))?,
        );
        p.insert(
            format!("patch_embed.{conv}.bias"),
            take(raw, &format!("patch_embed.proj_in.{conv_idx}.bias"))?,
        );
        p.insert(
            format!("patch_embed.gn{gn}.weight"),
            take(raw, &format!("patch_embed.proj_in.{gn_idx}.weight"))?,
        );
        p.insert(
            format!("patch_embed.gn{gn}.bias"),
            take(raw, &format!("patch_embed.proj_in.{gn_idx}.bias"))?,
        );
    }
    p.insert("freq_embed.fc1.weight".into(), take_linear_w(raw, "freq_embed.frequency_to_embed.fc1.weight")?);
    p.insert("freq_embed.fc1.bias".into(),   take(raw, "freq_embed.frequency_to_embed.fc1.bias")?);
    p.insert("freq_embed.fc2.weight".into(), take_linear_w(raw, "freq_embed.frequency_to_embed.fc2.weight")?);
    p.insert("freq_embed.fc2.bias".into(),   take(raw, "freq_embed.frequency_to_embed.fc2.bias")?);
    p.insert("chan_loc.fc1.weight".into(),  take_linear_w(raw, "channel_location_embedder.0.fc1.weight")?);
    p.insert("chan_loc.fc1.bias".into(),    take(raw, "channel_location_embedder.0.fc1.bias")?);
    p.insert("chan_loc.norm.weight".into(), take(raw, "channel_location_embedder.0.norm.weight")?);
    p.insert("chan_loc.norm.bias".into(),   take(raw, "channel_location_embedder.0.norm.bias")?);
    p.insert("chan_loc.fc2.weight".into(),  take_linear_w(raw, "channel_location_embedder.0.fc2.weight")?);
    p.insert("chan_loc.fc2.bias".into(),    take(raw, "channel_location_embedder.0.fc2.bias")?);
    if let Ok(t) = take(raw, "channel_emb.embeddings.weight") {
        // Embedding table [vocab, D] — gathered on the host for decoder queries.
        p.insert("channel_emb.weight".into(), t);
    }
    // Linear classifier head (BasicLinearClassifier): applied host-side on the
    // mean-pooled encoder latent. Loaded only if present in the checkpoint.
    if raw.contains_key("classifier.fc1.weight") {
        p.insert("classifier.fc1.weight".into(), take_linear_w(raw, "classifier.fc1.weight")?);
        p.insert("classifier.fc1.bias".into(), take(raw, "classifier.fc1.bias")?);
    }
    Ok(p)
}

/// Bind every parameter in `params` into a compiled graph by name.
pub fn apply_params(compiled: &mut rlx::CompiledGraph, params: &ParamMap) {
    for (name, buf) in params {
        compiled.set_param(name, &buf.data);
    }
}