inferencelayer 0.2.4

//! Parakeet-TDT ASR (`mlx-community/parakeet-tdt-0.6b-v3`) — FastConformer encoder + a
//! token-and-duration transducer, ported operation-for-operation from `parakeet-mlx` (the package
//! the Python `parakeet-asr` service runs). CPU, f32.
//!
//! The reference computes in f32 **with bf16-rounded weights** (`from_pretrained` casts every
//! parameter to bf16; the mel path stays f32) — the export stores the tensors AS bf16, and
//! `LazySt::tensor_f32` widens them, so this module computes with the exact same values.
//!
//! Reference quirks that are part of the contract, not bugs to fix:
//!
//! * The spectrogram magnitude is **|Re| + |Im|, then squared** — `parakeet_mlx.audio.get_logmel`
//!   views the complex STFT as float pairs and sums `abs` over the last axis in pairs. That is the
//!   L1 magnitude, not `sqrt(Re² + Im²)`; the mel filterbank eats (|Re|+|Im|)².
//! * The hann window is numpy's PERIODIC hann (`np.hanning(N+1)[:-1]`), 400 samples zero-padded to
//!   `n_fft` 512 — so each frame reads 400 samples even though the FFT is 512-point.
//! * The prediction network's start symbol is a **zeros embedding**, not the blank row.
//! * `mx.std` is the POPULATION std (ddof 0), and the normalizer divides by `std + 1e-5`.
//!
//! Layer-by-layer parity: `tests/parakeet_parity.rs` against
//! `tests/fixtures/export_parakeet.py`'s oracle — mel, encoder features, and the DISCRETE greedy
//! decode trace (step/token/duration triples must be identical).

use std::cell::{Cell, RefCell};
use std::f64::consts::PI;
use std::path::Path;

use anyhow::{Context, Result};

use crate::cpu_gemm::{PackedWeight, gemm_packed};
use crate::weights::LazySt;

thread_local! {
    static PROFILE: Cell<bool> = Cell::new(std::env::var("OSFKB_PARAKEET_PROFILE").as_deref() == Ok("1"));
    /// [ffn, attn, conv] accumulated ms — printed by the profile test.
    static STAGE_MS: RefCell<[f64; 3]> = const { RefCell::new([0.0; 3]) };
}

/// Read-and-reset the per-stage encoder timers (ffn, attn, conv ms). Test-only.
pub fn take_stage_ms() -> [f64; 3] {
    STAGE_MS.with(|s| std::mem::take(&mut *s.borrow_mut()))
}

// ---------------------------------------------------------------------------------------------
// Primitives
// ---------------------------------------------------------------------------------------------

/// One `nn.Linear`: `y = x·Wᵀ + b`, MLX/HF layout (`W` is `[n, k]` row-major). Prepacked GEMM.
pub(crate) struct Linear {
    pub(crate) w: Vec<f32>,
    pub(crate) b: Vec<f32>,
    pub(crate) n: usize,
    pub(crate) k: usize,
    packed: std::sync::OnceLock<PackedWeight>,
}

impl Linear {
    fn load(st: &LazySt, prefix: &str, n: usize, k: usize) -> Result<Self> {
        let w = st.tensor_f32(&format!("{prefix}.weight"))?;
        anyhow::ensure!(w.len() == n * k, "{prefix}.weight: {} != {n}x{k}", w.len());
        let b = if st.has(&format!("{prefix}.bias")) {
            st.tensor_f32(&format!("{prefix}.bias"))?
        } else {
            vec![0.0; n]
        };
        Ok(Self {
            w,
            b,
            n,
            k,
            packed: std::sync::OnceLock::new(),
        })
    }

    /// `x` is `[rows, k]` → `[rows, n]`.
    fn forward(&self, x: &[f32]) -> Vec<f32> {
        let m = x.len() / self.k;
        assert_eq!(x.len(), m * self.k, "lhs shape");
        let mut out = vec![0f32; m * self.n];
        let packed = self
            .packed
            .get_or_init(|| PackedWeight::new(&self.w, self.n, self.k));
        gemm_packed(&mut out, x, packed, m, Some(&self.b));
        out
    }
}

/// `nn.LayerNorm(d, eps=1e-5)`, in place over the last axis.
pub(crate) struct LayerNorm {
    pub(crate) w: Vec<f32>,
    pub(crate) b: Vec<f32>,
}

impl LayerNorm {
    fn load(st: &LazySt, prefix: &str) -> Result<Self> {
        Ok(Self {
            w: st.tensor_f32(&format!("{prefix}.weight"))?,
            b: st.tensor_f32(&format!("{prefix}.bias"))?,
        })
    }

    fn forward(&self, x: &[f32]) -> Vec<f32> {
        let d = self.w.len();
        let mut out = vec![0f32; x.len()];
        for (row, orow) in x.chunks_exact(d).zip(out.chunks_exact_mut(d)) {
            let mean = row.iter().sum::<f32>() / d as f32;
            let var = row.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / d as f32;
            let inv = 1.0 / (var + 1e-5).sqrt();
            for i in 0..d {
                orow[i] = (row[i] - mean) * inv * self.w[i] + self.b[i];
            }
        }
        out
    }
}

fn silu(v: f32) -> f32 {
    v / (1.0 + (-v).exp())
}

fn sigmoid(v: f32) -> f32 {
    1.0 / (1.0 + (-v).exp())
}

// ---------------------------------------------------------------------------------------------
// Config
// ---------------------------------------------------------------------------------------------

struct PreprocessorCfg {
    sample_rate: usize,
    win_length: usize,
    hop_length: usize,
    n_fft: usize,
    features: usize,
    preemph: f32,
}

struct EncoderCfg {
    n_layers: usize,
    d_model: usize,
    n_heads: usize,
    ff_expansion_factor: usize,
    subsampling_factor: usize,
    subsampling_conv_channels: usize,
    conv_kernel_size: usize,
    feat_in: usize,
}

struct Cfg {
    preprocessor: PreprocessorCfg,
    encoder: EncoderCfg,
    durations: Vec<usize>,
    max_symbols: Option<usize>,
    vocabulary: Vec<String>,
}

impl Cfg {
    fn parse(v: &serde_json::Value) -> Result<Self> {
        let u = |v: &serde_json::Value, k: &str| -> Result<usize> {
            v.get(k)
                .and_then(|x| x.as_u64())
                .map(|x| x as usize)
                .with_context(|| format!("config field {k}"))
        };
        let p = v.get("preprocessor").context("config.preprocessor")?;
        let e = v.get("encoder").context("config.encoder")?;
        Ok(Self {
            preprocessor: PreprocessorCfg {
                sample_rate: u(p, "sample_rate")?,
                win_length: u(p, "win_length")?,
                hop_length: u(p, "hop_length")?,
                n_fft: u(p, "n_fft")?,
                features: u(p, "features")?,
                preemph: p.get("preemph").and_then(|x| x.as_f64()).unwrap_or(0.97) as f32,
            },
            encoder: EncoderCfg {
                n_layers: u(e, "n_layers")?,
                d_model: u(e, "d_model")?,
                n_heads: u(e, "n_heads")?,
                ff_expansion_factor: u(e, "ff_expansion_factor")?,
                subsampling_factor: u(e, "subsampling_factor")?,
                subsampling_conv_channels: u(e, "subsampling_conv_channels")?,
                conv_kernel_size: u(e, "conv_kernel_size")?,
                feat_in: u(e, "feat_in")?,
            },
            durations: v["durations"]
                .as_array()
                .context("config.durations")?
                .iter()
                .map(|x| x.as_u64().unwrap_or(0) as usize)
                .collect(),
            max_symbols: v
                .get("max_symbols")
                .and_then(|x| x.as_u64())
                .map(|x| x as usize),
            vocabulary: v["vocabulary"]
                .as_array()
                .context("config.vocabulary")?
                .iter()
                .map(|x| x.as_str().unwrap_or_default().to_string())
                .collect(),
        })
    }
}

// ---------------------------------------------------------------------------------------------
// Encoder pieces
// ---------------------------------------------------------------------------------------------

/// One 3×3 (or 1×1) Conv2d of the dw-striding subsampler, MLX weight layout `[out, kh, kw, in/g]`,
/// NHWC activations `[h, w, c]`.
struct Conv2d {
    w: Vec<f32>,
    b: Vec<f32>,
    out_c: usize,
    in_per_group: usize,
    k: usize,
    stride: usize,
    pad: usize,
    depthwise: bool,
}

impl Conv2d {
    fn load(st: &LazySt, name: &str, stride: usize, depthwise: bool) -> Result<Self> {
        let shape = st.shape(&format!("{name}.weight"))?.to_vec();
        let (out_c, kh, _kw, in_per_group) = (shape[0], shape[1], shape[2], shape[3]);
        Ok(Self {
            w: st.tensor_f32(&format!("{name}.weight"))?,
            b: st.tensor_f32(&format!("{name}.bias"))?,
            out_c,
            in_per_group,
            k: kh,
            stride,
            pad: (kh - 1) / 2,
            depthwise,
        })
    }

    fn out_dim(&self, d: usize) -> usize {
        (d + 2 * self.pad - self.k) / self.stride + 1
    }

    /// `x` is `[h, w, c_in]` → `[h', w', out_c]`. Output ROWS are independent — computed across the
    /// pool (deterministic: each `oy` writes a disjoint row band).
    fn forward(&self, x: &[f32], h: usize, w: usize, c_in: usize) -> Vec<f32> {
        use rayon::prelude::*;
        let (ho, wo) = (self.out_dim(h), self.out_dim(w));
        let mut out = vec![0f32; ho * wo * self.out_c];
        out.par_chunks_mut(wo * self.out_c)
            .enumerate()
            .for_each(|(oy, orow_all)| {
                for ox in 0..wo {
                    let orow = &mut orow_all[ox * self.out_c..][..self.out_c];
                    for oc in 0..self.out_c {
                        let mut acc = self.b[oc];
                        for ky in 0..self.k {
                            let iy = (oy * self.stride + ky) as isize - self.pad as isize;
                            if iy < 0 || iy as usize >= h {
                                continue;
                            }
                            for kx in 0..self.k {
                                let ix = (ox * self.stride + kx) as isize - self.pad as isize;
                                if ix < 0 || ix as usize >= w {
                                    continue;
                                }
                                let base = ((iy as usize) * w + ix as usize) * c_in;
                                let wbase = ((oc * self.k + ky) * self.k + kx) * self.in_per_group;
                                if self.depthwise {
                                    acc += self.w[wbase] * x[base + oc];
                                } else {
                                    for ic in 0..self.in_per_group {
                                        acc += self.w[wbase + ic] * x[base + ic];
                                    }
                                }
                            }
                        }
                        orow[oc] = acc;
                    }
                }
            });
        out
    }
}

/// The conformer conv module: pointwise (GLU) → depthwise k9 → batchnorm → SiLU → pointwise.
pub(crate) struct ConvModule {
    pub(crate) pointwise1: Linear, // Conv1d k=1 ≡ Linear over channels
    pub(crate) dw: Vec<f32>,       // [k][d] — depthwise taps, transposed for contiguous channel ops
    pub(crate) bn_scale: Vec<f32>, // γ / sqrt(var + 1e-5)
    pub(crate) bn_shift: Vec<f32>, // β − mean·scale
    pub(crate) pointwise2: Linear,
    pub(crate) k: usize,
    pub(crate) d: usize,
}

impl ConvModule {
    fn load(st: &LazySt, prefix: &str, d: usize, k: usize) -> Result<Self> {
        // Conv1d k=1 weights are [out, 1, in] — a Linear in disguise.
        let pointwise1 = Linear::load(st, &format!("{prefix}.pointwise_conv1"), 2 * d, d)?;
        let pointwise2 = Linear::load(st, &format!("{prefix}.pointwise_conv2"), d, d)?;
        let dw_raw = st.tensor_f32(&format!("{prefix}.depthwise_conv.weight"))?; // [d, k, 1]
        // use_bias=false for this model — the fold below then reduces to the plain batchnorm shift
        let dw_b = if st.has(&format!("{prefix}.depthwise_conv.bias")) {
            st.tensor_f32(&format!("{prefix}.depthwise_conv.bias"))?
        } else {
            vec![0.0; d]
        };
        let mut dw = vec![0f32; k * d];
        for c in 0..d {
            for t in 0..k {
                dw[t * d + c] = dw_raw[c * k + t];
            }
        }
        let gamma = st.tensor_f32(&format!("{prefix}.batch_norm.weight"))?;
        let beta = st.tensor_f32(&format!("{prefix}.batch_norm.bias"))?;
        let mean = st.tensor_f32(&format!("{prefix}.batch_norm.running_mean"))?;
        let var = st.tensor_f32(&format!("{prefix}.batch_norm.running_var"))?;
        let bn_scale: Vec<f32> = (0..d).map(|i| gamma[i] / (var[i] + 1e-5).sqrt()).collect();
        // The depthwise bias folds into the batchnorm shift: bn((conv + b)) = conv·s + (b−mean)·s + β.
        let bn_shift: Vec<f32> = (0..d)
            .map(|i| (dw_b[i] - mean[i]) * bn_scale[i] + beta[i])
            .collect();
        Ok(Self {
            pointwise1,
            dw,
            bn_scale,
            bn_shift,
            pointwise2,
            k,
            d,
        })
    }

    fn forward(&self, x: &[f32], t: usize) -> Vec<f32> {
        let gated = self.pointwise1.forward(x); // [t, 2d]
        let conv = conv_glu_dw_bn_silu(
            &gated,
            &self.dw,
            &self.bn_scale,
            &self.bn_shift,
            t,
            self.d,
            self.k,
        );
        self.pointwise2.forward(&conv)
    }
}

/// The conformer conv module's non-GEMM middle, shared by the CPU forward and the GPU-kernel parity
/// oracle: GLU (`x[:d]·σ(x[d:])`) → symmetric depthwise conv (pad `(k−1)/2`) → folded batchnorm →
/// SiLU. `gated` is the pointwise1 output `[t, 2d]`; `dw` is `[k·d]` tap-major (`dw[tap·d+c]`);
/// `bn_scale`/`bn_shift` are `[d]`. Returns `[t, d]` (the pointwise2 input). The GPU kernels in
/// `parakeet_gpu.rs` reproduce this exactly.
pub fn conv_glu_dw_bn_silu(
    gated: &[f32],
    dw: &[f32],
    bn_scale: &[f32],
    bn_shift: &[f32],
    t: usize,
    d: usize,
    k: usize,
) -> Vec<f32> {
    let mut glu = vec![0f32; t * d];
    for i in 0..t {
        let row = &gated[i * 2 * d..];
        for c in 0..d {
            glu[i * d + c] = row[c] * sigmoid(row[d + c]);
        }
    }
    let pad = (k - 1) / 2;
    let mut conv = vec![0f32; t * d];
    for i in 0..t {
        let orow = &mut conv[i * d..][..d];
        for (tap, w_tap) in dw.chunks_exact(d).enumerate() {
            let j = i as isize + tap as isize - pad as isize;
            if j < 0 || j as usize >= t {
                continue;
            }
            let xrow = &glu[(j as usize) * d..][..d];
            for c in 0..d {
                orow[c] += w_tap[c] * xrow[c];
            }
        }
        for c in 0..d {
            orow[c] = silu(orow[c] * bn_scale[c] + bn_shift[c]);
        }
    }
    conv
}

pub(crate) struct FeedForward {
    pub(crate) linear1: Linear,
    pub(crate) linear2: Linear,
}

impl FeedForward {
    fn load(st: &LazySt, prefix: &str, d: usize, ff: usize) -> Result<Self> {
        Ok(Self {
            linear1: Linear::load(st, &format!("{prefix}.linear1"), ff, d)?,
            linear2: Linear::load(st, &format!("{prefix}.linear2"), d, ff)?,
        })
    }

    fn forward(&self, x: &[f32]) -> Vec<f32> {
        let mut h = self.linear1.forward(x);
        for v in h.iter_mut() {
            *v = silu(*v);
        }
        self.linear2.forward(&h)
    }
}

pub(crate) struct RelPosAttention {
    pub(crate) linear_q: Linear,
    pub(crate) linear_k: Linear,
    pub(crate) linear_v: Linear,
    pub(crate) linear_out: Linear,
    pub(crate) linear_pos: Linear,
    pub(crate) pos_bias_u: Vec<f32>, // [heads, head_dim]
    pub(crate) pos_bias_v: Vec<f32>,
    pub(crate) heads: usize,
    pub(crate) head_dim: usize,
}

impl RelPosAttention {
    fn load(st: &LazySt, prefix: &str, d: usize, heads: usize) -> Result<Self> {
        Ok(Self {
            linear_q: Linear::load(st, &format!("{prefix}.linear_q"), d, d)?,
            linear_k: Linear::load(st, &format!("{prefix}.linear_k"), d, d)?,
            linear_v: Linear::load(st, &format!("{prefix}.linear_v"), d, d)?,
            linear_out: Linear::load(st, &format!("{prefix}.linear_out"), d, d)?,
            linear_pos: Linear::load(st, &format!("{prefix}.linear_pos"), d, d)?,
            pos_bias_u: st.tensor_f32(&format!("{prefix}.pos_bias_u"))?,
            pos_bias_v: st.tensor_f32(&format!("{prefix}.pos_bias_v"))?,
            heads,
            head_dim: d / heads,
        })
    }

    /// `x` is `[t, d]`, `p` is `linear_pos(pos_emb)` `[2t−1, d]` (shared across layers per input).
    fn forward(&self, x: &[f32], p: &[f32], t: usize) -> Vec<f32> {
        let q = self.linear_q.forward(x);
        let k = self.linear_k.forward(x);
        let v = self.linear_v.forward(x);
        let out = relpos_attention(
            &q,
            &k,
            &v,
            p,
            &self.pos_bias_u,
            &self.pos_bias_v,
            self.heads,
            self.head_dim,
            t,
        );
        self.linear_out.forward(&out)
    }
}

/// Transformer-XL relative-position attention core, shared by the CPU forward and the GPU-kernel
/// parity oracle. `q`/`k`/`v` are `[t, d]`, `p = linear_pos(pos_emb)` is `[2t−1, d]`, the biases are
/// `[heads, head_dim]`. Returns the pre-`linear_out` context `[t, d]` (heads interleaved).
///
/// The rel-shift reduces to a plain index: `score[i,j]` reads `p[(t−1)−i+j]` (derived from the
/// pad-left-1 → reshape → drop-row shuffle) — no materialized `[t, 2t−1]` `bd` matrix or padded
/// buffer. The GPU kernel in `parakeet_gpu.rs` uses the SAME `m = (t−1)−i+j`.
#[allow(clippy::too_many_arguments)]
pub fn relpos_attention(
    q: &[f32],
    k: &[f32],
    v: &[f32],
    p: &[f32],
    pos_bias_u: &[f32],
    pos_bias_v: &[f32],
    heads: usize,
    head_dim: usize,
    t: usize,
) -> Vec<f32> {
    let (h, hd) = (heads, head_dim);
    let d = h * hd;
    let _plen = 2 * t - 1;
    let scale = 1.0 / (hd as f32).sqrt();
    use rayon::prelude::*;
    let head_out: Vec<Vec<f32>> = (0..h)
        .into_par_iter()
        .map(|head| {
            let u = &pos_bias_u[head * hd..][..hd];
            let vb = &pos_bias_v[head * hd..][..hd];
            let mut ohead = vec![0f32; t * hd];
            let mut srow = vec![0f32; t];
            for i in 0..t {
                let qi = &q[i * d + head * hd..][..hd];
                for j in 0..t {
                    let kj = &k[j * d + head * hd..][..hd];
                    let m = (t - 1) - i + j; // rel-shift index into p (always in [0, plen))
                    let pj = &p[m * d + head * hd..][..hd];
                    let mut ac = 0f32;
                    let mut bd = 0f32;
                    for c in 0..hd {
                        ac += (qi[c] + u[c]) * kj[c];
                        bd += (qi[c] + vb[c]) * pj[c];
                    }
                    srow[j] = (ac + bd) * scale;
                }
                let max = srow.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
                let mut sum = 0f32;
                for s in srow.iter_mut() {
                    *s = (*s - max).exp();
                    sum += *s;
                }
                let inv = 1.0 / sum;
                let orow = &mut ohead[i * hd..][..hd];
                for j in 0..t {
                    let w = srow[j] * inv;
                    let vj = &v[j * d + head * hd..][..hd];
                    for c in 0..hd {
                        orow[c] += w * vj[c];
                    }
                }
            }
            ohead
        })
        .collect();
    let mut out = vec![0f32; t * d];
    for (head, ohead) in head_out.iter().enumerate() {
        for i in 0..t {
            out[i * d + head * hd..][..hd].copy_from_slice(&ohead[i * hd..][..hd]);
        }
    }
    out
}

pub(crate) struct Block {
    pub(crate) norm_ff1: LayerNorm,
    pub(crate) ff1: FeedForward,
    pub(crate) norm_attn: LayerNorm,
    pub(crate) attn: RelPosAttention,
    pub(crate) norm_conv: LayerNorm,
    pub(crate) conv: ConvModule,
    pub(crate) norm_ff2: LayerNorm,
    pub(crate) ff2: FeedForward,
    pub(crate) norm_out: LayerNorm,
}

impl Block {
    fn load(st: &LazySt, prefix: &str, cfg: &EncoderCfg) -> Result<Self> {
        let (d, ff) = (cfg.d_model, cfg.d_model * cfg.ff_expansion_factor);
        Ok(Self {
            norm_ff1: LayerNorm::load(st, &format!("{prefix}.norm_feed_forward1"))?,
            ff1: FeedForward::load(st, &format!("{prefix}.feed_forward1"), d, ff)?,
            norm_attn: LayerNorm::load(st, &format!("{prefix}.norm_self_att"))?,
            attn: RelPosAttention::load(st, &format!("{prefix}.self_attn"), d, cfg.n_heads)?,
            norm_conv: LayerNorm::load(st, &format!("{prefix}.norm_conv"))?,
            conv: ConvModule::load(st, &format!("{prefix}.conv"), d, cfg.conv_kernel_size)?,
            norm_ff2: LayerNorm::load(st, &format!("{prefix}.norm_feed_forward2"))?,
            ff2: FeedForward::load(st, &format!("{prefix}.feed_forward2"), d, ff)?,
            norm_out: LayerNorm::load(st, &format!("{prefix}.norm_out"))?,
        })
    }

    fn forward(&self, x: &mut Vec<f32>, p: &[f32], t: usize) {
        let add = |x: &mut Vec<f32>, y: Vec<f32>, s: f32| {
            for (a, b) in x.iter_mut().zip(y) {
                *a += s * b;
            }
        };
        // OSFKB_PARAKEET_PROFILE=1 accumulates per-stage encoder time into thread-local counters.
        macro_rules! timed {
            ($slot:expr, $e:expr) => {{
                if PROFILE.with(|p| p.get()) {
                    let t0 = std::time::Instant::now();
                    let r = $e;
                    STAGE_MS.with(|s| s.borrow_mut()[$slot] += t0.elapsed().as_secs_f64() * 1000.0);
                    r
                } else {
                    $e
                }
            }};
        }
        let ff1 = timed!(0, self.ff1.forward(&self.norm_ff1.forward(x)));
        add(x, ff1, 0.5);
        let attn = timed!(1, self.attn.forward(&self.norm_attn.forward(x), p, t));
        add(x, attn, 1.0);
        let conv = timed!(2, self.conv.forward(&self.norm_conv.forward(x), t));
        add(x, conv, 1.0);
        let ff2 = timed!(0, self.ff2.forward(&self.norm_ff2.forward(x)));
        add(x, ff2, 0.5);
        *x = self.norm_out.forward(x);
    }
}

// ---------------------------------------------------------------------------------------------
// Decoder (prediction LSTM + joint)
// ---------------------------------------------------------------------------------------------

pub(crate) struct LstmLayer {
    pub(crate) wx: Vec<f32>, // [4H, I]
    pub(crate) wh: Vec<f32>, // [4H, H]
    pub(crate) bias: Vec<f32>,
    pub(crate) hidden: usize,
    pub(crate) input: usize,
}

impl LstmLayer {
    fn load(st: &LazySt, prefix: &str, input: usize, hidden: usize) -> Result<Self> {
        Ok(Self {
            wx: st.tensor_f32(&format!("{prefix}.Wx"))?,
            wh: st.tensor_f32(&format!("{prefix}.Wh"))?,
            bias: st.tensor_f32(&format!("{prefix}.bias"))?,
            hidden,
            input,
        })
    }

    /// One step. MLX gate order is i, f, g, o; a `None` initial state is h=0 (no Wh term) and
    /// `c = i∘g` (equivalently c₀ = 0).
    fn step(&self, x: &[f32], state: Option<&(Vec<f32>, Vec<f32>)>) -> (Vec<f32>, Vec<f32>) {
        let hn = self.hidden;
        let mut gates = self.bias.clone();
        for (r, g) in gates.iter_mut().enumerate() {
            let wrow = &self.wx[r * self.input..][..self.input];
            let mut acc = 0f32;
            for c in 0..self.input {
                acc += wrow[c] * x[c];
            }
            *g += acc;
            if let Some((h, _)) = state {
                let wrow = &self.wh[r * hn..][..hn];
                let mut acc = 0f32;
                for c in 0..hn {
                    acc += wrow[c] * h[c];
                }
                *g += acc;
            }
        }
        let mut h_out = vec![0f32; hn];
        let mut c_out = vec![0f32; hn];
        for j in 0..hn {
            let i = sigmoid(gates[j]);
            let f = sigmoid(gates[hn + j]);
            let g = gates[2 * hn + j].tanh();
            let o = sigmoid(gates[3 * hn + j]);
            let c = match state {
                Some((_, cell)) => f * cell[j] + i * g,
                None => i * g,
            };
            c_out[j] = c;
            h_out[j] = o * c.tanh();
        }
        (h_out, c_out)
    }
}

/// One decode-side LSTM state: (h, c) per layer.
pub type DecoderState = Vec<(Vec<f32>, Vec<f32>)>;

/// One emitted token, in ENCODER-FRAME units (the caller owns the seconds conversion).
#[derive(Clone, Debug)]
pub struct TokenEvent {
    pub id: usize,
    pub step: usize,
    pub duration_frames: usize,
    pub text: String,
}

/// One greedy-decode iteration, for the parity trace: what was argmaxed where.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TraceStep {
    pub step: usize,
    pub token: usize,
    pub duration: usize,
}

// ---------------------------------------------------------------------------------------------
// The model
// ---------------------------------------------------------------------------------------------

pub struct Parakeet {
    // frontend
    filterbanks: Vec<f32>, // [features, n_fft/2+1]
    window: Vec<f32>,      // [win_length]
    dft_cos: Vec<f32>,     // [n_fft/2+1, win_length] — window folded out of the frame loop
    dft_sin: Vec<f32>,
    pre: PreprocessorCfg,
    // subsampler
    sub_convs: Vec<Conv2d>,
    pre_out: Linear,
    // conformer
    pub(crate) layers: Vec<Block>,
    pub(crate) d_model: usize,
    subsampling: usize,
    // decoder
    pub(crate) embed: Vec<f32>, // [vocab+1, pred_hidden]
    pub(crate) lstm: Vec<LstmLayer>,
    pub(crate) joint_enc: Linear,
    pub(crate) joint_pred: Linear,
    pub(crate) joint_out: Linear,
    pub(crate) pred_hidden: usize,
    pub vocabulary: Vec<String>,
    pub durations: Vec<usize>,
    pub max_symbols: Option<usize>,
}

impl Parakeet {
    /// Seconds per encoder frame (`time_ratio`): subsampling × hop / sample-rate.
    pub fn time_ratio(&self) -> f64 {
        // The reference computes `subsampling / sample_rate * hop` — keep ITS operation order so
        // downstream `step * time_ratio` timestamps are bit-identical (they are serialized).
        self.subsampling as f64 / self.pre.sample_rate as f64 * self.pre.hop_length as f64
    }

    pub fn load(dir: &Path) -> Result<Self> {
        let config = std::fs::read(dir.join("config.json")).context("parakeet config.json")?;
        let st = LazySt::open(dir)?;
        Self::from_st(&config, st)
    }

    /// [`Self::load`] depuis des octets — le chemin navigateur/wasm (OPFS/fetch).
    pub fn from_bytes(config: &[u8], model: Vec<u8>) -> Result<Self> {
        Self::from_st(config, LazySt::from_bytes(vec![model])?)
    }

    /// [`Self::load`] depuis un [`LazySt`] arbitraire — avec [`LazySt::from_readers`], le
    /// checkpoint reste HORS mémoire (OPFS par plages) : le pic n'est plus « blob 1,25 Go +
    /// tenseurs f32 » mais les tenseurs f32 seuls, ce qui fait tenir Parakeet dans la limite
    /// wasm32 de 4 Go.
    pub fn from_lazyst(config: &[u8], st: LazySt) -> Result<Self> {
        Self::from_st(config, st)
    }

    fn from_st(config: &[u8], st: LazySt) -> Result<Self> {
        let cfg = Cfg::parse(&serde_json::from_slice::<serde_json::Value>(config)?)?;
        let pre = cfg.preprocessor;
        let enc = cfg.encoder;

        let filterbanks = st.tensor_f32("preprocessor.filterbanks")?;
        let window = st.tensor_f32("preprocessor.window")?;
        let bins = pre.n_fft / 2 + 1;
        // Frames are win_length real samples zero-padded to n_fft — bake the window into the DFT
        // tables so a frame is two [bins, win] mat-vecs.
        let mut dft_cos = vec![0f32; bins * pre.win_length];
        let mut dft_sin = vec![0f32; bins * pre.win_length];
        for b in 0..bins {
            for j in 0..pre.win_length {
                let ang = 2.0 * PI * b as f64 * j as f64 / pre.n_fft as f64;
                dft_cos[b * pre.win_length + j] = (ang.cos() * window[j] as f64) as f32;
                dft_sin[b * pre.win_length + j] = (ang.sin() * window[j] as f64) as f32;
            }
        }

        // dw-striding subsampler: conv.0 full, then (depthwise, pointwise) pairs — module-list
        // indices 0, 2, 3, 5, 6 (ReLUs hold 1, 4, 7).
        let n_stages = (enc.subsampling_factor as f64).log2() as usize;
        let mut sub_convs = vec![Conv2d::load(&st, "encoder.pre_encode.conv.0", 2, false)?];
        let mut idx = 2;
        for _ in 0..n_stages - 1 {
            sub_convs.push(Conv2d::load(
                &st,
                &format!("encoder.pre_encode.conv.{idx}"),
                2,
                true,
            )?);
            sub_convs.push(Conv2d::load(
                &st,
                &format!("encoder.pre_encode.conv.{}", idx + 1),
                1,
                false,
            )?);
            idx += 3;
        }
        let mut freq = enc.feat_in;
        for _ in 0..n_stages {
            freq = (freq + 2 - 3) / 2 + 1;
        }
        let pre_out = Linear::load(
            &st,
            "encoder.pre_encode.out",
            enc.d_model,
            enc.subsampling_conv_channels * freq,
        )?;

        let layers = (0..enc.n_layers)
            .map(|i| Block::load(&st, &format!("encoder.layers.{i}"), &enc))
            .collect::<Result<Vec<_>>>()?;

        let pred_hidden = st.shape("decoder.prediction.embed.weight")?[1];
        let embed = st.tensor_f32("decoder.prediction.embed.weight")?;
        let lstm = (0..2)
            .map(|i| {
                LstmLayer::load(
                    &st,
                    &format!("decoder.prediction.dec_rnn.lstm.{i}"),
                    pred_hidden,
                    pred_hidden,
                )
            })
            .collect::<Result<Vec<_>>>()?;
        let joint_hidden = st.shape("joint.enc.weight")?[0];
        let n_out = st.shape("joint.joint_net.2.weight")?[0];
        Ok(Self {
            filterbanks,
            window,
            dft_cos,
            dft_sin,
            pre,
            sub_convs,
            pre_out,
            layers,
            d_model: enc.d_model,
            subsampling: enc.subsampling_factor,
            embed,
            lstm,
            joint_enc: Linear::load(&st, "joint.enc", joint_hidden, enc.d_model)?,
            joint_pred: Linear::load(&st, "joint.pred", joint_hidden, pred_hidden)?,
            joint_out: Linear::load(&st, "joint.joint_net.2", n_out, joint_hidden)?,
            pred_hidden,
            vocabulary: cfg.vocabulary,
            durations: cfg.durations,
            max_symbols: cfg.max_symbols,
        })
    }

    /// `get_logmel`: 16 kHz mono f32 → `[frames, features]` (row-major), normalized per feature.
    pub fn logmel(&self, samples: &[f32]) -> (Vec<f32>, usize) {
        let p = &self.pre;
        // pre-emphasis
        let mut x = Vec::with_capacity(samples.len());
        x.push(samples[0]);
        for i in 1..samples.len() {
            x.push(samples[i] - p.preemph * samples[i - 1]);
        }
        // reflect pad n_fft/2 both sides
        let pad = p.n_fft / 2;
        let n = x.len();
        let mut padded = Vec::with_capacity(n + 2 * pad);
        padded.extend(x[1..=pad].iter().rev());
        padded.extend_from_slice(&x);
        padded.extend(x[n - pad - 1..n - 1].iter().rev());

        let frames = (padded.len() - p.win_length + p.hop_length) / p.hop_length;
        let bins = p.n_fft / 2 + 1;
        let features = p.features;

        // spec[b, t] = (|Re| + |Im|)² — the reference's L1-magnitude quirk (module note)
        let mut spec = vec![0f32; bins * frames];
        for t in 0..frames {
            let frame = &padded[t * p.hop_length..][..p.win_length];
            for b in 0..bins {
                let cos = &self.dft_cos[b * p.win_length..][..p.win_length];
                let sin = &self.dft_sin[b * p.win_length..][..p.win_length];
                let mut re = 0f64;
                let mut im = 0f64;
                for j in 0..p.win_length {
                    re += frame[j] as f64 * cos[j] as f64;
                    im -= frame[j] as f64 * sin[j] as f64;
                }
                let mag = (re.abs() + im.abs()) as f32;
                spec[b * frames + t] = mag * mag;
            }
        }
        // mel = log(fb·spec + 1e-5), then per-feature (population) mean/std over time
        let mut mel = vec![0f32; features * frames];
        for m in 0..features {
            let fb = &self.filterbanks[m * bins..][..bins];
            for t in 0..frames {
                let mut acc = 0f32;
                for b in 0..bins {
                    acc += fb[b] * spec[b * frames + t];
                }
                mel[m * frames + t] = (acc + 1e-5).ln();
            }
        }
        let mut out = vec![0f32; frames * features];
        for m in 0..features {
            let row = &mel[m * frames..][..frames];
            let mean = row.iter().sum::<f32>() / frames as f32;
            let var = row.iter().map(|v| (v - mean) * (v - mean)).sum::<f32>() / frames as f32;
            let inv = 1.0 / (var.sqrt() + 1e-5);
            for t in 0..frames {
                out[t * features + m] = (row[t] - mean) * inv;
            }
        }
        (out, frames)
    }

    /// The dw-striding subsampler on `[t, feat, 1]` NHWC. ReLU placement matters: the module list
    /// is `[conv, ReLU, dw, pw, ReLU, dw, pw, ReLU]` — there is NO ReLU between a depthwise and
    /// its pointwise, so the loop must not clamp after the depthwise stages.
    fn forward_subsample(&self, mel: &[f32], t: usize) -> (Vec<f32>, usize) {
        let mut h = t;
        let mut w = self.pre.features;
        let mut c = 1usize;
        let mut x = mel.to_vec();
        for (i, conv) in self.sub_convs.iter().enumerate() {
            x = conv.forward(&x, h, w, c);
            h = conv.out_dim(h);
            w = conv.out_dim(w);
            c = conv.out_c;
            // ReLU after the full conv (index 0) and after each POINTWISE (odd module pairs) —
            // never between depthwise and pointwise.
            let is_dw = i > 0 && i % 2 == 1;
            if !is_dw {
                for v in x.iter_mut() {
                    *v = v.max(0.0);
                }
            }
        }
        // [t', freq', C] → channel-major flatten [t', C·freq'] (the reference swaps C ahead of freq)
        let (tp, fp) = (h, w);
        let mut flat = vec![0f32; tp * c * fp];
        for ti in 0..tp {
            for fi in 0..fp {
                for ci in 0..c {
                    flat[ti * c * fp + ci * fp + fi] = x[(ti * fp + fi) * c + ci];
                }
            }
        }
        (self.pre_out.forward(&flat), tp)
    }

    /// Full encoder: mel → conformer features.
    pub fn forward_encoder(&self, mel: &[f32], t: usize) -> (Vec<f32>, usize) {
        let (mut x, pe, t_enc) = self.subsample_and_posemb(mel, t);
        for layer in &self.layers {
            let p = layer.attn.linear_pos.forward(&pe);
            layer.forward(&mut x, &p, t_enc);
        }
        (x, t_enc)
    }

    /// The pre-conformer work the GPU path shares: dw-striding subsampler → `x` `[t_enc, d]`, plus
    /// the relative positional encoding `pe` `[2·t_enc−1, d]` (positions `t_enc−1 … −(t_enc−1)`).
    pub(crate) fn subsample_and_posemb(
        &self,
        mel: &[f32],
        t: usize,
    ) -> (Vec<f32>, Vec<f32>, usize) {
        let (x, t_enc) = self.forward_subsample(mel, t);
        let d = self.d_model;
        let plen = 2 * t_enc - 1;
        let mut pe = vec![0f32; plen * d];
        for (i, row) in pe.chunks_exact_mut(d).enumerate() {
            let pos = (t_enc as isize - 1 - i as isize) as f64;
            for j in 0..d / 2 {
                let div = (-(2.0 * j as f64) * (10000f64).ln() / d as f64).exp();
                row[2 * j] = (pos * div).sin() as f32;
                row[2 * j + 1] = (pos * div).cos() as f32;
            }
        }
        (x, pe, t_enc)
    }

    fn decoder_step(
        &self,
        last_token: Option<usize>,
        state: Option<&DecoderState>,
    ) -> (Vec<f32>, DecoderState) {
        let x0 = match last_token {
            Some(id) => self.embed[id * self.pred_hidden..][..self.pred_hidden].to_vec(),
            None => vec![0f32; self.pred_hidden], // zeros, NOT the blank embedding
        };
        let mut x = x0;
        let mut new_state = Vec::with_capacity(self.lstm.len());
        for (i, layer) in self.lstm.iter().enumerate() {
            let (ho, co) = layer.step(&x, state.map(|s| &s[i]));
            x = ho.clone();
            new_state.push((ho, co));
        }
        (x, new_state)
    }

    /// Greedy TDT decode over encoder features. Returns the emitted tokens and, for parity, the
    /// full per-iteration trace.
    pub fn decode_greedy(
        &self,
        features: &[f32],
        t_enc: usize,
    ) -> (Vec<TokenEvent>, Vec<TraceStep>) {
        let n_vocab = self.vocabulary.len(); // blank = n_vocab
        let enc_proj = self.joint_enc.forward(features); // [t_enc, joint_hidden]
        let jh = self.joint_enc.n;

        let mut tokens = Vec::new();
        let mut trace = Vec::new();
        let mut last_token: Option<usize> = None;
        let mut state: Option<DecoderState> = None;
        let mut cached: Option<(Option<usize>, Vec<f32>, DecoderState)> = None;
        let mut step = 0usize;
        let mut new_symbols = 0usize;

        while step < t_enc {
            // the decoder output only changes when last_token/state change — cache it
            let (dec_out, dec_state) = match &cached {
                Some((tok, out, st)) if *tok == last_token => (out.clone(), st.clone()),
                _ => {
                    let (out, st) = self.decoder_step(last_token, state.as_ref());
                    cached = Some((last_token, out.clone(), st.clone()));
                    (out, st)
                }
            };
            let pred_proj = self.joint_pred.forward(&dec_out); // [jh]
            let mut hidden = vec![0f32; jh];
            for i in 0..jh {
                hidden[i] = (enc_proj[step * jh + i] + pred_proj[i]).max(0.0);
            }
            let logits = self.joint_out.forward(&hidden); // [n_vocab + 1 + n_durations]

            let mut pred_token = 0usize;
            let mut best = f32::NEG_INFINITY;
            for (i, &v) in logits[..n_vocab + 1].iter().enumerate() {
                if v > best {
                    best = v;
                    pred_token = i;
                }
            }
            let mut decision = 0usize;
            let mut best_d = f32::NEG_INFINITY;
            for (i, &v) in logits[n_vocab + 1..].iter().enumerate() {
                if v > best_d {
                    best_d = v;
                    decision = i;
                }
            }
            let duration = self.durations[decision];
            trace.push(TraceStep {
                step,
                token: pred_token,
                duration,
            });

            if pred_token != n_vocab {
                tokens.push(TokenEvent {
                    id: pred_token,
                    step,
                    duration_frames: duration,
                    text: self.vocabulary[pred_token].replace('▁', " "),
                });
                last_token = Some(pred_token);
                state = Some(dec_state);
                cached = None; // token changed → recompute next round
            }
            step += duration;
            new_symbols += 1;
            if duration != 0 {
                new_symbols = 0;
            } else if let Some(max) = self.max_symbols
                && max <= new_symbols
            {
                step += 1;
                new_symbols = 0;
            }
        }
        (tokens, trace)
    }

    /// Full pipeline: 16 kHz mono f32 samples → tokens (frame units) + the joined text.
    pub fn transcribe(&self, samples_16k: &[f32]) -> (Vec<TokenEvent>, String) {
        let (mel, t) = self.logmel(samples_16k);
        let (features, t_enc) = self.forward_encoder(&mel, t);
        let (tokens, _) = self.decode_greedy(&features, t_enc);
        let text: String = tokens.iter().map(|t| t.text.as_str()).collect();
        (tokens, text.trim().to_string())
    }
}