inferencelayer 0.2.10

//! OpenAI Whisper — speech-to-text, CPU f32 (the parity oracle and the portable fallback; the wgpu
//! path reuses the encoder's kernels, since Whisper's encoder is a VANILLA bidirectional transformer
//! — plain scaled-dot attention + GELU MLP, no rel-pos, no conv module).
//!
//! This module starts with the log-mel frontend and the encoder (conv stem + N transformer blocks),
//! parity-gated against `transformers`'s `WhisperModel.encoder` via `tests/fixtures/export_whisper.py`.
//! The decoder (causal self-attention + cross-attention + autoregressive generation) and the
//! tokenizer come next.
//!
//! Whisper specifics that are contract, not choices:
//! * mel: pad/trim to 30 s (480 000 samples), STFT n_fft 400 / hop 160 / **periodic** hann, POWER
//!   spectrum, the checkpoint's mel filterbank, then `log10(max(x, 1e-10))`, floor at `max − 8`, and
//!   `(x + 4) / 4`. Exactly `WhisperFeatureExtractor`.
//! * the conv stem is conv1 (k3 s1 pad1) → GELU → conv2 (k3 s2 pad1) → GELU, halving 3000 → 1500.
//! * attention has NO bias on `k_proj` (q/v/out do); q is scaled by `head_dim^-0.5`.
//! * GELU is the EXACT (erf) gelu.

use std::path::Path;

use anyhow::{Context, Result};

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

const SR: usize = 16000;
const N_FFT: usize = 400;
const HOP: usize = 160;
const CHUNK_SAMPLES: usize = 30 * SR; // 480 000
/// Overlap between long-form windows. Long enough that a sentence clause straddling a boundary
/// appears whole in one of the two windows, short enough not to waste much compute.
pub const WINDOW_OVERLAP_SAMPLES: usize = 5 * SR;
const N_FRAMES: usize = CHUNK_SAMPLES / HOP; // 3000

/// `y = x·Wᵀ + b`, `[n, k]` row-major (HF Linear). Optional bias.
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, bias: bool) -> 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 bias {
            st.tensor_f32(&format!("{prefix}.bias"))?
        } else {
            vec![0.0; n]
        };
        Ok(Self {
            w,
            b,
            n,
            k,
            packed: std::sync::OnceLock::new(),
        })
    }

    fn forward(&self, x: &[f32]) -> Vec<f32> {
        let m = x.len() / self.k;
        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
    }
}

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
    }
}

/// Exact (erf) GELU, matching HF's `"gelu"`.
fn gelu(v: f32) -> f32 {
    0.5 * v * (1.0 + libm::erff(v * std::f32::consts::FRAC_1_SQRT_2))
}

/// One Whisper attention block (bidirectional; q scaled, k_proj biasless).
pub(crate) struct Attention {
    pub(crate) q: Linear,
    pub(crate) k: Linear,
    pub(crate) v: Linear,
    pub(crate) out: Linear,
    pub(crate) heads: usize,
    pub(crate) hd: usize,
}

impl Attention {
    fn load(st: &LazySt, prefix: &str, d: usize, heads: usize) -> Result<Self> {
        Ok(Self {
            q: Linear::load(st, &format!("{prefix}.q_proj"), d, d, true)?,
            k: Linear::load(st, &format!("{prefix}.k_proj"), d, d, false)?, // NO bias
            v: Linear::load(st, &format!("{prefix}.v_proj"), d, d, true)?,
            out: Linear::load(st, &format!("{prefix}.out_proj"), d, d, true)?,
            heads,
            hd: d / heads,
        })
    }

    /// Bidirectional self-attention over `x` `[t, d]`.
    fn forward(&self, x: &[f32], t: usize) -> Vec<f32> {
        let k = self.k.forward(x);
        let v = self.v.forward(x);
        self.forward_qkv(x, &k, &v, t, t, false)
    }

    /// Project the K/V of an input once (for cross-attention, whose encoder K/V are step-invariant).
    fn project_kv(&self, kv_in: &[f32]) -> (Vec<f32>, Vec<f32>) {
        (self.k.forward(kv_in), self.v.forward(kv_in))
    }

    /// Attention with a caller-supplied Q input and PRE-PROJECTED K/V. `causal` masks j > i (for the
    /// decoder self-attention). `tq`/`tk` are the query/key lengths.
    fn forward_qkv(
        &self,
        q_in: &[f32],
        k: &[f32],
        v: &[f32],
        tq: usize,
        tk: usize,
        causal: bool,
    ) -> Vec<f32> {
        let (h, hd) = (self.heads, self.hd);
        let d = h * hd;
        let q = self.q.forward(q_in);
        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 mut ohead = vec![0f32; tq * hd];
                let mut srow = vec![0f32; tk];
                for i in 0..tq {
                    let qi = &q[i * d + head * hd..][..hd];
                    let jmax = if causal { i + 1 } else { tk };
                    for j in 0..jmax {
                        let kj = &k[j * d + head * hd..][..hd];
                        let mut acc = 0f32;
                        for c in 0..hd {
                            acc += qi[c] * kj[c];
                        }
                        srow[j] = acc * scale;
                    }
                    let max = srow[..jmax]
                        .iter()
                        .cloned()
                        .fold(f32::NEG_INFINITY, f32::max);
                    let mut sum = 0f32;
                    for s in srow[..jmax].iter_mut() {
                        *s = (*s - max).exp();
                        sum += *s;
                    }
                    let inv = 1.0 / sum;
                    let orow = &mut ohead[i * hd..][..hd];
                    for j in 0..jmax {
                        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 ctx = vec![0f32; tq * d];
        for (head, oh) in head_out.iter().enumerate() {
            for i in 0..tq {
                ctx[i * d + head * hd..][..hd].copy_from_slice(&oh[i * hd..][..hd]);
            }
        }
        self.out.forward(&ctx)
    }

    /// Causal self-attention against a growing KV cache: query `i` sits at absolute position
    /// `q_off + i` and attends to keys `0..=(q_off + i)`. `q_in` is the NEW tokens `[tq, d]`; `k`/`v`
    /// are the FULL cache `[tk, d]` (tk = q_off + tq). With q_off=0 and tq=tk this is exactly the
    /// non-cached causal path — the cache just avoids recomputing earlier steps' K/V. O(n²)→O(n).
    fn forward_qkv_cached(
        &self,
        q_in: &[f32],
        k: &[f32],
        v: &[f32],
        tq: usize,
        tk: usize,
        q_off: usize,
    ) -> Vec<f32> {
        let (h, hd) = (self.heads, self.hd);
        let d = h * hd;
        let q = self.q.forward(q_in);
        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 mut ohead = vec![0f32; tq * hd];
                let mut srow = vec![0f32; tk];
                for i in 0..tq {
                    let qi = &q[i * d + head * hd..][..hd];
                    let jmax = (q_off + i + 1).min(tk);
                    for j in 0..jmax {
                        let kj = &k[j * d + head * hd..][..hd];
                        let mut acc = 0f32;
                        for c in 0..hd {
                            acc += qi[c] * kj[c];
                        }
                        srow[j] = acc * scale;
                    }
                    let max = srow[..jmax]
                        .iter()
                        .cloned()
                        .fold(f32::NEG_INFINITY, f32::max);
                    let mut sum = 0f32;
                    for s in srow[..jmax].iter_mut() {
                        *s = (*s - max).exp();
                        sum += *s;
                    }
                    let inv = 1.0 / sum;
                    let orow = &mut ohead[i * hd..][..hd];
                    for j in 0..jmax {
                        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 ctx = vec![0f32; tq * d];
        for (head, oh) in head_out.iter().enumerate() {
            for i in 0..tq {
                ctx[i * d + head * hd..][..hd].copy_from_slice(&oh[i * hd..][..hd]);
            }
        }
        self.out.forward(&ctx)
    }
}

pub(crate) struct EncBlock {
    pub(crate) norm_attn: LayerNorm,
    pub(crate) attn: Attention,
    pub(crate) norm_ff: LayerNorm,
    pub(crate) fc1: Linear,
    pub(crate) fc2: Linear,
}

impl EncBlock {
    fn forward(&self, x: &mut [f32], t: usize) {
        let a = self.attn.forward(&self.norm_attn.forward(x), t);
        for (xi, ai) in x.iter_mut().zip(a) {
            *xi += ai;
        }
        let mut h = self.fc1.forward(&self.norm_ff.forward(x));
        for v in h.iter_mut() {
            *v = gelu(*v);
        }
        let f = self.fc2.forward(&h);
        for (xi, fi) in x.iter_mut().zip(f) {
            *xi += fi;
        }
    }
}

/// A whole Conv1d `[out, in, k]` + bias, stride `s`, symmetric pad 1. `x` is `[t, in]` → `[t', out]`.
struct Conv1d {
    w: Vec<f32>,
    b: Vec<f32>,
    out_c: usize,
    in_c: usize,
    stride: usize,
}

impl Conv1d {
    fn load(st: &LazySt, prefix: &str, out_c: usize, in_c: usize, stride: usize) -> Result<Self> {
        Ok(Self {
            w: st.tensor_f32(&format!("{prefix}.weight"))?, // [out, in, 3]
            b: st.tensor_f32(&format!("{prefix}.bias"))?,
            out_c,
            in_c,
            stride,
        })
    }

    /// k=3, pad=1. `x` is `[t, in]` → `[ceil after stride, out]`.
    fn forward(&self, x: &[f32], t: usize) -> (Vec<f32>, usize) {
        let k = 3;
        let pad = 1;
        let to = (t + 2 * pad - k) / self.stride + 1;
        use rayon::prelude::*;
        let mut out = vec![0f32; to * self.out_c];
        out.par_chunks_mut(self.out_c)
            .enumerate()
            .for_each(|(i, orow)| {
                for oc in 0..self.out_c {
                    let mut acc = self.b[oc];
                    let wbase = oc * self.in_c * k;
                    for tap in 0..k {
                        let j = (i * self.stride + tap) as isize - pad as isize;
                        if j < 0 || j as usize >= t {
                            continue;
                        }
                        let xrow = &x[(j as usize) * self.in_c..][..self.in_c];
                        for ic in 0..self.in_c {
                            acc += self.w[wbase + ic * k + tap] * xrow[ic];
                        }
                    }
                    orow[oc] = acc;
                }
            });
        (out, to)
    }
}

struct Cfg {
    d_model: usize,
    encoder_layers: usize,
    encoder_attention_heads: usize,
    num_mel_bins: usize,
}

pub struct WhisperEncoder {
    mel_filters: Vec<f32>, // [n_mels, n_fft/2+1]
    window: Vec<f32>,      // periodic hann [n_fft]
    dft_cos: Vec<f64>,     // [bins, n_fft] DFT basis — precomputed to kill per-term cos/sin
    dft_sin: Vec<f64>,     // [bins, n_fft]
    conv1: Conv1d,
    conv2: Conv1d,
    pos: Vec<f32>, // embed_positions [1500, d]
    pub(crate) blocks: Vec<EncBlock>,
    /// Absent when the host model replaced it with Identity — see [`Self::load_embedded`].
    pub(crate) ln_post: Option<LayerNorm>,
    n_mels: usize,
    pub(crate) d: usize,
}

impl WhisperEncoder {
    pub fn load(dir: &Path) -> Result<Self> {
        Self::load_prefixed(dir, "model.encoder")
    }

    /// Load the encoder from an arbitrary tensor prefix.
    ///
    /// The 2026 ASR leaderboard is dominated by "audio encoder + LLM decoder" models, and several
    /// of them embed a stock Whisper encoder under their own namespace — ARK-ASR-3B keeps it at
    /// `audio_encoder.whisper.encoder.*`. Parameterizing the prefix is what lets those reuse this
    /// encoder (and its wgpu path) instead of each shipping a copy.
    pub fn load_prefixed(dir: &Path, prefix: &str) -> Result<Self> {
        Self::load_embedded(dir, prefix, None)
    }

    /// Load an encoder embedded in a larger model.
    ///
    /// `section` names a nested config object holding the ENCODER's dimensions — ARK-ASR-3B keeps
    /// them under `whisper_config` while the top level describes its LLM, so reading the top level
    /// would silently build a 2048-wide encoder for a 1280-wide checkpoint.
    ///
    /// The final `layer_norm` is loaded only if the checkpoint HAS one. Host models routinely
    /// replace it with Identity and apply their own norm before their adapter (ARK does exactly
    /// this); applying ours as well would quietly corrupt every embedding it produces.
    pub fn load_embedded(dir: &Path, prefix: &str, section: Option<&str>) -> Result<Self> {
        let root: serde_json::Value =
            serde_json::from_slice(&std::fs::read(dir.join("config.json"))?).context("config")?;
        let st = LazySt::open(dir)?;
        Self::from_st(&root, &st, prefix, section)
    }

    /// [`Self::load_embedded`] sans système de fichiers — config déjà parsée + checkpoint déjà
    /// ouvert (le chemin navigateur/wasm, où les octets viennent d'OPFS ou d'un fetch).
    pub fn from_st(
        root: &serde_json::Value,
        st: &LazySt,
        prefix: &str,
        section: Option<&str>,
    ) -> Result<Self> {
        let v = match section {
            Some(k) => root.get(k).cloned().unwrap_or(root.clone()),
            None => root.clone(),
        };
        let g = |k: &str| v.get(k).and_then(|x| x.as_u64()).unwrap_or(0) as usize;
        let cfg = Cfg {
            d_model: g("d_model"),
            encoder_layers: g("encoder_layers"),
            encoder_attention_heads: g("encoder_attention_heads"),
            num_mel_bins: g("num_mel_bins"),
        };
        let d = cfg.d_model;
        let window: Vec<f32> = (0..N_FFT)
            .map(|i| {
                (0.5 - 0.5 * (2.0 * std::f64::consts::PI * i as f64 / N_FFT as f64).cos()) as f32
            })
            .collect();
        // DFT basis [bins, n_fft], the SAME f64 cos/sin the per-term loop used to recompute — hoisted
        // to load so `logmel` is a couple of GEMMs, not 240M transcendental evaluations.
        let bins = N_FFT / 2 + 1;
        let mut dft_cos = vec![0f64; bins * N_FFT];
        let mut dft_sin = vec![0f64; bins * N_FFT];
        for b in 0..bins {
            for j in 0..N_FFT {
                let ang = 2.0 * std::f64::consts::PI * b as f64 * j as f64 / N_FFT as f64;
                dft_cos[b * N_FFT + j] = ang.cos();
                dft_sin[b * N_FFT + j] = ang.sin();
            }
        }
        let blocks = (0..cfg.encoder_layers)
            .map(|i| {
                let p = format!("{prefix}.layers.{i}");
                Ok(EncBlock {
                    norm_attn: LayerNorm::load(&st, &format!("{p}.self_attn_layer_norm"))?,
                    attn: Attention::load(
                        &st,
                        &format!("{p}.self_attn"),
                        d,
                        cfg.encoder_attention_heads,
                    )?,
                    norm_ff: LayerNorm::load(&st, &format!("{p}.final_layer_norm"))?,
                    fc1: Linear::load(&st, &format!("{p}.fc1"), g("encoder_ffn_dim"), d, true)?,
                    fc2: Linear::load(&st, &format!("{p}.fc2"), d, g("encoder_ffn_dim"), true)?,
                })
            })
            .collect::<Result<Vec<_>>>()?;
        Ok(Self {
            // Prefer the checkpoint's own table when an export supplied one; otherwise BUILD it.
            // Generating it is what lets a stock HF checkpoint load with no Python export step —
            // see `mel_filter_bank`.
            mel_filters: st
                .tensor_f32("mel_filters")
                .unwrap_or_else(|_| mel_filter_bank(cfg.num_mel_bins, N_FFT, SR)),
            window,
            dft_cos,
            dft_sin,
            conv1: Conv1d::load(&st, &format!("{prefix}.conv1"), d, cfg.num_mel_bins, 1)?,
            conv2: Conv1d::load(&st, &format!("{prefix}.conv2"), d, d, 2)?,
            pos: st.tensor_f32(&format!("{prefix}.embed_positions.weight"))?,
            blocks,
            ln_post: LayerNorm::load(&st, &format!("{prefix}.layer_norm")).ok(),
            n_mels: cfg.num_mel_bins,
            d,
        })
    }

    /// 16 kHz mono f32 → log-mel `[n_mels, 3000]` (pad/trim to 30 s), matching WhisperFeatureExtractor.
    pub fn logmel(&self, samples: &[f32]) -> Vec<f32> {
        let mut audio = samples.to_vec();
        audio.resize(CHUNK_SAMPLES, 0.0); // pad/trim to 30 s
        // center pad reflect n_fft/2
        let pad = N_FFT / 2;
        let n = audio.len();
        let mut padded = Vec::with_capacity(n + 2 * pad);
        padded.extend(audio[1..=pad].iter().rev());
        padded.extend_from_slice(&audio);
        padded.extend(audio[n - pad - 1..n - 1].iter().rev());

        use rayon::prelude::*;
        let bins = N_FFT / 2 + 1;
        // windowed frames [N_FRAMES, N_FFT] in f64 (the same per-term product the old loop formed)
        let mut wframes = vec![0f64; N_FRAMES * N_FFT];
        wframes
            .par_chunks_mut(N_FFT)
            .enumerate()
            .for_each(|(t, row)| {
                let frame = &padded[t * HOP..][..N_FFT];
                for j in 0..N_FFT {
                    row[j] = frame[j] as f64 * self.window[j] as f64;
                }
            });
        // power spectrum [bins, N_FRAMES] via the DFT tables — parallel over bins (contiguous rows).
        // Same f64 products in the same j-order as before → bit-identical to the transcendental loop.
        let mut power = vec![0f32; bins * N_FRAMES];
        power
            .par_chunks_mut(N_FRAMES)
            .enumerate()
            .for_each(|(b, prow)| {
                let cb = &self.dft_cos[b * N_FFT..][..N_FFT];
                let sb = &self.dft_sin[b * N_FFT..][..N_FFT];
                for t in 0..N_FRAMES {
                    let wf = &wframes[t * N_FFT..][..N_FFT];
                    let mut re = 0f64;
                    let mut im = 0f64;
                    for j in 0..N_FFT {
                        re += wf[j] * cb[j];
                        im -= wf[j] * sb[j];
                    }
                    prow[t] = (re * re + im * im) as f32;
                }
            });
        // mel = filters · power → [n_mels, frames], then log10, floor, normalize
        let mut mel = vec![0f32; self.n_mels * N_FRAMES];
        mel.par_chunks_mut(N_FRAMES)
            .enumerate()
            .for_each(|(m, mrow)| {
                let fb = &self.mel_filters[m * bins..][..bins];
                for t in 0..N_FRAMES {
                    let mut acc = 0f32;
                    for b in 0..bins {
                        acc += fb[b] * power[b * N_FRAMES + t];
                    }
                    mrow[t] = acc.max(1e-10).log10();
                }
            });
        let gmax = mel.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
        for v in mel.iter_mut() {
            *v = ((v.max(gmax - 8.0)) + 4.0) / 4.0;
        }
        mel
    }

    /// The conv stem + positional embedding: log-mel `[n_mels, 3000]` → `[1500, d]` (the input to
    /// the transformer blocks). Shared by the CPU forward and the GPU path (the conv stem stays CPU).
    pub(crate) fn stem(&self, mel: &[f32]) -> (Vec<f32>, usize) {
        let mut x = vec![0f32; N_FRAMES * self.n_mels];
        for m in 0..self.n_mels {
            for t in 0..N_FRAMES {
                x[t * self.n_mels + m] = mel[m * N_FRAMES + t];
            }
        }
        let (mut h, t1) = self.conv1.forward(&x, N_FRAMES);
        for v in h.iter_mut() {
            *v = gelu(*v);
        }
        let (mut h, t2) = self.conv2.forward(&h, t1);
        for v in h.iter_mut() {
            *v = gelu(*v);
        }
        for (i, v) in h.iter_mut().enumerate() {
            *v += self.pos[i];
        }
        (h, t2)
    }

    /// Encoder forward: log-mel `[n_mels, 3000]` → `[1500, d]`.
    pub fn forward(&self, mel: &[f32]) -> Vec<f32> {
        let (mut h, t2) = self.stem(mel);
        for b in &self.blocks {
            b.forward(&mut h, t2);
        }
        match &self.ln_post {
            Some(ln) => ln.forward(&h),
            None => h,
        }
    }

    /// Full: samples → encoder features `[1500, d]`.
    pub fn encode(&self, samples: &[f32]) -> Vec<f32> {
        self.forward(&self.logmel(samples))
    }
}

// ---------------------------------------------------------------------------------------------
// Decoder — causal self-attn + cross-attn + autoregressive greedy generation
// ---------------------------------------------------------------------------------------------

pub(crate) struct DecBlock {
    pub(crate) norm_self: LayerNorm,
    pub(crate) self_attn: Attention,
    pub(crate) norm_cross: LayerNorm,
    pub(crate) cross_attn: Attention,
    pub(crate) norm_ff: LayerNorm,
    pub(crate) fc1: Linear,
    pub(crate) fc2: Linear,
}

/// Growing per-block self-attention KV cache for autoregressive decode. `k[b]`/`v[b]` are `[pos, d]`
/// row-major, appended one batch at a time (the prompt prefill, then one token per step).
struct DecCache {
    k: Vec<Vec<f32>>,
    v: Vec<Vec<f32>>,
}

impl DecCache {
    fn new(layers: usize) -> Self {
        Self {
            k: vec![Vec::new(); layers],
            v: vec![Vec::new(); layers],
        }
    }
}

/// Whisper (encoder + decoder + greedy generation). The tokenizer (ids → text) is separate.
pub struct Whisper {
    pub encoder: WhisperEncoder,
    pub(crate) embed_tokens: Vec<f32>, // [vocab, d]
    pub(crate) embed_pos: Vec<f32>,    // [max_target, d]
    pub(crate) blocks: Vec<DecBlock>,
    pub(crate) ln_post: LayerNorm,
    pub(crate) lm_head: Linear, // tied: w = embed_tokens, [vocab, d]
    pub(crate) d: usize,
    pub(crate) vocab: usize,
    /// The decode prompt template (SOT, language, task, notimestamps). The language slot is filled by
    /// auto-detection when `lang_tokens` is non-empty; otherwise `prompt` is used verbatim.
    pub prompt: Vec<usize>,
    /// Candidate `<|lang|>` token ids for auto-detection (empty ⇒ detection off, use `prompt` as-is).
    pub(crate) lang_tokens: Vec<usize>,
    pub(crate) transcribe_token: usize,
    pub(crate) notimestamps_token: usize,
    pub(crate) suppress: Vec<usize>,
    pub(crate) begin_suppress: Vec<usize>,
    pub(crate) eos: usize,
    pub(crate) max_target: usize,
}

impl Whisper {
    pub fn load(dir: &Path) -> Result<Self> {
        let v: serde_json::Value =
            serde_json::from_slice(&std::fs::read(dir.join("config.json"))?).context("config")?;
        let gcfg = std::fs::read(dir.join("generation_config.json")).ok();
        let st = LazySt::open(dir)?;
        Self::from_st(&v, gcfg.as_deref(), st)
    }

    /// [`Self::load`] depuis des octets — le chemin navigateur/wasm (OPFS/fetch).
    /// `generation_config` est le contenu de `generation_config.json` s'il existe.
    pub fn from_bytes(
        config: &[u8],
        generation_config: Option<&[u8]>,
        model: Vec<u8>,
    ) -> Result<Self> {
        let v: serde_json::Value = serde_json::from_slice(config).context("config")?;
        Self::from_st(&v, generation_config, LazySt::from_bytes(vec![model])?)
    }

    fn from_st(v: &serde_json::Value, gcfg_bytes: Option<&[u8]>, st: LazySt) -> Result<Self> {
        let encoder = WhisperEncoder::from_st(v, &st, "model.encoder", None)?;
        let g = |k: &str| v.get(k).and_then(|x| x.as_u64()).unwrap_or(0) as usize;
        let list = |k: &str| -> Vec<usize> {
            v.get(k)
                .and_then(|x| x.as_array())
                .map(|a| {
                    a.iter()
                        .filter_map(|x| x.as_u64().map(|u| u as usize))
                        .collect()
                })
                .unwrap_or_default()
        };
        let d = g("d_model");
        let heads = g("decoder_attention_heads");
        let ffn = g("decoder_ffn_dim");
        let vocab = g("vocab_size");
        let blocks = (0..g("decoder_layers"))
            .map(|i| {
                let p = format!("model.decoder.layers.{i}");
                Ok(DecBlock {
                    norm_self: LayerNorm::load(&st, &format!("{p}.self_attn_layer_norm"))?,
                    self_attn: Attention::load(&st, &format!("{p}.self_attn"), d, heads)?,
                    norm_cross: LayerNorm::load(&st, &format!("{p}.encoder_attn_layer_norm"))?,
                    cross_attn: Attention::load(&st, &format!("{p}.encoder_attn"), d, heads)?,
                    norm_ff: LayerNorm::load(&st, &format!("{p}.final_layer_norm"))?,
                    fc1: Linear::load(&st, &format!("{p}.fc1"), ffn, d, true)?,
                    fc2: Linear::load(&st, &format!("{p}.fc2"), d, ffn, true)?,
                })
            })
            .collect::<Result<Vec<_>>>()?;
        let embed_tokens = st.tensor_f32("model.decoder.embed_tokens.weight")?; // [vocab, d]
        // The default en/transcribe/notimestamps prompt; overridable via config `prompt`.
        // `generation_config.json` — HF ships it with every Whisper checkpoint and it carries the
        // special-token ids. Reading it is what makes a STOCK HF directory work: the hardcoded
        // fallback below is the v2 layout, and large-v3 added a language (`<|yue|>`), which shifts
        // every task token by one. Guessing there would silently prompt the decoder with a
        // LANGUAGE token where `<|transcribe|>` belongs.
        let gcfg: serde_json::Value = gcfg_bytes
            .and_then(|b| serde_json::from_slice(b).ok())
            .unwrap_or(serde_json::Value::Null);
        let gid = |k: &str| gcfg.get(k).and_then(|x| x.as_u64()).map(|v| v as usize);
        let gmap = |k: &str| -> Vec<usize> {
            gcfg.get(k)
                .and_then(|m| m.as_object())
                .map(|m| {
                    m.values()
                        .filter_map(|v| v.as_u64())
                        .map(|v| v as usize)
                        .collect()
                })
                .unwrap_or_default()
        };
        let prompt = {
            let p = list("prompt");
            if !p.is_empty() {
                p
            } else if let (Some(sot), Some(nots)) =
                (gid("decoder_start_token_id"), gid("no_timestamps_token_id"))
            {
                let en = gcfg
                    .get("lang_to_id")
                    .and_then(|m| m.get("<|en|>"))
                    .and_then(|v| v.as_u64())
                    .map(|v| v as usize)
                    .unwrap_or(sot + 1);
                let task = gcfg
                    .get("task_to_id")
                    .and_then(|m| m.get("transcribe"))
                    .and_then(|v| v.as_u64())
                    .map(|v| v as usize)
                    .unwrap_or(nots - 4);
                vec![sot, en, task, nots]
            } else {
                vec![50258, 50259, 50359, 50363]
            }
        };
        Ok(Self {
            lm_head: Linear {
                w: embed_tokens.clone(),
                b: vec![0.0; vocab],
                n: vocab,
                k: d,
                packed: std::sync::OnceLock::new(),
            },
            embed_pos: st.tensor_f32("model.decoder.embed_positions.weight")?,
            ln_post: LayerNorm::load(&st, "model.decoder.layer_norm")?,
            embed_tokens,
            encoder,
            blocks,
            d,
            vocab,
            transcribe_token: {
                let t = g("transcribe_token");
                if t == 0 {
                    *prompt.get(2).unwrap_or(&50359)
                } else {
                    t
                }
            },
            notimestamps_token: {
                let t = g("notimestamps_token");
                if t == 0 {
                    *prompt.get(3).unwrap_or(&50363)
                } else {
                    t
                }
            },
            prompt,
            lang_tokens: {
                let l = list("lang_tokens");
                if l.is_empty() { gmap("lang_to_id") } else { l }
            },
            suppress: list("suppress_tokens"),
            begin_suppress: {
                let b = list("begin_suppress_tokens");
                if b.is_empty() { vec![220, 50257] } else { b }
            },
            eos: {
                let e = g("eos_token_id");
                if e == 0 { 50257 } else { e }
            },
            max_target: g("max_target_positions").max(448),
        })
    }

    /// Logits `[vocab]` for the last of `new_ids`, extending the per-block self-attention KV cache.
    /// `start` is the absolute position of `new_ids[0]` (0 for the prompt prefill, then the running
    /// length). Self-attn K/V for earlier tokens live in `cache` and are never recomputed; cross-attn
    /// K/V (`cross_kv`) are the step-invariant encoder projections. Identical outputs to a full
    /// recompute — the cache only skips work — so greedy tokens are unchanged. O(n²)→O(n) decode.
    fn last_logits_cached(
        &self,
        new_ids: &[usize],
        start: usize,
        cache: &mut DecCache,
        cross_kv: &[(Vec<f32>, Vec<f32>)],
    ) -> Vec<f32> {
        let (d, n) = (self.d, new_ids.len());
        let mut x = vec![0f32; n * d];
        for (i, &id) in new_ids.iter().enumerate() {
            let pos = start + i;
            for c in 0..d {
                x[i * d + c] = self.embed_tokens[id * d + c] + self.embed_pos[pos * d + c];
            }
        }
        let tk = start + n; // total cached keys after appending this batch
        for (b, blk) in self.blocks.iter().enumerate() {
            let normed = blk.norm_self.forward(&x);
            cache.k[b].extend_from_slice(&blk.self_attn.k.forward(&normed));
            cache.v[b].extend_from_slice(&blk.self_attn.v.forward(&normed));
            let a =
                blk.self_attn
                    .forward_qkv_cached(&normed, &cache.k[b], &cache.v[b], n, tk, start);
            for (xi, ai) in x.iter_mut().zip(a) {
                *xi += ai;
            }
            let normed = blk.norm_cross.forward(&x);
            let (ck, cv) = &cross_kv[b];
            let a = blk
                .cross_attn
                .forward_qkv(&normed, ck, cv, n, N_FRAMES / 2, false);
            for (xi, ai) in x.iter_mut().zip(a) {
                *xi += ai;
            }
            let normed = blk.norm_ff.forward(&x);
            let mut h = blk.fc1.forward(&normed);
            for v in h.iter_mut() {
                *v = gelu(*v);
            }
            let f = blk.fc2.forward(&h);
            for (xi, fi) in x.iter_mut().zip(f) {
                *xi += fi;
            }
        }
        let x = self.ln_post.forward(&x);
        self.lm_head.forward(&x[(n - 1) * d..][..d]) // [vocab]
    }

    /// Greedy transcribe: samples → token ids (INCLUDING the forced prompt and the final EOS),
    /// matching HF's `generate(num_beams=1)` with the suppress-token processors. KV-cached: the
    /// prompt is prefilled in one pass, then one token per step.
    pub fn generate(&self, samples: &[f32], max_new: usize) -> Vec<usize> {
        self.generate_from_enc(&self.encoder.encode(samples), max_new)
    }

    /// Auto-detect the spoken language: one decoder step from SOT, argmax over the `<|lang|>` token
    /// set (exactly HF's `detect_language`). Returns the language token id; when detection is disabled
    /// (no `lang_tokens` in config) returns the template's language slot.
    pub fn detect_language(&self, enc: &[f32]) -> usize {
        let cross_kv: Vec<(Vec<f32>, Vec<f32>)> = self
            .blocks
            .iter()
            .map(|b| b.cross_attn.project_kv(enc))
            .collect();
        self.detect_language_kv(&cross_kv)
    }

    fn detect_language_kv(&self, cross_kv: &[(Vec<f32>, Vec<f32>)]) -> usize {
        if self.lang_tokens.is_empty() {
            return *self.prompt.get(1).unwrap_or(&50259);
        }
        let mut cache = DecCache::new(self.blocks.len());
        let logits = self.last_logits_cached(&[self.prompt[0]], 0, &mut cache, cross_kv);
        *self
            .lang_tokens
            .iter()
            .max_by(|&&a, &&b| {
                logits[a]
                    .partial_cmp(&logits[b])
                    .unwrap_or(std::cmp::Ordering::Equal)
            })
            .expect("lang_tokens non-empty")
    }

    /// The decode prompt with the language slot resolved by detection (or the template verbatim).
    fn build_prompt(&self, cross_kv: &[(Vec<f32>, Vec<f32>)]) -> Vec<usize> {
        if self.lang_tokens.is_empty() {
            return self.prompt.clone();
        }
        vec![
            self.prompt[0],
            self.detect_language_kv(cross_kv),
            self.transcribe_token,
            self.notimestamps_token,
        ]
    }

    /// Greedy decode from PRE-COMPUTED encoder features `[1500, d]` — lets the caller run the encoder
    /// on the GPU ([`crate::whisper_gpu::WhisperEncoderGpu`]) while the (small, autoregressive) decode
    /// stays on the CPU. Identical to [`Self::generate`] given the same encoder output.
    pub fn generate_from_enc(&self, enc: &[f32], max_new: usize) -> Vec<usize> {
        let cross_kv: Vec<(Vec<f32>, Vec<f32>)> = self
            .blocks
            .iter()
            .map(|b| b.cross_attn.project_kv(enc))
            .collect();
        let mut ids = self.build_prompt(&cross_kv);
        let mut cache = DecCache::new(self.blocks.len());
        let cap = max_new.min(self.max_target.saturating_sub(ids.len()));
        let mut pending = ids.clone(); // step 0 prefills the whole prompt; later steps feed one token
        let mut pos = 0usize;
        for step in 0..cap {
            let mut logits = self.last_logits_cached(&pending, pos, &mut cache, &cross_kv);
            pos += pending.len();
            for &t in &self.suppress {
                if t < self.vocab {
                    logits[t] = f32::NEG_INFINITY;
                }
            }
            if step == 0 {
                for &t in &self.begin_suppress {
                    if t < self.vocab {
                        logits[t] = f32::NEG_INFINITY;
                    }
                }
            }
            let mut nxt = 0;
            let mut best = f32::NEG_INFINITY;
            for (i, &l) in logits.iter().enumerate() {
                if l > best {
                    best = l;
                    nxt = i;
                }
            }
            ids.push(nxt);
            if nxt == self.eos {
                break;
            }
            pending = vec![nxt]; // the next step feeds only the token just produced
        }
        ids
    }

    /// Greedy transcribe to TEXT, using the checkpoint's BPE tokenizer (`tokenizer.json` in the
    /// model dir). Special tokens (prompt/EOS, all ≥ 50257) are dropped before decoding.
    ///
    /// Audio longer than the model's 30 s window is transcribed in overlapping windows and
    /// stitched — see [`Self::transcribe_windowed`]. It used to be silently TRUNCATED to the first
    /// 30 s, which is the worst possible failure for a transcriber: a confident, complete-looking
    /// transcript of the first half minute of a meeting.
    #[cfg(feature = "cli")]
    pub fn transcribe(&self, dir: &Path, samples: &[f32]) -> Result<String> {
        if samples.len() <= CHUNK_SAMPLES {
            return self.decode_text(dir, &self.generate(samples, self.max_target));
        }
        self.transcribe_windowed(dir, samples)
    }

    /// Long-form transcription: 30 s windows with [`WINDOW_OVERLAP_SAMPLES`] of overlap, each
    /// decoded independently, then stitched by deleting the duplicated words in the seam.
    ///
    /// Why overlap-and-stitch rather than the reference's sequential timestamp walk: this decoder
    /// is prompted with `<|notimestamps|>`, so there are no segment times to advance by. Fixed
    /// windows need no timestamps, and the overlap is what makes the seam recoverable — a word
    /// straddling a boundary is cut in one window but whole in its neighbour.
    ///
    /// Windows are INDEPENDENT, so they run in parallel; long audio therefore costs about
    /// `ceil(len / stride)` encoder passes spread across cores rather than serialized.
    #[cfg(feature = "cli")]
    pub fn transcribe_windowed(&self, dir: &Path, samples: &[f32]) -> Result<String> {
        use rayon::prelude::*;
        let stride = CHUNK_SAMPLES - WINDOW_OVERLAP_SAMPLES;
        let starts: Vec<usize> = (0..samples.len()).step_by(stride).collect();
        let texts = starts
            .par_iter()
            .map(|&start| {
                let end = (start + CHUNK_SAMPLES).min(samples.len());
                let ids = self.generate(&samples[start..end], self.max_target);
                self.decode_text(dir, &ids)
            })
            .collect::<Result<Vec<String>>>()?;

        let mut out: Vec<String> = Vec::new();
        for t in texts {
            let next: Vec<String> = t.split_whitespace().map(str::to_string).collect();
            stitch(&mut out, &next);
        }
        Ok(out.join(" "))
    }

    /// Detokenize greedy ids → text (drops special tokens ≥ eos). Split out so a caller that ran the
    /// encoder on the GPU ([`Self::generate_from_enc`]) can finish without recomputing the encoder.
    #[cfg(feature = "cli")]
    pub fn decode_text(&self, dir: &Path, ids: &[usize]) -> Result<String> {
        let tok = tokenizers::Tokenizer::from_file(dir.join("tokenizer.json"))
            .map_err(|e| anyhow::anyhow!("load tokenizer.json: {e}"))?;
        let content: Vec<u32> = ids
            .iter()
            .filter(|&&t| t < self.eos)
            .map(|&t| t as u32)
            .collect();
        let text = tok
            .decode(&content, true)
            .map_err(|e| anyhow::anyhow!("decode: {e}"))?;
        Ok(text.trim().to_string())
    }

    /// The decoder's max target length (the greedy cap), for callers driving [`Self::generate_from_enc`].
    pub fn max_target(&self) -> usize {
        self.max_target
    }
}

/// Append `next` to `acc`, deleting the words the two share at the seam.
///
/// Adjacent windows overlap by [`WINDOW_OVERLAP_SAMPLES`], so the same speech is transcribed
/// twice and the duplicate has to go. Longest match wins: find the largest `k` where the last `k`
/// words of `acc` equal the first `k` of `next`, and drop those `k` from `next`. Comparison is on
/// a normalized form (lowercase, no punctuation) because the two windows often differ only in
/// capitalization or a comma — `"today. Tell"` vs `"today, tell"`.
///
/// With no match the whole of `next` is appended: a false seam repeats a few words, while an
/// over-eager match would DELETE speech, and losing words is the worse failure.
#[cfg(feature = "cli")]
fn stitch(acc: &mut Vec<String>, next: &[String]) {
    fn norm(w: &str) -> String {
        w.chars()
            .filter(|c| c.is_alphanumeric())
            .flat_map(char::to_lowercase)
            .collect()
    }
    if acc.is_empty() {
        acc.extend_from_slice(next);
        return;
    }
    // Bounded by the overlap's worth of words (~5 s of speech is well under 40).
    let cap = acc.len().min(next.len()).min(40);
    for k in (1..=cap).rev() {
        let tail = &acc[acc.len() - k..];
        if tail
            .iter()
            .zip(&next[..k])
            .all(|(a, b)| !norm(a).is_empty() && norm(a) == norm(b))
        {
            acc.extend_from_slice(&next[k..]);
            return;
        }
    }
    acc.extend_from_slice(next);
}

/// The Slaney mel filterbank HF's `WhisperFeatureExtractor` uses, built in Rust.
///
/// Whisper checkpoints do not ship this table — the Python export baked it into the safetensors.
/// Building it here removes that step, so a stock `openai/whisper-*` or `distil-whisper/*`
/// directory (config.json + model.safetensors + tokenizer.json) loads directly, including the
/// 128-bin large-v3 family. `tests/whisper_parity.rs` pins it against the exported table.
///
/// Slaney scale (not HTK): linear below 1 kHz at 200/3 Hz per mel, logarithmic above, and each
/// triangle normalized by `2 / (hz[i+2] - hz[i])` so filters are equal-AREA rather than
/// equal-height. Returns `[n_mels, n_fft/2 + 1]` row-major.
pub fn mel_filter_bank(n_mels: usize, n_fft: usize, sample_rate: usize) -> Vec<f32> {
    const F_SP: f64 = 200.0 / 3.0;
    const MIN_LOG_HZ: f64 = 1000.0;
    let min_log_mel = MIN_LOG_HZ / F_SP; // 15.0
    let logstep = (6.4f64).ln() / 27.0;
    let hz_to_mel = |hz: f64| -> f64 {
        if hz < MIN_LOG_HZ {
            hz / F_SP
        } else {
            min_log_mel + (hz / MIN_LOG_HZ).ln() / logstep
        }
    };
    let mel_to_hz = |mel: f64| -> f64 {
        if mel < min_log_mel {
            mel * F_SP
        } else {
            MIN_LOG_HZ * ((mel - min_log_mel) * logstep).exp()
        }
    };

    let bins = n_fft / 2 + 1;
    // Whisper's front end covers 0 .. 8 kHz regardless of the mel count.
    let (mel_min, mel_max) = (hz_to_mel(0.0), hz_to_mel(8000.0));
    let points: Vec<f64> = (0..n_mels + 2)
        .map(|i| mel_to_hz(mel_min + (mel_max - mel_min) * i as f64 / (n_mels + 1) as f64))
        .collect();
    let fft_hz: Vec<f64> = (0..bins)
        .map(|i| i as f64 * sample_rate as f64 / n_fft as f64)
        .collect();

    let mut fb = vec![0f32; n_mels * bins];
    for m in 0..n_mels {
        let (lo, ctr, hi) = (points[m], points[m + 1], points[m + 2]);
        let enorm = 2.0 / (hi - lo); // Slaney area normalization
        for (b, &f) in fft_hz.iter().enumerate() {
            let ramp = if f <= ctr {
                (f - lo) / (ctr - lo)
            } else {
                (hi - f) / (hi - ctr)
            };
            if ramp > 0.0 {
                fb[m * bins + b] = (ramp * enorm) as f32;
            }
        }
    }
    fb
}