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

//! LuMamba forward pass: cross-attention channel unification → bidirectional
//! Mamba (FEMBA) encoder → reconstruction / classification head.
//!
//! The LUNA front-end (cross-attention + query self-attention) and the
//! reconstruction head are structurally identical to `luna-rs`; the temporal
//! encoder is the part that differs — LUNA's rotary Transformer blocks are
//! replaced by `num_blocks` pre-norm residual bidirectional Mamba blocks.

#![allow(clippy::too_many_arguments)]

use rlx::ir::GraphExt;
use rlx::prelude::*;

/// Resolved per-shape dimensions for one forward pass. Fields are terse
/// dimension shorthands (documented inline); built from [`crate::ModelConfig`]
/// plus the input `(batch, channels, time)`.
#[derive(Clone, Copy, Debug)]
#[allow(missing_docs)]
pub struct ForwardSpec {
    pub b: usize,        // batch (always 1 in the per-epoch path)
    pub c: usize,        // channels
    pub s: usize,        // patches  (T / patch_size)
    pub bt: usize,       // b * s    (token rows)
    pub d: usize,        // embed_dim E
    pub q: usize,        // num_queries Q
    pub hidden: usize,   // d_model = Q * E
    pub nh_ca: usize,    // cross-attention heads
    pub dh_ca: usize,    // cross-attention head dim
    pub ff_ca: usize,    // cross-attention FFN hidden
    pub patch_size: usize,
    pub norm_eps: f32,
    // ── Mamba ──────────────────────────────────────────────────────────────
    pub num_blocks: usize,
    pub d_inner: usize,  // expand * d_model
    pub d_state: usize,  // N
    pub d_conv: usize,   // depthwise conv kernel width
    pub dt_rank: usize,
    pub bidir_multiply: bool,
    /// `0` = reconstruction head, `>0` = classification logits.
    pub num_classes: usize,
    pub nh_cls: usize,
}

fn s1(d: usize) -> Shape { Shape::new(&[d], DType::F32) }
fn s2(a: usize, b: usize) -> Shape { Shape::new(&[a, b], DType::F32) }
fn s3(a: usize, b: usize, c: usize) -> Shape { Shape::new(&[a, b, c], DType::F32) }
fn s4(a: usize, b: usize, c: usize, d: usize) -> Shape { Shape::new(&[a, b, c, d], DType::F32) }

/// `[B,S,H,D]` → `[B,H,S,D]` (canonical attention layout, required by MLX SDPA).
fn to_bhsd(g: &mut Graph, t: NodeId) -> NodeId {
    g.transpose_(t, vec![0, 2, 1, 3])
}

fn ln(g: &mut Graph, x: NodeId, w: NodeId, b: NodeId, eps: f32) -> NodeId {
    g.ln(x, w, b, eps)
}

// ── Multi-head attention (split Q/K/V projections + RLX fused SDPA) ──────────
fn mha(
    g: &mut Graph,
    q_in: NodeId,
    k_in: NodeId,
    v_in: NodeId,
    prefix: &str,
    batch: usize,
    s_q: usize,
    s_kv: usize,
    d: usize,
    nh: usize,
    dh: usize,
) -> NodeId {
    let h_total = nh * dh;
    let wq = g.param(format!("{prefix}.wq.weight"), s2(d, h_total));
    let wk = g.param(format!("{prefix}.wk.weight"), s2(d, h_total));
    let wv = g.param(format!("{prefix}.wv.weight"), s2(d, h_total));
    let wo = g.param(format!("{prefix}.wo.weight"), s2(h_total, d));
    let wq_b = g.param(format!("{prefix}.wq.bias"), s1(h_total));
    let wk_b = g.param(format!("{prefix}.wk.bias"), s1(h_total));
    let wv_b = g.param(format!("{prefix}.wv.bias"), s1(h_total));
    let wo_b = g.param(format!("{prefix}.wo.bias"), s1(d));

    let qm = g.mm(q_in, wq);
    let q = g.add(qm, wq_b);
    let km = g.mm(k_in, wk);
    let k = g.add(km, wk_b);
    let vm = g.mm(v_in, wv);
    let v = g.add(vm, wv_b);

    let q4 = g.reshape_(q, vec![batch as i64, s_q as i64, nh as i64, dh as i64]);
    let k4 = g.reshape_(k, vec![batch as i64, s_kv as i64, nh as i64, dh as i64]);
    let v4 = g.reshape_(v, vec![batch as i64, s_kv as i64, nh as i64, dh as i64]);

    let q_bhsd = to_bhsd(g, q4);
    let k_bhsd = to_bhsd(g, k4);
    let v_bhsd = to_bhsd(g, v4);
    let attn = g.attention_kind(q_bhsd, k_bhsd, v_bhsd, nh, dh, MaskKind::None, s4(batch, nh, s_q, dh));
    let attn_bshd = to_bhsd(g, attn);
    let attn_3 = g.reshape_(attn_bshd, vec![batch as i64, s_q as i64, h_total as i64]);
    let om = g.mm(attn_3, wo);
    g.add(om, wo_b)
}

fn transformer_encoder_layer(
    g: &mut Graph,
    x: NodeId,
    prefix: &str,
    batch: usize,
    s: usize,
    d: usize,
    nh: usize,
    dh: usize,
    ff: usize,
    eps: f32,
) -> NodeId {
    let n1w = g.param(format!("{prefix}.norm1.weight"), s1(d));
    let n1b = g.param(format!("{prefix}.norm1.bias"), s1(d));
    let normed = ln(g, x, n1w, n1b, eps);
    let attn = mha(g, normed, normed, normed, &format!("{prefix}.self_attn"), batch, s, s, d, nh, dh);
    let x = g.add(x, attn);

    let n2w = g.param(format!("{prefix}.norm2.weight"), s1(d));
    let n2b = g.param(format!("{prefix}.norm2.bias"), s1(d));
    let normed = ln(g, x, n2w, n2b, eps);
    let l1w = g.param(format!("{prefix}.linear1.weight"), s2(d, ff));
    let l1b = g.param(format!("{prefix}.linear1.bias"), s1(ff));
    let l2w = g.param(format!("{prefix}.linear2.weight"), s2(ff, d));
    let l2b = g.param(format!("{prefix}.linear2.bias"), s1(d));
    let l1m = g.mm(normed, l1w);
    let l1b_ = g.add(l1m, l1b);
    let h = g.gelu(l1b_);
    let l2m = g.mm(h, l2w);
    let ff_out = g.add(l2m, l2b);
    g.add(x, ff_out)
}

fn cross_attention_block(g: &mut Graph, x: NodeId, queries: NodeId, spec: &ForwardSpec) -> NodeId {
    let d = spec.d;
    let qn = spec.q;
    let bt = spec.bt;
    let c = spec.c;
    let nh = spec.nh_ca;
    let dh = spec.dh_ca;
    let eps = spec.norm_eps;

    let qnw = g.param("cross_attn.queries_norm.weight", s1(d));
    let qnb = g.param("cross_attn.queries_norm.bias", s1(d));
    let knw = g.param("cross_attn.keys_norm.weight", s1(d));
    let knb = g.param("cross_attn.keys_norm.bias", s1(d));
    let vnw = g.param("cross_attn.values_norm.weight", s1(d));
    let vnb = g.param("cross_attn.values_norm.bias", s1(d));

    let q = ln(g, queries, qnw, qnb, eps);
    let k = ln(g, x, knw, knb, eps);
    let v = ln(g, x, vnw, vnb, eps);

    let mut out = mha(g, q, k, v, "cross_attn.cross_attention", bt, qn, c, d, nh, dh);

    let f1w = g.param("cross_attn.ffn.fc1.weight", s2(d, spec.ff_ca));
    let f1b = g.param("cross_attn.ffn.fc1.bias", s1(spec.ff_ca));
    let fnw = g.param("cross_attn.ffn.norm.weight", s1(spec.ff_ca));
    let fnb = g.param("cross_attn.ffn.norm.bias", s1(spec.ff_ca));
    let f2w = g.param("cross_attn.ffn.fc2.weight", s2(spec.ff_ca, d));
    let f2b = g.param("cross_attn.ffn.fc2.bias", s1(d));

    let residual = out;
    let f1m = g.mm(out, f1w);
    let f1h = g.add(f1m, f1b);
    let h = g.gelu(f1h);
    let h = ln(g, h, fnw, fnb, eps);
    let f2m = g.mm(h, f2w);
    let f2h = g.add(f2m, f2b);
    out = g.add(residual, f2h);

    for i in 0..3 {
        out = transformer_encoder_layer(
            g, out, &format!("cross_attn.query_self_attn.layers.{i}"),
            bt, qn, d, nh, dh, spec.ff_ca, eps,
        );
    }
    out
}

// ── Bidirectional Mamba encoder ──────────────────────────────────────────────

/// `softplus(x) = log(1 + exp(x))`. Composed from existing RLX ops; exact in
/// f32 across the `dt` range Mamba produces (the `dt_proj` bias keeps it small).
fn softplus(g: &mut Graph, x: NodeId, shape: Shape) -> NodeId {
    let e = g.exp(x);
    let one = g.constant(1.0, DType::F32);
    let ep1 = g.add(e, one);
    g.activation(Activation::Log, ep1, shape)
}

/// Reverse a `[b, s, h]` tensor along the sequence axis (`Op::Reverse`).
fn flip_seq(g: &mut Graph, x: NodeId, b: usize, s: usize, h: usize) -> NodeId {
    g.add_node(Op::Reverse { axes: vec![1] }, vec![x], s3(b, s, h))
}

/// Causal depthwise conv1d over the sequence axis, emulated with a grouped
/// `conv2d`: left-pad `k-1` zeros, run a `groups = d_inner` conv, add bias.
fn causal_conv1d(g: &mut Graph, x: NodeId, prefix: &str, spec: &ForwardSpec) -> NodeId {
    let b = spec.b;
    let s = spec.s;
    let di = spec.d_inner;
    let k = spec.d_conv;

    // Prepend (k-1) zeros along the sequence axis → causal receptive field.
    let pad_elems = b * (k - 1) * di;
    let zeros = g.add_node(
        Op::Constant { data: vec![0u8; pad_elems * 4] },
        vec![],
        Shape::new(&[b, k - 1, di], DType::F32),
    );
    let xpad = g.concat_(vec![zeros, x], 1); // [b, s+k-1, di]

    // [b, s+k-1, di] → NCHW [b, di, 1, s+k-1]
    let xt = g.transpose_(xpad, vec![0, 2, 1]); // [b, di, s+k-1]
    let x4 = g.reshape_(xt, vec![b as i64, di as i64, 1, (s + k - 1) as i64]);

    let w = g.param(format!("{prefix}.conv1d.weight"), s4(di, 1, 1, k));
    let y4 = g.conv2d(x4, w, [1, k], [1, 1], [0, 0], [1, 1], di); // [b, di, 1, s]

    let y = g.reshape_(y4, vec![b as i64, di as i64, s as i64]);
    let yt = g.transpose_(y, vec![0, 2, 1]); // [b, s, di]
    let cb = g.param(format!("{prefix}.conv1d.bias"), s1(di));
    g.add(yt, cb)
}

/// A single (forward or reverse) Mamba block on `[b, s, hidden]` input.
/// When `dbg` is `Some`, key intermediates are pushed in a fixed order for
/// stage-by-stage validation (see `build_encoder_debug_graph`).
fn mamba_inner(
    g: &mut Graph,
    x: NodeId,
    prefix: &str,
    spec: &ForwardSpec,
    mut dbg: Option<&mut Vec<NodeId>>,
) -> NodeId {
    let b = spec.b;
    let s = spec.s;
    let di = spec.d_inner;
    let n = spec.d_state;
    let dtr = spec.dt_rank;
    let hidden = spec.hidden;

    macro_rules! rec {
        ($node:expr) => {
            if let Some(v) = dbg.as_deref_mut() {
                v.push($node);
            }
        };
    }

    // in_proj → [x_ssm | z]
    let in_w = g.param(format!("{prefix}.in_proj.weight"), s2(hidden, 2 * di));
    let xz = g.mm(x, in_w);                 // [b, s, 2*di]
    let x_ssm = g.narrow_(xz, 2, 0, di);    // [b, s, di]
    let z = g.narrow_(xz, 2, di, di);       // [b, s, di]
    rec!(x_ssm);

    // causal depthwise conv + SiLU
    let x_conv = causal_conv1d(g, x_ssm, prefix, spec);
    rec!(x_conv);
    let x_act = g.silu(x_conv);
    rec!(x_act);

    // x_proj → dt, B, C
    let xproj_w = g.param(format!("{prefix}.x_proj.weight"), s2(di, dtr + 2 * n));
    let x_dbl = g.mm(x_act, xproj_w);            // [b, s, dt_rank + 2N]
    let dt_in = g.narrow_(x_dbl, 2, 0, dtr);     // [b, s, dt_rank]
    let bmat = g.narrow_(x_dbl, 2, dtr, n);      // [b, s, N]
    let cmat = g.narrow_(x_dbl, 2, dtr + n, n);  // [b, s, N]
    rec!(bmat);
    rec!(cmat);

    // delta = softplus(dt_proj(dt) + bias)
    let dtw = g.param(format!("{prefix}.dt_proj.weight"), s2(dtr, di));
    let dtb = g.param(format!("{prefix}.dt_proj.bias"), s1(di));
    let dt = g.mm(dt_in, dtw);
    let dt = g.add(dt, dtb);
    let delta = softplus(g, dt, s3(b, s, di));
    rec!(delta);

    // selective scan: y = C·h, with h_t = exp(delta·A)·h_{t-1} + delta·B·x
    let a = g.param(format!("{prefix}.A"), s2(di, n)); // A = -exp(A_log), pre-computed
    let y = g.selective_scan(x_act, delta, a, bmat, cmat, n, s3(b, s, di));
    rec!(y);

    // + D ⊙ x   (skip term, applied outside the scan op)
    let dparam = g.param(format!("{prefix}.D"), s1(di));
    let dx = g.mul(x_act, dparam);
    let y = g.add(y, dx);

    // gate by SiLU(z)
    let zg = g.silu(z);
    let y = g.mul(y, zg);
    rec!(y);

    // out_proj → [b, s, hidden]
    let ow = g.param(format!("{prefix}.out_proj.weight"), s2(di, hidden));
    let out = g.mm(y, ow);
    rec!(out);
    out
}

/// Bidirectional wrapper: `out = mamba_fwd(x) (⊕) flip(mamba_rev(flip(x)))`,
/// where `⊕` is add (default) or elementwise multiply.
fn mamba_wrapper(g: &mut Graph, x: NodeId, idx: usize, spec: &ForwardSpec) -> NodeId {
    let out_fwd = mamba_inner(g, x, &format!("mamba_blocks.{idx}.mamba_fwd"), spec, None);

    let x_flip = flip_seq(g, x, spec.b, spec.s, spec.hidden);
    let out_rev0 = mamba_inner(g, x_flip, &format!("mamba_blocks.{idx}.mamba_rev"), spec, None);
    let out_rev = flip_seq(g, out_rev0, spec.b, spec.s, spec.hidden);

    if spec.bidir_multiply {
        g.mul(out_fwd, out_rev)
    } else {
        g.add(out_fwd, out_rev)
    }
}

// ── Heads ────────────────────────────────────────────────────────────────────

fn reconstruction_head(g: &mut Graph, enc: NodeId, tgt: NodeId, spec: &ForwardSpec) -> NodeId {
    let d = spec.d;
    let qn = spec.q;
    let b = spec.b;
    let s = spec.s;
    let bt = spec.bt;
    let c = spec.c;
    let p = spec.patch_size;
    let nh = spec.nh_ca;
    let dh = spec.dh_ca;
    let eps = spec.norm_eps;
    let dp = "decoder_head.decoder_pred.layers.0";

    let memory = g.reshape_(enc, vec![bt as i64, qn as i64, d as i64]);
    let mut x = tgt;

    let n1w = g.param(format!("{dp}.norm1.weight"), s1(d));
    let n1b = g.param(format!("{dp}.norm1.bias"), s1(d));
    let normed = ln(g, x, n1w, n1b, eps);
    let sa = mha(g, normed, normed, normed, &format!("{dp}.self_attn"), bt, c, c, d, nh, dh);
    x = g.add(x, sa);

    let n2w = g.param(format!("{dp}.norm2.weight"), s1(d));
    let n2b = g.param(format!("{dp}.norm2.bias"), s1(d));
    let normed = ln(g, x, n2w, n2b, eps);
    let ca = mha(g, normed, memory, memory, &format!("{dp}.multihead_attn"), bt, c, qn, d, nh, dh);
    x = g.add(x, ca);

    let n3w = g.param(format!("{dp}.norm3.weight"), s1(d));
    let n3b = g.param(format!("{dp}.norm3.bias"), s1(d));
    let normed = ln(g, x, n3w, n3b, eps);
    let l1w = g.param(format!("{dp}.linear1.weight"), s2(d, d * 4));
    let l1b = g.param(format!("{dp}.linear1.bias"), s1(d * 4));
    let l2w = g.param(format!("{dp}.linear2.weight"), s2(d * 4, d));
    let l2b = g.param(format!("{dp}.linear2.bias"), s1(d));
    let l1m = g.mm(normed, l1w);
    let l1b_ = g.add(l1m, l1b);
    let ff = g.gelu(l1b_);
    let l2m = g.mm(ff, l2w);
    let l2b_ = g.add(l2m, l2b);
    x = g.add(x, l2b_);

    let onw = g.param("decoder_head.norm.weight", s1(d));
    let onb = g.param("decoder_head.norm.bias", s1(d));
    x = ln(g, x, onw, onb, eps);

    let of1w = g.param("decoder_head.decoder_linear.fc1.weight", s2(d, d * 4));
    let of1b = g.param("decoder_head.decoder_linear.fc1.bias", s1(d * 4));
    let of2w = g.param("decoder_head.decoder_linear.fc2.weight", s2(d * 4, p));
    let of2b = g.param("decoder_head.decoder_linear.fc2.bias", s1(p));

    let o1m = g.mm(x, of1w);
    let o1b_ = g.add(o1m, of1b);
    let h = g.gelu(o1b_);
    let o2m = g.mm(h, of2w);
    let out = g.add(o2m, of2b);

    // [bt, C, P] → [B, C, T]
    let out5 = g.reshape_(out, vec![b as i64, s as i64, c as i64, p as i64]);
    let out5 = g.transpose_(out5, vec![0, 2, 1, 3]);
    g.reshape_(out5, vec![b as i64, c as i64, (s * p) as i64])
}

fn classification_head(g: &mut Graph, enc: NodeId, agg: NodeId, spec: &ForwardSpec) -> NodeId {
    let hidden = spec.hidden;
    let b = spec.b;
    let nh = spec.nh_cls;
    let dh = hidden / nh;
    let nc = spec.num_classes;

    let ca = mha(g, agg, enc, enc, "classifier.decoder_attn", b, 1, spec.s, hidden, nh, dh);
    let logits_in = g.reshape_(ca, vec![b as i64, hidden as i64]);
    let f1w = g.param("classifier.decoder_ffn.fc1.weight", s2(hidden, hidden * 4));
    let f1b = g.param("classifier.decoder_ffn.fc1.bias", s1(hidden * 4));
    let f2w = g.param("classifier.decoder_ffn.fc2.weight", s2(hidden * 4, nc));
    let f2b = g.param("classifier.decoder_ffn.fc2.bias", s1(nc));
    let l1m = g.mm(logits_in, f1w);
    let l1b_ = g.add(l1m, f1b);
    let h = g.gelu(l1b_);
    let l2m = g.mm(h, f2w);
    g.add(l2m, f2b)
}

/// The bidirectional-Mamba temporal encoder: `num_blocks` pre-norm residual
/// blocks over a `[B, S, hidden]` latent. NO trailing encoder norm (LuMamba
/// deletes LUNA's `self.norm`).
fn mamba_stack(g: &mut Graph, mut h: NodeId, spec: &ForwardSpec) -> NodeId {
    for i in 0..spec.num_blocks {
        let res = h;
        let nw = g.param(format!("norm_layers.{i}.weight"), s1(spec.hidden));
        let nb = g.param(format!("norm_layers.{i}.bias"), s1(spec.hidden));
        let hn = ln(g, h, nw, nb, spec.norm_eps);
        let m = mamba_wrapper(g, hn, i, spec);
        h = g.add(res, m);
    }
    h
}

/// Build a graph for the temporal encoder alone: input `h_in` `[B, S, hidden]`
/// (already channel-unified) → encoder latent `[B, S, hidden]`. Used to run /
/// validate the FEMBA stack in isolation.
pub fn build_encoder_graph(spec: &ForwardSpec) -> Graph {
    let mut g = Graph::new("lumamba_encoder");
    let h_in = g.input("h_in", s3(spec.b, spec.s, spec.hidden));
    let h = mamba_stack(&mut g, h_in, spec);
    g.set_outputs(vec![h]);
    g
}

/// Debug graph: run ONLY `mamba_blocks.0.mamba_fwd` on `h_in` (no norm, no
/// reverse stream, no residual) and expose per-stage intermediates. Output
/// order: `x_ssm, x_conv, x_act, B, C, delta, scan_y, y_gated, out`.
#[cfg(feature = "validation")]
pub fn build_encoder_debug_graph(spec: &ForwardSpec) -> Graph {
    let mut g = Graph::new("lumamba_encoder_debug");
    let h_in = g.input("h_in", s3(spec.b, spec.s, spec.hidden));
    let mut dbg = Vec::new();
    let _ = mamba_inner(&mut g, h_in, "mamba_blocks.0.mamba_fwd", spec, Some(&mut dbg));
    g.set_outputs(dbg);
    g
}

/// Debug graph #2: bidirectional wrapper of block 0 on `h_in` (no norm, no
/// residual). Output order: `flip_in, fwd, rev_out, wrapper`.
#[cfg(feature = "validation")]
pub fn build_encoder_debug2_graph(spec: &ForwardSpec) -> Graph {
    let mut g = Graph::new("lumamba_encoder_debug2");
    let h_in = g.input("h_in", s3(spec.b, spec.s, spec.hidden));
    let fwd = mamba_inner(&mut g, h_in, "mamba_blocks.0.mamba_fwd", spec, None);
    let flip_in = flip_seq(&mut g, h_in, spec.b, spec.s, spec.hidden);
    let rev_raw = mamba_inner(&mut g, flip_in, "mamba_blocks.0.mamba_rev", spec, None);
    let rev_out = flip_seq(&mut g, rev_raw, spec.b, spec.s, spec.hidden);
    let wrapper = g.add(fwd, rev_out);
    g.set_outputs(vec![flip_in, fwd, rev_out, wrapper]);
    g
}

/// Build the LuMamba forward graph for one `(batch, channels, time)` shape.
///
/// Inputs:
/// * `x_tokenized` — `[B*S, C, E]`
/// * `queries` — `[B*S, Q, E]` (broadcast cross-attention query prototypes)
/// * `decoder_queries` — `[B*S, C, E]` (reconstruction only)
/// * `agg_query` — `[B, 1, Q*E]` (classification only)
///
/// Outputs: `[head_output, encoder_latent]`, where `encoder_latent` is the
/// `[B, S, Q*E]` foundation embedding (usable directly for downstream tasks).
pub fn build_forward_graph(spec: &ForwardSpec) -> Graph {
    let mut g = Graph::new("lumamba_forward");

    let x = g.input("x_tokenized", s3(spec.bt, spec.c, spec.d));
    let queries = g.input("queries", s3(spec.bt, spec.q, spec.d));

    // 1. Channel unification → [B*S, Q, E]
    let unified = cross_attention_block(&mut g, x, queries, spec);
    // 2. reshape to the temporal sequence [B, S, Q*E]
    let mut h = g.reshape_(unified, vec![spec.b as i64, spec.s as i64, spec.hidden as i64]);

    // 3. Bidirectional Mamba stack (pre-norm residual; NO trailing encoder norm)
    h = mamba_stack(&mut g, h, spec);

    let out = if spec.num_classes > 0 {
        let agg = g.input("agg_query", s3(spec.b, 1, spec.hidden));
        classification_head(&mut g, h, agg, spec)
    } else {
        let dec_q = g.input("decoder_queries", s3(spec.bt, spec.c, spec.d));
        reconstruction_head(&mut g, h, dec_q, spec)
    };

    // Expose both the head output and the encoder latent `h`.
    g.set_outputs(vec![out, h]);
    g
}