car-inference 0.47.0

//! Native Parakeet-TDT speech-to-text backend for Apple Silicon via mlx-rs.
//!
//! Implements the full Parakeet-TDT-0.6B-v3 architecture:
//! - Mel spectrogram preprocessor (128 features, 16kHz)
//! - Conformer encoder (24 layers, d_model=1024)
//! - RNN-T decoder with TDT (Token-and-Duration Transducer)
//!
//! Model weights: `mlx-community/parakeet-tdt-0.6b-v3` (model.safetensors, unquantized).
//! This eliminates the last Python dependency for STT on Apple Silicon.

use std::collections::HashMap;
use std::path::Path;

use mlx_rs::nn;
use mlx_rs::ops;
use mlx_rs::ops::indexing::IndexOp;
use mlx_rs::Array;
use tracing::{info, warn};

use super::mlx::load_all_tensors;
use crate::InferenceError;

// ─── Constants ─────────────────────────────────────────────────────────────

const SAMPLE_RATE: usize = 16000;
const N_MEL: usize = 128;
const N_FFT: usize = 512;
const WIN_LEN_SAMPLES: usize = 400; // 25ms at 16kHz
const HOP_LEN_SAMPLES: usize = 160; // 10ms at 16kHz
/// NeMo's `preemph`: `y[n] = x[n] - 0.97*x[n-1]`, applied to the raw waveform
/// before framing. Training used it, so inference must too.
const PREEMPH: f32 = 0.97;
/// NeMo's `log_zero_guard_value`, added *inside* the log (`log(x + 2^-24)`),
/// not a floor under it (`log(max(x, eps))`). The two differ everywhere the
/// mel energy is small, which is most of a spectrogram.
const LOG_ZERO_GUARD: f32 = 5.960_464_5e-8; // 2^-24

const D_MODEL: usize = 1024;
const NUM_HEADS: usize = 8;
const HEAD_DIM: usize = D_MODEL / NUM_HEADS; // 128
const NUM_ENCODER_LAYERS: usize = 24;
/// Channels inside the subsampling stack, and the mel bins left after its three
/// stride-2 convolutions (128 → 64 → 32 → 16). Their product is the width of
/// the linear that projects into `d_model`.
const SUBSAMPLE_CHANNELS: usize = 256;
const SUBSAMPLE_FREQ_OUT: usize = 16;

const PRED_HIDDEN: usize = 640;
const PRED_LAYERS: usize = 2;
/// Real vocabulary size. **The blank is `VOCAB_SIZE`, not 0** — NeMo's RNN-T
/// appends the blank after the vocabulary, and the checkpoint's fused joint
/// head is `[8198, 640]` = 8192 vocab + 1 blank + 5 duration bins. Decoding
/// with blank=0 emits `<unk>` forever.
const VOCAB_SIZE: u32 = 8192;
const BLANK_ID: u32 = VOCAB_SIZE;
/// The TDT duration bins (`decoding.durations` in the config). The joint's
/// duration head picks an index into this table; the decoder advances the
/// encoder frame cursor by that many frames — including **zero**, which is what
/// lets TDT emit several tokens on one frame.
const TDT_DURATIONS: [usize; 5] = [0, 1, 2, 3, 4];
/// `decoding.greedy.max_symbols`. Without it a run of zero-duration emissions
/// never advances `time` and the decode does not terminate.
const MAX_SYMBOLS_PER_FRAME: usize = 10;

// ─── Mel Spectrogram ───────────────────────────────────────────────────────

/// Compute the mel filterbank matrix of shape `(n_mel, n_fft/2+1)`.
///
/// **Slaney mel scale with Slaney area normalization**, matching the
/// `mel_filters(..., norm="slaney", mel_scale="slaney")` the reference calls.
/// This is not the HTK formula: below 1 kHz Slaney is linear in frequency and
/// only becomes logarithmic above it, so the two disagree across the whole
/// filterbank, and the area normalization rescales every filter on top of that.
/// A mel front end that differs from the one the model was trained with does
/// not fail loudly — it produces a plausible spectrogram and a wrong transcript.
fn mel_filterbank() -> Result<Array, mlx_rs::error::Exception> {
    const MIN_LOG_HZ: f64 = 1000.0;
    const F_SP: f64 = 200.0 / 3.0;

    let n_freqs = N_FFT / 2 + 1; // 257
    let f_max = (SAMPLE_RATE / 2) as f64;
    let min_log_mel = MIN_LOG_HZ / F_SP;
    let logstep = (6.4f64).ln() / 27.0;

    let hz_to_mel = |freq: f64| -> f64 {
        if freq >= MIN_LOG_HZ {
            min_log_mel + (freq / MIN_LOG_HZ).ln() / logstep
        } else {
            freq / F_SP
        }
    };
    let mel_to_hz = |mel: f64| -> f64 {
        if mel >= min_log_mel {
            MIN_LOG_HZ * (logstep * (mel - min_log_mel)).exp()
        } else {
            F_SP * mel
        }
    };

    // `linspace(0, sample_rate // 2, n_freqs)` — endpoint included, so the step
    // is 8000/256 = 31.25 Hz, i.e. exactly the rfft bin spacing.
    let all_freqs: Vec<f64> = (0..n_freqs)
        .map(|i| f_max * i as f64 / (n_freqs - 1) as f64)
        .collect();

    let m_min = hz_to_mel(0.0);
    let m_max = hz_to_mel(f_max);
    let n_pts = N_MEL + 2;
    let f_pts: Vec<f64> = (0..n_pts)
        .map(|i| mel_to_hz(m_min + (m_max - m_min) * i as f64 / (n_pts - 1) as f64))
        .collect();
    let f_diff: Vec<f64> = f_pts.windows(2).map(|w| w[1] - w[0]).collect();

    let mut fb = vec![0.0f32; N_MEL * n_freqs];
    for m in 0..N_MEL {
        // Slaney area normalization: each filter integrates to a constant, so
        // wide high-frequency filters do not dominate narrow low ones.
        let enorm = 2.0 / (f_pts[m + 2] - f_pts[m]);
        for (k, &freq) in all_freqs.iter().enumerate() {
            let down = -(f_pts[m] - freq) / f_diff[m];
            let up = (f_pts[m + 2] - freq) / f_diff[m + 1];
            fb[m * n_freqs + k] = (down.min(up).max(0.0) * enorm) as f32;
        }
    }

    Ok(Array::from_slice(&fb, &[N_MEL as i32, n_freqs as i32]))
}

/// Compute the log-mel spectrogram from raw 16 kHz PCM samples, reproducing
/// NeMo's `AudioToMelSpectrogramPreprocessor` as the reference implements it.
/// Input: f32 samples. Output: shape `(1, n_frames, N_MEL)`.
///
/// Four details here are load-bearing and were all absent or wrong before, each
/// of which silently shifts the feature distribution away from the one the
/// encoder was trained on:
///
/// 1. **Pre-emphasis** on the raw waveform.
/// 2. **A 400-sample Hann window centred inside the 512-point FFT frame**, not
///    left-aligned with zeros trailing it. Off-centring the window rotates the
///    phase of every bin.
/// 3. **Centred framing** — the signal is padded by `n_fft/2` on both sides, so
///    frame *i* is centred on sample `i * hop`. This is what makes the frame
///    count `1 + len/hop` rather than `1 + (len - win)/hop`.
/// 4. **Per-feature normalization** — each mel bin is standardized over time
///    using the *sample* variance (`n-1`). Skipping it feeds the encoder
///    unnormalized log-energies, which is the single largest discrepancy.
///
/// `dither` is deliberately not applied: it adds training-time noise and the
/// reference's inference path does not use it either.
fn compute_log_mel(samples: &[f32]) -> Result<Array, InferenceError> {
    let map_err = |e: mlx_rs::error::Exception| InferenceError::InferenceFailed(e.to_string());

    if samples.is_empty() {
        return Err(InferenceError::InferenceFailed(
            "empty audio: nothing to transcribe".into(),
        ));
    }

    // (1) Pre-emphasis: y[0] = x[0], y[n] = x[n] - 0.97*x[n-1].
    let mut emphasized = Vec::with_capacity(samples.len());
    emphasized.push(samples[0]);
    for i in 1..samples.len() {
        emphasized.push(samples[i] - PREEMPH * samples[i - 1]);
    }

    // (3) Centre-pad by n_fft/2 with zeros (`pad_mode="constant"`).
    let half = N_FFT / 2;
    let mut padded = vec![0.0f32; half];
    padded.extend_from_slice(&emphasized);
    padded.extend(std::iter::repeat_n(0.0f32, half));

    let num_frames = 1 + (padded.len() - N_FFT) / HOP_LEN_SAMPLES;

    // (2) Symmetric (non-periodic) Hann of win_length, centred in the FFT frame.
    let hann: Vec<f32> = (0..WIN_LEN_SAMPLES)
        .map(|n| {
            0.5 * (1.0
                - (2.0 * std::f32::consts::PI * n as f32 / (WIN_LEN_SAMPLES - 1) as f32).cos())
        })
        .collect();
    let win_offset = (N_FFT - WIN_LEN_SAMPLES) / 2;
    let mut window = vec![0.0f32; N_FFT];
    window[win_offset..win_offset + WIN_LEN_SAMPLES].copy_from_slice(&hann);

    let mut framed = vec![0.0f32; num_frames * N_FFT];
    for f in 0..num_frames {
        let start = f * HOP_LEN_SAMPLES;
        for s in 0..N_FFT {
            framed[f * N_FFT + s] = padded[start + s] * window[s];
        }
    }

    let framed_arr = Array::from_slice(&framed, &[num_frames as i32, N_FFT as i32]);
    let spectrum = mlx_rs::fft::rfft(&framed_arr, N_FFT as i32, -1).map_err(map_err)?;

    // Power spectrum |X|^2, as (num_frames, 257) real.
    let mag = ops::abs(&spectrum).map_err(map_err)?;
    let power = ops::square(&mag).map_err(map_err)?;
    let power = power.as_dtype(mlx_rs::Dtype::Float32).map_err(map_err)?;

    // power (T, 257) @ filters^T (257, 128) => (T, 128).
    //
    // The reference computes `filters @ power.T`, i.e. mel-major `(128, T)`,
    // and transposes back at the very end. Doing it time-major throughout is
    // the same arithmetic and keeps the result **contiguous in the layout the
    // encoder wants**, so no transpose sits between here and the caller. That
    // matters: a transpose immediately before the final `reshape` produced an
    // array whose logical shape was `(1, T, 128)` while its buffer was still in
    // `(128, T)` order, which reads as a plausible spectrogram with time and
    // frequency swapped — values individually correct, every one in the wrong
    // place.
    let fb = mel_filterbank().map_err(map_err)?;
    let fb_t = ops::transpose_axes(&fb, &[1, 0]).map_err(map_err)?;
    let mel = ops::matmul(&power, &fb_t).map_err(map_err)?;

    // log(x + 2^-24) — the guard is inside the log, not a floor.
    let guard = Array::from_f32(LOG_ZERO_GUARD);
    let log_mel = ops::log(&ops::add(&mel, &guard).map_err(map_err)?).map_err(map_err)?;

    // (4) Per-feature standardization: reduce over **time** (axis 0 here), with
    // the n-1 (sample) denominator the reference uses.
    let mean = log_mel.mean_axes(&[0], true).map_err(map_err)?;
    let centered = ops::subtract(&log_mel, &mean).map_err(map_err)?;
    let denom = Array::from_f32((num_frames.saturating_sub(1)).max(1) as f32);
    let variance = ops::divide(
        &centered
            .square()
            .map_err(map_err)?
            .sum_axes(&[0], true)
            .map_err(map_err)?,
        &denom,
    )
    .map_err(map_err)?;
    let std = ops::sqrt(&variance).map_err(map_err)?;
    let std = ops::add(&std, Array::from_f32(1e-5)).map_err(map_err)?;
    let normalized = ops::divide(&centered, &std).map_err(map_err)?;

    // (T, 128) -> (1, T, 128): a pure leading-axis insert, no data movement.
    ops::reshape(&normalized, &[1, num_frames as i32, N_MEL as i32]).map_err(map_err)
}

/// Load 16kHz mono WAV from file path and return f32 samples.
fn load_wav(path: &Path) -> Result<Vec<f32>, InferenceError> {
    let data = std::fs::read(path)
        .map_err(|e| InferenceError::InferenceFailed(format!("read wav: {e}")))?;

    // Minimal WAV parser for PCM format
    if data.len() < 44 {
        return Err(InferenceError::InferenceFailed("WAV file too short".into()));
    }
    if &data[0..4] != b"RIFF" || &data[8..12] != b"WAVE" {
        return Err(InferenceError::InferenceFailed(
            "not a valid WAV file".into(),
        ));
    }

    // Find fmt chunk
    let mut pos = 12;
    let mut sample_rate = 0u32;
    let mut bits_per_sample = 0u16;
    let mut num_channels = 0u16;
    let mut audio_format = 0u16;
    let mut data_start = 0usize;
    let mut data_len = 0usize;

    while pos + 8 <= data.len() {
        let chunk_id = &data[pos..pos + 4];
        let chunk_size =
            u32::from_le_bytes([data[pos + 4], data[pos + 5], data[pos + 6], data[pos + 7]])
                as usize;

        if chunk_id == b"fmt " && chunk_size >= 16 {
            audio_format = u16::from_le_bytes([data[pos + 8], data[pos + 9]]);
            num_channels = u16::from_le_bytes([data[pos + 10], data[pos + 11]]);
            sample_rate = u32::from_le_bytes([
                data[pos + 12],
                data[pos + 13],
                data[pos + 14],
                data[pos + 15],
            ]);
            bits_per_sample = u16::from_le_bytes([data[pos + 22], data[pos + 23]]);
        } else if chunk_id == b"data" {
            data_start = pos + 8;
            data_len = chunk_size;
        }

        pos += 8 + chunk_size;
        // WAV chunks are word-aligned
        if !chunk_size.is_multiple_of(2) {
            pos += 1;
        }
    }

    if audio_format != 1 {
        return Err(InferenceError::InferenceFailed(format!(
            "unsupported WAV format {audio_format}, only PCM (1) supported"
        )));
    }
    if sample_rate == 0 {
        return Err(InferenceError::InferenceFailed(
            "WAV declares a zero sample rate".into(),
        ));
    }
    if data_start == 0 || data_len == 0 {
        return Err(InferenceError::InferenceFailed(
            "no data chunk found in WAV".into(),
        ));
    }

    let end = (data_start + data_len).min(data.len());
    let raw = &data[data_start..end];

    let samples: Vec<f32> = match bits_per_sample {
        16 => raw
            .chunks_exact(2 * num_channels as usize)
            .map(|frame| {
                // Take first channel only (mono mix)
                let s = i16::from_le_bytes([frame[0], frame[1]]);
                s as f32 / 32768.0
            })
            .collect(),
        32 => raw
            .chunks_exact(4 * num_channels as usize)
            .map(|frame| f32::from_le_bytes([frame[0], frame[1], frame[2], frame[3]]))
            .collect(),
        _ => {
            return Err(InferenceError::InferenceFailed(format!(
                "unsupported bits_per_sample: {bits_per_sample}"
            )))
        }
    };

    // Resample rather than refuse. The model is 16 kHz-only, but callers hand
    // this whatever they have — `car speech smoke` feeds it Kokoro's 24 kHz TTS
    // output, and rejecting that made the native backend unreachable on the one
    // command that exercises it end to end. The managed Python fallback never
    // hit this because `mlx_audio`'s loader resamples internally, so the gap
    // only showed once the native path started loading. `car-whisper` is
    // already a direct dependency and owns the same conversion for its own
    // 16 kHz input, so this reuses that resampler instead of adding one.
    if sample_rate == SAMPLE_RATE as u32 {
        return Ok(samples);
    }
    info!(
        from = sample_rate,
        to = SAMPLE_RATE,
        "resampling audio for the Parakeet front end"
    );
    car_whisper::resample_to_16k(&samples, sample_rate)
        .map_err(|e| InferenceError::InferenceFailed(format!("resample to 16kHz: {e}")))
}

// ─── Layer Norm ────────────────────────────────────────────────────────────

struct LayerNorm {
    weight: Array,
    bias: Array,
    eps: f32,
}

impl LayerNorm {
    fn forward(&self, x: &Array) -> Result<Array, mlx_rs::error::Exception> {
        let mean = x.mean_axes(&[-1], true)?;
        let centered = ops::subtract(x, &mean)?;
        let var = centered.square()?.mean_axes(&[-1], true)?;
        let eps_arr = Array::from_f32(self.eps);
        let norm = ops::rsqrt(&ops::add(&var, &eps_arr)?)?;
        let normed = ops::multiply(&centered, &norm)?;
        let scaled = ops::multiply(&normed, &self.weight)?;
        ops::add(&scaled, &self.bias)
    }

    fn all_arrays(&self) -> Vec<&Array> {
        vec![&self.weight, &self.bias]
    }
}

fn load_layer_norm(
    tensors: &HashMap<String, Array>,
    prefix: &str,
    eps: f32,
) -> Result<LayerNorm, InferenceError> {
    let weight = get_tensor(tensors, &format!("{prefix}.weight"))?;
    let bias = get_tensor(tensors, &format!("{prefix}.bias"))?;
    Ok(LayerNorm { weight, bias, eps })
}

// ─── Dense Linear ──────────────────────────────────────────────────────────

struct Linear {
    weight: Array,
    bias: Option<Array>,
}

impl Linear {
    fn forward(&self, x: &Array) -> Result<Array, mlx_rs::error::Exception> {
        let w_t = ops::transpose_axes(&self.weight, &[1, 0])?;
        let out = ops::matmul(x, &w_t)?;
        if let Some(ref b) = self.bias {
            ops::add(&out, b)
        } else {
            Ok(out)
        }
    }

    fn all_arrays(&self) -> Vec<&Array> {
        let mut v = vec![&self.weight];
        if let Some(ref b) = self.bias {
            v.push(b);
        }
        v
    }
}

fn load_linear(tensors: &HashMap<String, Array>, prefix: &str) -> Result<Linear, InferenceError> {
    let weight = get_tensor(tensors, &format!("{prefix}.weight"))?;
    let bias = tensors.get(&format!("{prefix}.bias")).cloned();
    Ok(Linear { weight, bias })
}

fn load_linear_with_bias(
    tensors: &HashMap<String, Array>,
    prefix: &str,
) -> Result<Linear, InferenceError> {
    let weight = get_tensor(tensors, &format!("{prefix}.weight"))?;
    let bias = get_tensor(tensors, &format!("{prefix}.bias"))?;
    Ok(Linear {
        weight,
        bias: Some(bias),
    })
}

fn get_tensor(tensors: &HashMap<String, Array>, key: &str) -> Result<Array, InferenceError> {
    tensors
        .get(key)
        .cloned()
        .ok_or_else(|| InferenceError::InferenceFailed(format!("missing tensor: {key}")))
}

// ─── Depthwise Striding Subsampling ────────────────────────────────────────

/// One convolution in the subsampling stack, with its `groups` (depthwise when
/// `groups == in_channels`, pointwise/dense when 1).
struct SubsampleConv {
    weight: Array,
    bias: Array,
    stride: i32,
    padding: i32,
    groups: i32,
}

impl SubsampleConv {
    fn forward(&self, x: &Array) -> Result<Array, mlx_rs::error::Exception> {
        let h = ops::conv2d(
            x,
            &self.weight,
            (self.stride, self.stride),
            (self.padding, self.padding),
            (1, 1),
            self.groups,
        )?;
        ops::add(&h, &self.bias)
    }
}

/// NeMo's `dw_striding` subsampling: a **2-D** depthwise-separable stack that
/// treats the spectrogram as a single-channel image and strides *both* time and
/// frequency by 2, three times (subsampling_factor 8).
///
/// This is not a stack of three Conv1ds over time — that was the previous
/// shape here, and it is a different operator with different weights. The
/// checkpoint carries five conv tensors, not three: after the first dense
/// `1→256` conv, each further stage is a depthwise `3×3` (groups=256) followed
/// by a pointwise `1×1`. Frequency shrinks 128 → 64 → 32 → 16 alongside time,
/// and the surviving `256 × 16` is what `out` projects into `d_model`. There is
/// no LayerNorm in this block.
///
/// Input `(batch, time, n_mel)` → output `(batch, time/8, d_model)`.
struct DepthwiseSubsampling {
    convs: Vec<(SubsampleConv, bool)>, // (conv, relu_after)
    out: Linear,
}

impl DepthwiseSubsampling {
    fn forward(&self, x: &Array) -> Result<Array, mlx_rs::error::Exception> {
        let shape = x.shape().to_vec();
        let (batch, time) = (shape[0], shape[1]);

        // (B, T, F) -> NHWC with H=time, W=freq, C=1.
        let mut h = ops::reshape(x, &[batch, time, N_MEL as i32, 1])?;
        for (conv, relu_after) in &self.convs {
            h = conv.forward(&h)?;
            if *relu_after {
                h = nn::relu(&h)?;
            }
        }

        // (B, T', F', C) -> (B, T', C, F') -> (B, T', C*F'). Channel-major
        // before frequency, matching the reference's transpose+swapaxes pair.
        let h = ops::transpose_axes(&h, &[0, 1, 3, 2])?;
        let out_time = h.shape()[1];
        let flat = (SUBSAMPLE_CHANNELS * SUBSAMPLE_FREQ_OUT) as i32;
        let h = ops::reshape(&h, &[batch, out_time, flat])?;
        self.out.forward(&h)
    }

    fn all_arrays(&self) -> Vec<&Array> {
        let mut v = Vec::new();
        for (conv, _) in &self.convs {
            v.push(&conv.weight);
            v.push(&conv.bias);
        }
        v.extend(self.out.all_arrays());
        v
    }
}

fn load_subsampling(
    tensors: &HashMap<String, Array>,
    prefix: &str,
) -> Result<DepthwiseSubsampling, InferenceError> {
    // (checkpoint index, stride, padding, groups, relu after). Indices 1, 4 and
    // 7 are the ReLUs and carry no weights, which is why the saved indices skip.
    const STAGES: [(u32, i32, i32, i32, bool); 5] = [
        (0, 2, 1, 1, true),                          // dense 1 -> 256, stride 2
        (2, 2, 1, SUBSAMPLE_CHANNELS as i32, false), // depthwise 3x3, stride 2
        (3, 1, 0, 1, true),                          // pointwise 1x1
        (5, 2, 1, SUBSAMPLE_CHANNELS as i32, false), // depthwise 3x3, stride 2
        (6, 1, 0, 1, true),                          // pointwise 1x1
    ];

    let mut convs = Vec::with_capacity(STAGES.len());
    for (idx, stride, padding, groups, relu_after) in STAGES {
        let conv = SubsampleConv {
            weight: get_tensor(tensors, &format!("{prefix}.conv.{idx}.weight"))?,
            bias: get_tensor(tensors, &format!("{prefix}.conv.{idx}.bias"))?,
            stride,
            padding,
            groups,
        };
        convs.push((conv, relu_after));
    }

    Ok(DepthwiseSubsampling {
        convs,
        out: load_linear_with_bias(tensors, &format!("{prefix}.out"))?,
    })
}

// ─── Conformer Feed-Forward Module ─────────────────────────────────────────

/// Macaron-style feed-forward: LayerNorm -> Linear -> SiLU -> Dropout -> Linear -> Dropout
/// With half-step residual in the conformer block.
struct FeedForwardModule {
    norm: LayerNorm,
    linear1: Linear,
    linear2: Linear,
}

impl FeedForwardModule {
    fn forward(&self, x: &Array) -> Result<Array, mlx_rs::error::Exception> {
        let h = self.norm.forward(x)?;
        let h = self.linear1.forward(&h)?;
        let h = nn::silu(&h)?;
        self.linear2.forward(&h)
    }

    fn all_arrays(&self) -> Vec<&Array> {
        let mut v = self.norm.all_arrays();
        v.extend(self.linear1.all_arrays());
        v.extend(self.linear2.all_arrays());
        v
    }
}

/// The norm is a *sibling* of the feed-forward block in the checkpoint
/// (`norm_feed_forward1` next to `feed_forward1`), not nested inside it, so the
/// two prefixes are passed separately. The linears are weight-only: the config
/// sets `use_bias: false` for the whole conformer, so requiring a `.bias` here
/// is what produced the original hard miss.
fn load_ff_module(
    tensors: &HashMap<String, Array>,
    norm_prefix: &str,
    ff_prefix: &str,
) -> Result<FeedForwardModule, InferenceError> {
    Ok(FeedForwardModule {
        norm: load_layer_norm(tensors, norm_prefix, 1e-5)?,
        linear1: load_linear(tensors, &format!("{ff_prefix}.linear1"))?,
        linear2: load_linear(tensors, &format!("{ff_prefix}.linear2"))?,
    })
}

// ─── Relative Positional Encoding ──────────────────────────────────────────

/// The Transformer-XL relative sinusoid block for one utterance.
///
/// Rows run from offset `+(T-1)` down to `-(T-1)` — descending and symmetric
/// about zero — so the `2T-1` rows cover every relative offset a length-`T`
/// sequence can produce, which is exactly what `rel_shift` then skews into a
/// `T x T` bias. The reference precomputes 5000 positions and slices the middle
/// out; generating the needed rows directly is the same values without the
/// 40 MB table or its length ceiling.
///
/// Returns `(1, 2T-1, d_model)`.
fn relative_positional_encoding(seq_len: usize) -> Result<Array, mlx_rs::error::Exception> {
    let rows = 2 * seq_len - 1;
    let mut pe = vec![0.0f32; rows * D_MODEL];
    let ln10000 = (10000.0f64).ln();
    // `exp(-(ln 10000 / d) * 2i)` for each pair index, hoisted out of the row
    // loop — it is the same for every position.
    let div_term: Vec<f64> = (0..D_MODEL / 2)
        .map(|i| (-(ln10000 / D_MODEL as f64) * (2 * i) as f64).exp())
        .collect();
    for (r, row) in pe.chunks_mut(D_MODEL).enumerate() {
        let pos = (seq_len - 1) as f64 - r as f64;
        for (i, &div) in div_term.iter().enumerate() {
            let angle = pos * div;
            row[2 * i] = angle.sin() as f32;
            row[2 * i + 1] = angle.cos() as f32;
        }
    }
    Ok(Array::from_slice(&pe, &[1, rows as i32, D_MODEL as i32]))
}

// ─── Multi-Head Self-Attention with Relative Position ──────────────────────

struct MultiHeadSelfAttention {
    norm: LayerNorm,
    q_proj: Linear,
    k_proj: Linear,
    v_proj: Linear,
    o_proj: Linear,
    /// Projects the relative sinusoids. Weight-only, like every other linear
    /// in the conformer.
    linear_pos: Linear,
    /// Per-head content and position biases, `[n_heads, head_dim]`.
    pos_bias_u: Array,
    pos_bias_v: Array,
}

impl MultiHeadSelfAttention {
    fn forward(&self, x: &Array, pos_emb: &Array) -> Result<Array, mlx_rs::error::Exception> {
        let shape = x.shape().to_vec();
        let batch = shape[0] as usize;
        let seq_len = shape[1] as usize;

        let h = self.norm.forward(x)?;

        let q = self.q_proj.forward(&h)?;
        let k = self.k_proj.forward(&h)?;
        let v = self.v_proj.forward(&h)?;

        // Transformer-XL relative attention: the query picks up two learned
        // per-head biases before the split — `pos_bias_u` for the content term
        // (against keys) and `pos_bias_v` for the position term (against the
        // projected sinusoids). The previous code modelled the whole thing as a
        // single `pos_bias` *linear* applied to raw `i-j` offsets, which is a
        // different operator with different weights; these two `[8, 128]`
        // vectors are what the checkpoint actually carries.
        let q_heads = ops::reshape(
            &q,
            &[
                batch as i32,
                seq_len as i32,
                NUM_HEADS as i32,
                HEAD_DIM as i32,
            ],
        )?;
        let q_u = ops::transpose_axes(&ops::add(&q_heads, &self.pos_bias_u)?, &[0, 2, 1, 3])?;
        let q_v = ops::transpose_axes(&ops::add(&q_heads, &self.pos_bias_v)?, &[0, 2, 1, 3])?;
        let q = q_u;
        let k = ops::transpose_axes(
            &ops::reshape(
                &k,
                &[
                    batch as i32,
                    seq_len as i32,
                    NUM_HEADS as i32,
                    HEAD_DIM as i32,
                ],
            )?,
            &[0, 2, 1, 3],
        )?;
        let v = ops::transpose_axes(
            &ops::reshape(
                &v,
                &[
                    batch as i32,
                    seq_len as i32,
                    NUM_HEADS as i32,
                    HEAD_DIM as i32,
                ],
            )?,
            &[0, 2, 1, 3],
        )?;

        // Scaled dot-product attention
        let scale = Array::from_f32(1.0 / (HEAD_DIM as f32).sqrt());
        let scores = ops::multiply(
            &ops::matmul(&q, &ops::transpose_axes(&k, &[0, 1, 3, 2])?)?,
            &scale,
        )?;

        // Position term: project the relative sinusoids, score them against
        // `q_v`, then `rel_shift` the (T, 2T-1) matrix down to (T, T) so entry
        // (i, j) lands on the offset i-j.
        let p = self.linear_pos.forward(pos_emb)?;
        let pos_len = p.shape()[1];
        let p = ops::transpose_axes(
            &ops::reshape(
                &p,
                &[batch as i32, pos_len, NUM_HEADS as i32, HEAD_DIM as i32],
            )?,
            &[0, 2, 1, 3],
        )?;
        let matrix_bd = ops::matmul(&q_v, &ops::transpose_axes(&p, &[0, 1, 3, 2])?)?;
        let matrix_bd = Self::rel_shift(&matrix_bd)?;
        let matrix_bd = matrix_bd.index((.., .., .., ..(seq_len as i32)));
        let matrix_bd = ops::multiply(&matrix_bd, &scale)?;

        let scores = ops::add(&scores, &matrix_bd)?;

        let attn = ops::softmax_axis(&scores, -1, None)?;
        let out = ops::matmul(&attn, &v)?;

        // Transpose back and reshape
        let out = ops::transpose_axes(&out, &[0, 2, 1, 3])?;
        let out = ops::reshape(&out, &[batch as i32, seq_len as i32, D_MODEL as i32])?;

        self.o_proj.forward(&out)
    }

    /// Skew a `(B, H, T, 2T-1)` relative-score matrix into `(B, H, T, 2T-1)`
    /// where column `j` of row `i` holds offset `i-j`. Pad one column, reinterpret
    /// the buffer with time and position transposed, drop the row the pad
    /// introduced, and reinterpret back — the standard Transformer-XL trick,
    /// which costs one copy instead of building an index matrix.
    fn rel_shift(x: &Array) -> Result<Array, mlx_rs::error::Exception> {
        let shape = x.shape().to_vec();
        let (b, h, tq, pos_len) = (shape[0], shape[1], shape[2], shape[3]);
        // One column of zeros on the low side of the last axis only.
        let widths: [(i32, i32); 4] = [(0, 0), (0, 0), (0, 0), (1, 0)];
        let padded = ops::pad(x, &widths[..], Array::from_f32(0.0), None)?;
        let reshaped = ops::reshape(&padded, &[b, h, pos_len + 1, tq])?;
        let dropped = reshaped.index((.., .., 1.., ..));
        ops::reshape(&dropped, &[b, h, tq, pos_len])
    }

    fn all_arrays(&self) -> Vec<&Array> {
        let mut v = self.norm.all_arrays();
        v.extend(self.q_proj.all_arrays());
        v.extend(self.k_proj.all_arrays());
        v.extend(self.v_proj.all_arrays());
        v.extend(self.o_proj.all_arrays());
        v.extend(self.linear_pos.all_arrays());
        v.push(&self.pos_bias_u);
        v.push(&self.pos_bias_v);
        v
    }
}

fn load_mhsa(
    tensors: &HashMap<String, Array>,
    norm_prefix: &str,
    attn_prefix: &str,
) -> Result<MultiHeadSelfAttention, InferenceError> {
    Ok(MultiHeadSelfAttention {
        norm: load_layer_norm(tensors, norm_prefix, 1e-5)?,
        q_proj: load_linear(tensors, &format!("{attn_prefix}.linear_q"))?,
        k_proj: load_linear(tensors, &format!("{attn_prefix}.linear_k"))?,
        v_proj: load_linear(tensors, &format!("{attn_prefix}.linear_v"))?,
        o_proj: load_linear(tensors, &format!("{attn_prefix}.linear_out"))?,
        linear_pos: load_linear(tensors, &format!("{attn_prefix}.linear_pos"))?,
        pos_bias_u: get_tensor(tensors, &format!("{attn_prefix}.pos_bias_u"))?,
        pos_bias_v: get_tensor(tensors, &format!("{attn_prefix}.pos_bias_v"))?,
    })
}

// ─── Conformer Convolution Module ──────────────────────────────────────────

/// Convolution module: LayerNorm -> Pointwise Conv -> GLU -> Depthwise Conv -> BatchNorm -> SiLU -> Pointwise Conv
struct ConvModule {
    norm: LayerNorm,
    /// The pointwise and depthwise convolutions are weight-only: the conformer
    /// is configured `use_bias: false`, and only the BatchNorm carries affine
    /// parameters. Demanding a `.bias` for these was part of the original
    /// mismatch.
    pointwise1_weight: Array,
    depthwise_weight: Array,
    batch_norm_weight: Array,
    batch_norm_bias: Array,
    batch_norm_mean: Array,
    batch_norm_var: Array,
    pointwise2_weight: Array,
}

impl ConvModule {
    fn forward(&self, x: &Array) -> Result<Array, mlx_rs::error::Exception> {
        let h = self.norm.forward(x)?;

        // Pointwise conv (1x1): expand channels 1024 -> 2048 for GLU
        let h = ops::conv1d(&h, &self.pointwise1_weight, 1, 0, 1, 1)?;

        // GLU activation: split in half along channel dim, sigmoid gate
        let ch = h.shape()[2] as usize;
        let half = (ch / 2) as i32;
        let gate_input = h.index((.., .., ..half));
        let gate = h.index((.., .., half..));
        let h = ops::multiply(&gate_input, &nn::sigmoid(&gate)?)?;

        // Depthwise conv with padding to preserve length
        // Kernel size is determined by weight shape
        let kernel_size = self.depthwise_weight.shape()[1];
        let pad = (kernel_size - 1) / 2;
        // Depthwise: the weight is `(C_out, K, C_in/groups)`, so its last axis
        // is 1 and the group count is the *input* channel count — not that
        // trailing 1, which would ask for a dense convolution over 1024
        // channels with a 1-channel kernel and fail the shape check.
        let groups = h.shape()[2] / self.depthwise_weight.shape()[2];
        let h = ops::conv1d(&h, &self.depthwise_weight, 1, pad, 1, groups)?;

        // Batch norm (inference mode: use running stats)
        let eps = Array::from_f32(1e-5);
        let bn_std = ops::rsqrt(&ops::add(&self.batch_norm_var, &eps)?)?;
        let h = ops::multiply(&ops::subtract(&h, &self.batch_norm_mean)?, &bn_std)?;
        let h = ops::multiply(&h, &self.batch_norm_weight)?;
        let h = ops::add(&h, &self.batch_norm_bias)?;

        let h = nn::silu(&h)?;

        // Pointwise conv (1x1): project back to d_model
        ops::conv1d(&h, &self.pointwise2_weight, 1, 0, 1, 1)
    }

    fn all_arrays(&self) -> Vec<&Array> {
        let mut v = self.norm.all_arrays();
        v.extend([
            &self.pointwise1_weight,
            &self.depthwise_weight,
            &self.batch_norm_weight,
            &self.batch_norm_bias,
            &self.batch_norm_mean,
            &self.batch_norm_var,
            &self.pointwise2_weight,
        ]);
        v
    }
}

fn load_conv_module(
    tensors: &HashMap<String, Array>,
    norm_prefix: &str,
    conv_prefix: &str,
) -> Result<ConvModule, InferenceError> {
    Ok(ConvModule {
        norm: load_layer_norm(tensors, norm_prefix, 1e-5)?,
        pointwise1_weight: get_tensor(tensors, &format!("{conv_prefix}.pointwise_conv1.weight"))?,
        depthwise_weight: get_tensor(tensors, &format!("{conv_prefix}.depthwise_conv.weight"))?,
        batch_norm_weight: get_tensor(tensors, &format!("{conv_prefix}.batch_norm.weight"))?,
        batch_norm_bias: get_tensor(tensors, &format!("{conv_prefix}.batch_norm.bias"))?,
        batch_norm_mean: get_tensor(tensors, &format!("{conv_prefix}.batch_norm.running_mean"))?,
        batch_norm_var: get_tensor(tensors, &format!("{conv_prefix}.batch_norm.running_var"))?,
        pointwise2_weight: get_tensor(tensors, &format!("{conv_prefix}.pointwise_conv2.weight"))?,
    })
}

// ─── Conformer Layer ───────────────────────────────────────────────────────

/// Single Conformer layer:
/// x + 0.5*ff1(x) -> + mhsa(x) -> + conv(x) -> + 0.5*ff2(x) -> layer_norm
struct ConformerLayer {
    ff1: FeedForwardModule,
    mhsa: MultiHeadSelfAttention,
    conv: ConvModule,
    ff2: FeedForwardModule,
    final_norm: LayerNorm,
}

impl ConformerLayer {
    fn forward(&self, x: &Array, pos_emb: &Array) -> Result<Array, mlx_rs::error::Exception> {
        let half = Array::from_f32(0.5);

        // Macaron FFN (first half-step)
        let ff1_out = self.ff1.forward(x)?;
        let x = ops::add(x, &ops::multiply(&half, &ff1_out)?)?;

        // Multi-head self-attention
        let mhsa_out = self.mhsa.forward(&x, pos_emb)?;
        let x = ops::add(&x, &mhsa_out)?;

        // Convolution module
        let conv_out = self.conv.forward(&x)?;
        let x = ops::add(&x, &conv_out)?;

        // Macaron FFN (second half-step)
        let ff2_out = self.ff2.forward(&x)?;
        let x = ops::add(&x, &ops::multiply(&half, &ff2_out)?)?;

        // Final layer norm
        self.final_norm.forward(&x)
    }

    fn all_arrays(&self) -> Vec<&Array> {
        let mut v = self.ff1.all_arrays();
        v.extend(self.mhsa.all_arrays());
        v.extend(self.conv.all_arrays());
        v.extend(self.ff2.all_arrays());
        v.extend(self.final_norm.all_arrays());
        v
    }
}

/// Each sub-module's LayerNorm is a *sibling* of the module it normalizes
/// (`norm_feed_forward1` / `feed_forward1`, `norm_self_att` / `self_attn`,
/// `norm_conv` / `conv`), and the block's output norm is `norm_out` — not the
/// nested `ff1.norm` / `mhsa.norm` / `final_norm` names this loader used to ask
/// for. `encoder.layers.0.ff1.norm.weight` was simply the first of those to be
/// looked up, which is why it read as one missing tensor rather than a whole
/// naming scheme that did not match.
fn load_conformer_layer(
    tensors: &HashMap<String, Array>,
    prefix: &str,
) -> Result<ConformerLayer, InferenceError> {
    Ok(ConformerLayer {
        ff1: load_ff_module(
            tensors,
            &format!("{prefix}.norm_feed_forward1"),
            &format!("{prefix}.feed_forward1"),
        )?,
        mhsa: load_mhsa(
            tensors,
            &format!("{prefix}.norm_self_att"),
            &format!("{prefix}.self_attn"),
        )?,
        conv: load_conv_module(
            tensors,
            &format!("{prefix}.norm_conv"),
            &format!("{prefix}.conv"),
        )?,
        ff2: load_ff_module(
            tensors,
            &format!("{prefix}.norm_feed_forward2"),
            &format!("{prefix}.feed_forward2"),
        )?,
        final_norm: load_layer_norm(tensors, &format!("{prefix}.norm_out"), 1e-5)?,
    })
}

// ─── Conformer Encoder ─────────────────────────────────────────────────────

/// There is deliberately **no top-level encoder norm**. The checkpoint has none
/// — each block already ends in its own `norm_out` — and the previous loader's
/// `encoder.final_norm` did not exist in any Parakeet conversion.
struct ConformerEncoder {
    subsampling: DepthwiseSubsampling,
    layers: Vec<ConformerLayer>,
}

impl ConformerEncoder {
    /// Forward pass: (batch, time, n_mel) -> (batch, time/8, d_model)
    fn forward(&self, mel: &Array) -> Result<Array, mlx_rs::error::Exception> {
        let mut h = self.subsampling.forward(mel)?;

        // One positional slice for the whole stack: every layer attends over
        // the same sequence, so the sinusoids are computed once and shared.
        let seq_len = h.shape()[1] as usize;
        let pos_emb = relative_positional_encoding(seq_len)?;

        for layer in &self.layers {
            h = layer.forward(&h, &pos_emb)?;
        }
        Ok(h)
    }

    fn all_arrays(&self) -> Vec<&Array> {
        let mut v = self.subsampling.all_arrays();
        for layer in &self.layers {
            v.extend(layer.all_arrays());
        }
        v
    }
}

fn load_encoder(tensors: &HashMap<String, Array>) -> Result<ConformerEncoder, InferenceError> {
    let mut layers = Vec::with_capacity(NUM_ENCODER_LAYERS);
    for i in 0..NUM_ENCODER_LAYERS {
        layers.push(load_conformer_layer(
            tensors,
            &format!("encoder.layers.{i}"),
        )?);
    }

    Ok(ConformerEncoder {
        subsampling: load_subsampling(tensors, "encoder.pre_encode")?,
        layers,
    })
}

// ─── RNN-T Prediction Network (2-layer LSTM) ──────────────────────────────

/// There is **no output projection**. The checkpoint's prediction network is an
/// embedding plus the LSTM stack; the joint's `pred` linear does the projection
/// into joint space. The old `prediction.proj` did not exist.
struct PredictionNetwork {
    embedding: Array, // (vocab_size + 1, pred_hidden)
    lstm_layers: Vec<LstmLayer>,
}

struct LstmLayer {
    wx: Array, // (4*hidden, input_dim)
    wh: Array, // (4*hidden, hidden)
    bias: Option<Array>,
}

impl LstmLayer {
    /// Single LSTM step: (batch, input_dim) + (h, c) -> (h, c)
    fn step(
        &self,
        x: &Array,
        h: &Array,
        c: &Array,
    ) -> Result<(Array, Array), mlx_rs::error::Exception> {
        // gates = x @ Wx^T + h @ Wh^T + bias
        let wx_t = ops::transpose_axes(&self.wx, &[1, 0])?;
        let wh_t = ops::transpose_axes(&self.wh, &[1, 0])?;
        let mut gates = ops::add(&ops::matmul(x, &wx_t)?, &ops::matmul(h, &wh_t)?)?;
        if let Some(ref bias) = self.bias {
            gates = ops::add(&gates, bias)?;
        }

        // Split into i, f, g, o (each of size hidden)
        let hidden = h.shape()[h.shape().len() - 1] as usize;
        let i_gate = nn::sigmoid(gates.index((.., ..(hidden as i32))))?;
        let f_gate = nn::sigmoid(gates.index((.., (hidden as i32)..(2 * hidden as i32))))?;
        let g_gate = ops::tanh(gates.index((.., (2 * hidden as i32)..(3 * hidden as i32))))?;
        let o_gate = nn::sigmoid(gates.index((.., (3 * hidden as i32)..)))?;

        let new_c = ops::add(
            &ops::multiply(&f_gate, c)?,
            &ops::multiply(&i_gate, &g_gate)?,
        )?;
        let new_h = ops::multiply(&o_gate, &ops::tanh(&new_c)?)?;

        Ok((new_h, new_c))
    }

    fn all_arrays(&self) -> Vec<&Array> {
        let mut v = vec![&self.wx, &self.wh];
        if let Some(ref b) = self.bias {
            v.push(b);
        }
        v
    }
}

impl PredictionNetwork {
    /// Step prediction: given previous token, return prediction output and updated states.
    /// token: scalar u32
    /// states: Vec of (h, c) per LSTM layer
    /// Returns: (pred_output (1, joint_dim), new_states)
    /// One decoder step. `token` is `None` at the start of a hypothesis — the
    /// reference feeds a **zero vector**, not the embedding of the blank id, and
    /// the two are different inputs (`embed[8192]` is a trained row).
    fn step(
        &self,
        token: Option<u32>,
        states: &[(Array, Array)],
    ) -> Result<(Array, Vec<(Array, Array)>), mlx_rs::error::Exception> {
        let mut x = match token {
            Some(token) => {
                let tok_arr = Array::from_slice(&[token as i32], &[1]);
                self.embedding.index((tok_arr, ..)) // (1, pred_hidden)
            }
            None => Array::zeros::<f32>(&[1, PRED_HIDDEN as i32])?,
        };
        let mut new_states = Vec::with_capacity(self.lstm_layers.len());

        for (i, lstm) in self.lstm_layers.iter().enumerate() {
            let (h, c) = &states[i];
            let (new_h, new_c) = lstm.step(&x, h, c)?;
            x = new_h.clone();
            new_states.push((new_h, new_c));
        }

        Ok((x, new_states))
    }

    fn initial_states(&self) -> Result<Vec<(Array, Array)>, mlx_rs::error::Exception> {
        let mut states = Vec::with_capacity(self.lstm_layers.len());
        for _ in &self.lstm_layers {
            let h = Array::zeros::<f32>(&[1, PRED_HIDDEN as i32])?;
            let c = Array::zeros::<f32>(&[1, PRED_HIDDEN as i32])?;
            states.push((h, c));
        }
        Ok(states)
    }

    fn all_arrays(&self) -> Vec<&Array> {
        let mut v = vec![&self.embedding];
        for lstm in &self.lstm_layers {
            v.extend(lstm.all_arrays());
        }
        v
    }
}

fn load_prediction_network(
    tensors: &HashMap<String, Array>,
) -> Result<PredictionNetwork, InferenceError> {
    let embedding = get_tensor(tensors, "decoder.prediction.embed.weight")?;

    let mut lstm_layers = Vec::with_capacity(PRED_LAYERS);
    for i in 0..PRED_LAYERS {
        // Capitalized `Wx`/`Wh` — MLX's `nn.LSTM` parameter names, carried
        // verbatim into the checkpoint. The gate order is `i, f, g, o`, which
        // is what `LstmLayer::step` already splits on.
        let pfx = format!("decoder.prediction.dec_rnn.lstm.{i}");
        let wx = get_tensor(tensors, &format!("{pfx}.Wx"))?;
        let wh = get_tensor(tensors, &format!("{pfx}.Wh"))?;
        let bias = tensors.get(&format!("{pfx}.bias")).cloned();
        lstm_layers.push(LstmLayer { wx, wh, bias });
    }

    Ok(PredictionNetwork {
        embedding,
        lstm_layers,
    })
}

// ─── Joint Network ─────────────────────────────────────────────────────────

/// The TDT duration head is **fused into the vocabulary head**, not a separate
/// projection: `joint_net.2` is `[8198, 640]` = 8192 vocab + 1 blank + 5
/// duration bins in one output layer, split after the matmul. The old separate
/// `duration_proj` did not exist in the checkpoint.
struct JointNetwork {
    encoder_proj: Linear,
    pred_proj: Linear,
    out: Linear,
}

impl JointNetwork {
    /// Compute joint output for a single (encoder_frame, pred_output) pair.
    /// encoder_frame: (1, d_model), pred_output: (1, joint_dim)
    /// Returns: (token_logits (1, vocab_size), duration_logits (1, num_durations))
    fn forward(
        &self,
        encoder_frame: &Array,
        pred_output: &Array,
    ) -> Result<(Array, Array), mlx_rs::error::Exception> {
        let enc = self.encoder_proj.forward(encoder_frame)?;
        let pred = self.pred_proj.forward(pred_output)?;
        let joint = nn::relu(&ops::add(&enc, &pred)?)?;
        let logits = self.out.forward(&joint)?;
        // Split the fused head: [0, vocab] inclusive is vocabulary + blank,
        // the tail is one logit per TDT duration bin.
        let split = (VOCAB_SIZE + 1) as i32;
        let token_logits = logits.index((.., ..split));
        let dur_logits = logits.index((.., split..));
        Ok((token_logits, dur_logits))
    }

    fn all_arrays(&self) -> Vec<&Array> {
        let mut v = self.encoder_proj.all_arrays();
        v.extend(self.pred_proj.all_arrays());
        v.extend(self.out.all_arrays());
        v
    }
}

fn load_joint_network(tensors: &HashMap<String, Array>) -> Result<JointNetwork, InferenceError> {
    Ok(JointNetwork {
        encoder_proj: load_linear(tensors, "joint.enc")?,
        pred_proj: load_linear(tensors, "joint.pred")?,
        out: load_linear(tensors, "joint.joint_net.2")?,
    })
}

// ─── Vocabulary ────────────────────────────────────────────────────────────

/// Load BPE vocabulary from tokenizer.vocab (one token per line).
fn load_vocabulary(model_dir: &Path) -> Result<Vec<String>, InferenceError> {
    let vocab_path = model_dir.join("tokenizer.vocab");
    if !vocab_path.exists() {
        // Try tokenizer.json as fallback
        let tokenizer_json_path = model_dir.join("tokenizer.json");
        if tokenizer_json_path.exists() {
            return load_vocabulary_from_json(&tokenizer_json_path);
        }
        return Err(InferenceError::InferenceFailed(
            "neither tokenizer.vocab nor tokenizer.json found".into(),
        ));
    }

    let content = std::fs::read_to_string(&vocab_path)
        .map_err(|e| InferenceError::InferenceFailed(format!("read tokenizer.vocab: {e}")))?;

    // SentencePiece `.vocab` is `<piece>\t<log-prob>` per line. Keeping the
    // whole line put the score into the transcript ("T\t-83est\t-91ing…"),
    // which still *looked* like a decode was happening. Split on the tab; the
    // resulting pieces are byte-identical to the `joint.vocabulary` list in
    // `config.json` that the reference decodes with, so the ids line up.
    let vocab: Vec<String> = content
        .lines()
        .map(|line| line.split('\t').next().unwrap_or(line).to_string())
        .collect();
    if vocab.is_empty() {
        return Err(InferenceError::InferenceFailed(
            "empty vocabulary file".into(),
        ));
    }
    Ok(vocab)
}

fn load_vocabulary_from_json(path: &Path) -> Result<Vec<String>, InferenceError> {
    let content = std::fs::read_to_string(path)
        .map_err(|e| InferenceError::InferenceFailed(format!("read tokenizer.json: {e}")))?;

    let json: serde_json::Value = serde_json::from_str(&content)
        .map_err(|e| InferenceError::InferenceFailed(format!("parse tokenizer.json: {e}")))?;

    // NeMo-style tokenizer.json has "model" -> "vocab" as a list
    if let Some(model) = json.get("model") {
        if let Some(vocab_arr) = model.get("vocab").and_then(|v| v.as_array()) {
            let vocab: Vec<String> = vocab_arr
                .iter()
                .filter_map(|v| v.as_str().map(|s| s.to_string()))
                .collect();
            if !vocab.is_empty() {
                return Ok(vocab);
            }
        }
    }

    // HuggingFace-style: "model" -> "vocab" as object {token: id}
    if let Some(model) = json.get("model") {
        if let Some(vocab_obj) = model.get("vocab").and_then(|v| v.as_object()) {
            let mut pairs: Vec<(String, u64)> = vocab_obj
                .iter()
                .filter_map(|(k, v)| v.as_u64().map(|id| (k.clone(), id)))
                .collect();
            pairs.sort_by_key(|(_, id)| *id);
            let vocab: Vec<String> = pairs.into_iter().map(|(k, _)| k).collect();
            if !vocab.is_empty() {
                return Ok(vocab);
            }
        }
    }

    Err(InferenceError::InferenceFailed(
        "could not extract vocabulary from tokenizer.json".into(),
    ))
}

// ─── TDT Greedy Decoding ──────────────────────────────────────────────────

/// Greedy TDT decode: iterate over encoder frames, emit tokens with duration skips.
/// One non-blank emission from the TDT decoder: the token id, the
/// encoder frame at which it was emitted (start_frame), and the
/// predicted duration in encoder frames. Frame duration in seconds
/// is `HOP_LEN_SAMPLES / SAMPLE_RATE * encoder_subsample_factor` =
/// `160/16000 * 8 = 0.08 s` (80 ms per encoder frame).
#[derive(Debug, Clone)]
pub(crate) struct TokenEmission {
    pub token_id: u32,
    pub start_frame: usize,
    pub duration_frames: usize,
}

/// Seconds per encoder frame. Equal to `HOP_LEN_SAMPLES (10 ms) ×
/// encoder subsample factor (8)`. Kept as `f64` to avoid drift when
/// multiplied by large frame indices on long audio; callers cast to
/// `f32` at the struct boundary.
pub(crate) const FRAME_DURATION_S: f64 = 0.08;

/// True for a vocabulary piece that carries no transcript text: NeMo's control
/// tags (`<|en|>`, `<|pnc|>`, `<|emo:neutral|>`, …) plus `<unk>` and `<pad>`.
/// This checkpoint's vocabulary opens with ~200 of them — a v3 multilingual
/// model emits several per utterance — so leaving them in puts literal
/// `<|startoftranscript|>` text in the transcript.
fn is_special_piece(piece: &str) -> bool {
    (piece.starts_with("<|") && piece.ends_with("|>")) || piece == "<unk>" || piece == "<pad>"
}

/// Greedy TDT decode over the encoder frames.
///
/// The token-and-duration transducer differs from plain RNN-T in that the joint
/// predicts *how far to advance* alongside what to emit, so the cursor jumps by
/// `TDT_DURATIONS[decision]` rather than always by one. Three consequences the
/// previous loop got wrong:
///
/// * the blank is `VOCAB_SIZE`, not 0;
/// * a duration of **zero** is legal and is how several tokens are emitted on
///   one frame, so the cursor is *not* clamped to advance — instead
///   `MAX_SYMBOLS_PER_FRAME` consecutive zero-duration steps force a
///   single-frame advance, which is what guarantees termination;
/// * the decoder state and `last_token` advance only on a non-blank emission.
fn greedy_tdt_decode(
    encoder_output: &Array,
    prediction: &PredictionNetwork,
    joint: &JointNetwork,
    vocab: &[String],
) -> Result<(String, Vec<TokenEmission>), InferenceError> {
    let map_err = |e: mlx_rs::error::Exception| InferenceError::InferenceFailed(e.to_string());

    let num_frames = encoder_output.shape()[1] as usize;

    let mut pred_states = prediction.initial_states().map_err(map_err)?;
    // `None` until the first non-blank: the decoder is fed zeros, not `embed[blank]`.
    let mut last_token: Option<u32> = None;
    let mut emissions: Vec<TokenEmission> = Vec::new();

    let mut time = 0usize;
    let mut new_symbols = 0usize;

    while time < num_frames {
        let enc_frame = encoder_output.index((.., time as i32..time as i32 + 1, ..));
        let enc_frame = ops::reshape(&enc_frame, &[1, D_MODEL as i32]).map_err(map_err)?;

        let (pred_out, new_states) = prediction.step(last_token, &pred_states).map_err(map_err)?;
        let (token_logits, dur_logits) = joint.forward(&enc_frame, &pred_out).map_err(map_err)?;

        let token_logits = ops::reshape(&token_logits, &[-1]).map_err(map_err)?;
        token_logits.eval().map_err(map_err)?;
        let token_id = argmax_f32(token_logits.as_slice()) as u32;

        let dur_logits = ops::reshape(&dur_logits, &[-1]).map_err(map_err)?;
        dur_logits.eval().map_err(map_err)?;
        let decision = argmax_f32(dur_logits.as_slice());
        let duration = TDT_DURATIONS
            .get(decision)
            .copied()
            .unwrap_or(1)
            .min(num_frames);

        if token_id != BLANK_ID {
            last_token = Some(token_id);
            pred_states = new_states;
            emissions.push(TokenEmission {
                token_id,
                start_frame: time,
                duration_frames: duration,
            });
        }

        time += duration;
        new_symbols += 1;
        if duration != 0 {
            new_symbols = 0;
        } else if new_symbols >= MAX_SYMBOLS_PER_FRAME {
            // Only a zero-duration run can stall; force progress rather than
            // spin. Matches the reference's `max_symbols` escape.
            time += 1;
            new_symbols = 0;
        }
    }

    // Special tags are dropped from the text but kept in `emissions`, so word
    // timing still sees the frames they occupied.
    let mut oov_count: usize = 0;
    let text: String = emissions
        .iter()
        .filter_map(|emission| {
            let piece = vocab.get(emission.token_id as usize);
            if piece.is_none() {
                oov_count += 1;
            }
            piece
        })
        .filter(|piece| !is_special_piece(piece))
        .map(|token| token.replace('\u{2581}', " "))
        .collect::<String>()
        .trim()
        .to_string();
    if oov_count > 0 {
        warn!(
            oov_count,
            vocab_size = vocab.len(),
            "parakeet: emitted token id outside vocab range — model/vocab mismatch"
        );
    }

    Ok((text, emissions))
}

/// Group per-token emissions into word-level timing spans. Words are
/// defined by the sentencepiece space marker (`▁`) — the first token
/// of a word carries it. Tokens between markers belong to the same
/// word and contribute to its extent.
///
/// Monotonicity is enforced on `end` (never less than the previous
/// emission's end) so downstream consumers can rely on non-decreasing
/// word boundaries even if the model predicts duration=0 for a token.
pub(crate) fn emissions_to_words(
    emissions: &[TokenEmission],
    vocab: &[String],
) -> Vec<crate::tasks::transcribe::TranscribedWord> {
    use crate::tasks::transcribe::TranscribedWord;

    let mut words: Vec<TranscribedWord> = Vec::new();
    let mut cur_text = String::new();
    let mut cur_start_frame: Option<usize> = None;
    let mut cur_end_frame: usize = 0;

    let flush =
        |words: &mut Vec<TranscribedWord>, start_frame: usize, end_frame: usize, text: String| {
            if text.is_empty() {
                return;
            }
            let start = (start_frame as f64 * FRAME_DURATION_S) as f32;
            let end = (end_frame as f64 * FRAME_DURATION_S) as f32;
            words.push(TranscribedWord { start, end, text });
        };

    for emission in emissions {
        let Some(piece) = vocab.get(emission.token_id as usize).map(String::as_str) else {
            continue;
        };
        let has_marker = piece.starts_with('\u{2581}');
        let clean = piece.trim_start_matches('\u{2581}');

        // Pure `▁` piece: skip — conveys no text, and its timing is
        // covered by the surrounding tokens.
        if clean.is_empty() && has_marker {
            continue;
        }

        let starts_new_word = has_marker || cur_start_frame.is_none();

        if starts_new_word {
            if let Some(start_frame) = cur_start_frame.take() {
                flush(
                    &mut words,
                    start_frame,
                    cur_end_frame,
                    std::mem::take(&mut cur_text),
                );
            }
            cur_end_frame = 0;
        }

        if cur_start_frame.is_none() {
            cur_start_frame = Some(emission.start_frame);
        }
        cur_text.push_str(clean);
        // Monotonic end: never regress even if the model predicts
        // duration=0 for a later token in the same word.
        let tok_end = emission
            .start_frame
            .saturating_add(emission.duration_frames);
        cur_end_frame = cur_end_frame.max(tok_end);
    }

    if let Some(start_frame) = cur_start_frame {
        flush(&mut words, start_frame, cur_end_frame, cur_text);
    }

    words
}

fn argmax_f32(data: &[f32]) -> usize {
    let mut best_idx = 0;
    let mut best_val = f32::NEG_INFINITY;
    for (i, &v) in data.iter().enumerate() {
        if v > best_val {
            best_val = v;
            best_idx = i;
        }
    }
    best_idx
}

// ─── ParakeetBackend ───────────────────────────────────────────────────────

/// Native Parakeet-TDT speech-to-text backend for Apple Silicon.
pub struct ParakeetBackend {
    encoder: ConformerEncoder,
    prediction: PredictionNetwork,
    joint: JointNetwork,
    vocab: Vec<String>,
}

// SAFETY: ParakeetBackend is only accessed through RwLock in InferenceEngine.
// mlx_rs::Array contains native handles which are not auto-Send/Sync.
unsafe impl Send for ParakeetBackend {}
unsafe impl Sync for ParakeetBackend {}

impl ParakeetBackend {
    /// Load Parakeet-TDT model from a directory containing model.safetensors and tokenizer.vocab.
    pub fn load(model_dir: &Path) -> Result<Self, InferenceError> {
        info!(model_dir = %model_dir.display(), "loading Parakeet-TDT model via MLX");

        // Pick the default MLX device (same logic as MlxBackend)
        #[cfg(feature = "mlx-metal")]
        let default_device = mlx_rs::Device::gpu();
        #[cfg(not(feature = "mlx-metal"))]
        let default_device = mlx_rs::Device::cpu();

        match std::env::var("CAR_MLX_DEVICE").ok().as_deref() {
            Some("cpu") => mlx_rs::Device::set_default(&mlx_rs::Device::cpu()),
            #[cfg(feature = "mlx-metal")]
            Some("gpu") => mlx_rs::Device::set_default(&mlx_rs::Device::gpu()),
            _ => mlx_rs::Device::set_default(&default_device),
        }

        // Load vocabulary
        let vocab = load_vocabulary(model_dir)?;
        info!(vocab_size = vocab.len(), "vocabulary loaded");

        // Load safetensors weights
        info!("loading safetensors weights");
        let tensors = load_all_tensors(model_dir)?;
        info!(tensors = tensors.len(), "tensors loaded");

        // Build model components
        let encoder = load_encoder(&tensors)?;
        info!(layers = NUM_ENCODER_LAYERS, "conformer encoder loaded");

        let prediction = load_prediction_network(&tensors)?;
        info!(
            lstm_layers = PRED_LAYERS,
            hidden = PRED_HIDDEN,
            "prediction network loaded"
        );

        let joint = load_joint_network(&tensors)?;
        info!("joint network loaded");

        // Evaluate all weights to materialize on device
        let mut all_params = encoder.all_arrays();
        all_params.extend(prediction.all_arrays());
        all_params.extend(joint.all_arrays());
        mlx_rs::transforms::eval(all_params)
            .map_err(|e| InferenceError::InferenceFailed(format!("eval weights: {e}")))?;

        info!("Parakeet-TDT model loaded successfully");
        Ok(Self {
            encoder,
            prediction,
            joint,
            vocab,
        })
    }

    /// Transcribe a 16kHz WAV audio file to text.
    pub fn transcribe(&self, audio_path: &Path) -> Result<String, InferenceError> {
        let (text, _words) = self.transcribe_detailed(audio_path)?;
        Ok(text)
    }

    /// Transcribe with per-word timing. Returns `(text, words)` where
    /// words are sentencepiece groupings of TDT token emissions.
    /// Empty `words` is not an error — it just means no non-blank tokens
    /// emerged (silent audio).
    pub fn transcribe_detailed(
        &self,
        audio_path: &Path,
    ) -> Result<(String, Vec<crate::tasks::transcribe::TranscribedWord>), InferenceError> {
        info!(path = %audio_path.display(), "transcribing audio (detailed)");

        let samples = load_wav(audio_path)?;
        if samples.is_empty() {
            return Ok((String::new(), Vec::new()));
        }
        info!(
            samples = samples.len(),
            duration_secs = samples.len() as f32 / SAMPLE_RATE as f32,
            "audio loaded"
        );

        let mel = compute_log_mel(&samples)?;
        let mel_frames = mel.shape()[1] as usize;
        info!(mel_frames = mel_frames, "mel spectrogram computed");

        let map_err = |e: mlx_rs::error::Exception| InferenceError::InferenceFailed(e.to_string());
        let encoder_output = self.encoder.forward(&mel).map_err(map_err)?;
        encoder_output.eval().map_err(map_err)?;
        let enc_frames = encoder_output.shape()[1] as usize;
        info!(encoder_frames = enc_frames, "encoder forward complete");

        let (text, emissions) =
            greedy_tdt_decode(&encoder_output, &self.prediction, &self.joint, &self.vocab)?;
        let words = emissions_to_words(&emissions, &self.vocab);

        info!(
            text_len = text.len(),
            word_count = words.len(),
            "transcription complete"
        );
        Ok((text, words))
    }

    /// Transcribe from raw 16kHz f32 PCM samples (no WAV header needed).
    pub fn transcribe_samples(&self, samples: &[f32]) -> Result<String, InferenceError> {
        let (text, _words) = self.transcribe_samples_detailed(samples)?;
        Ok(text)
    }

    /// Like [`transcribe_samples`] but also returns per-word timing spans.
    pub fn transcribe_samples_detailed(
        &self,
        samples: &[f32],
    ) -> Result<(String, Vec<crate::tasks::transcribe::TranscribedWord>), InferenceError> {
        if samples.is_empty() {
            return Ok((String::new(), Vec::new()));
        }

        let mel = compute_log_mel(samples)?;
        let map_err = |e: mlx_rs::error::Exception| InferenceError::InferenceFailed(e.to_string());
        let encoder_output = self.encoder.forward(&mel).map_err(map_err)?;
        encoder_output.eval().map_err(map_err)?;

        let (text, emissions) =
            greedy_tdt_decode(&encoder_output, &self.prediction, &self.joint, &self.vocab)?;
        let words = emissions_to_words(&emissions, &self.vocab);
        Ok((text, words))
    }

    /// Return the vocabulary size.
    pub fn vocab_size(&self) -> usize {
        self.vocab.len()
    }

    /// Return the expected sample rate (16000 Hz).
    pub fn sample_rate(&self) -> usize {
        SAMPLE_RATE
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn vocab() -> Vec<String> {
        // Indices: 0:hello (no marker, first-word fallback)
        //          1:▁world   2:▁   (pure marker, should be skipped)
        //          3:!       (continuation)     4:▁foo   5:bar
        vec![
            "hello".into(),
            "\u{2581}world".into(),
            "\u{2581}".into(),
            "!".into(),
            "\u{2581}foo".into(),
            "bar".into(),
        ]
    }

    fn emit(token_id: u32, start: usize, dur: usize) -> TokenEmission {
        TokenEmission {
            token_id,
            start_frame: start,
            duration_frames: dur,
        }
    }

    #[test]
    fn first_token_without_marker_starts_a_word() {
        let v = vocab();
        let em = vec![emit(0, 0, 5), emit(1, 5, 10)]; // hello, ▁world
        let words = emissions_to_words(&em, &v);
        assert_eq!(words.len(), 2);
        assert_eq!(words[0].text, "hello");
        assert!((words[0].start - 0.0).abs() < 1e-6);
        assert!((words[0].end - (5.0 * 0.08)).abs() < 1e-5);
        assert_eq!(words[1].text, "world");
    }

    #[test]
    fn pure_marker_token_is_skipped() {
        let v = vocab();
        // ▁foo, then lone ▁ (id=2) which should be dropped, then bar
        let em = vec![emit(4, 0, 4), emit(2, 4, 1), emit(5, 5, 3)];
        let words = emissions_to_words(&em, &v);
        assert_eq!(words.len(), 1, "pure ▁ should not split or emit a word");
        assert_eq!(words[0].text, "foobar");
    }

    #[test]
    fn punctuation_attaches_to_previous_word() {
        let v = vocab();
        // ▁world, !
        let em = vec![emit(1, 0, 4), emit(3, 4, 1)];
        let words = emissions_to_words(&em, &v);
        assert_eq!(words.len(), 1);
        assert_eq!(words[0].text, "world!");
    }

    #[test]
    fn zero_duration_does_not_regress_end() {
        let v = vocab();
        // hello at frame 0 dur 5; continuation `!` at frame 5 dur 0
        let em = vec![emit(0, 0, 5), emit(3, 5, 0)];
        let words = emissions_to_words(&em, &v);
        assert_eq!(words.len(), 1);
        assert!(words[0].end >= words[0].start);
        assert!((words[0].end - (5.0 * 0.08)).abs() < 1e-5);
    }

    #[test]
    fn empty_emissions_yields_empty() {
        let v = vocab();
        let words = emissions_to_words(&[], &v);
        assert!(words.is_empty());
    }

    #[test]
    fn out_of_vocab_token_is_silently_dropped() {
        let v = vocab();
        let em = vec![emit(0, 0, 5), emit(99, 5, 3), emit(1, 8, 4)];
        let words = emissions_to_words(&em, &v);
        assert_eq!(words.len(), 2);
        assert_eq!(words[0].text, "hello");
        assert_eq!(words[1].text, "world");
    }
}

/// Stage-by-stage comparison against the reference implementation.
///
/// The hazard this guards is the one car#660 called out: a loader wired up from
/// mismatched tensors will happily load and then emit garbage, which is worse
/// than a clean failure. Asserting only on the final transcript finds that late
/// and says nothing about *where* it went wrong, so these compare the mel front
/// end and the encoder output against tensors dumped from the `mlx_audio`
/// reference — the first stage that diverges is the one that is wrong.
///
/// Fixtures are not in the repo: they need the 2.4 GB checkpoint and the
/// managed Python runtime. Generate them with `scripts/parakeet-golden-dump.py`
/// and point `CAR_PARAKEET_GOLDEN` at the output directory; without it these
/// skip, so CI stays green and a developer on an Apple Silicon box can run the
/// real check on demand. Full context, including the layout bug these caught:
/// `docs/solutions/parakeet-native-mlx-port.md`.
#[cfg(all(test, target_os = "macos", target_arch = "aarch64", not(car_skip_mlx)))]
mod golden_tests {
    use super::*;

    struct Golden {
        dir: std::path::PathBuf,
        model_dir: std::path::PathBuf,
    }

    fn fixtures() -> Option<Golden> {
        let dir = std::path::PathBuf::from(std::env::var("CAR_PARAKEET_GOLDEN").ok()?);
        let model_dir = dirs_home()?.join(".car/models/Parakeet-TDT-0.6B-v3-MLX");
        if !dir.join("mel.bin").exists() || !model_dir.join("model.safetensors").exists() {
            return None;
        }
        Some(Golden { dir, model_dir })
    }

    fn dirs_home() -> Option<std::path::PathBuf> {
        std::env::var_os("HOME").map(std::path::PathBuf::from)
    }

    /// Read a `<name>.shape` + `<name>.bin` pair written by the golden harness.
    fn load_golden(g: &Golden, name: &str) -> (Vec<usize>, Vec<f32>) {
        let shape: Vec<usize> = serde_json::from_str(
            &std::fs::read_to_string(g.dir.join(format!("{name}.shape"))).unwrap(),
        )
        .unwrap();
        let bytes = std::fs::read(g.dir.join(format!("{name}.bin"))).unwrap();
        let data: Vec<f32> = bytes
            .chunks_exact(4)
            .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
            .collect();
        (shape, data)
    }

    /// Max absolute difference plus a correlation-free relative scale, so a
    /// systematically shifted tensor cannot pass on absolute tolerance alone.
    fn compare(name: &str, got: &[f32], want: &[f32], tol: f32) {
        assert_eq!(
            got.len(),
            want.len(),
            "{name}: element count differs — got {}, want {}",
            got.len(),
            want.len()
        );
        let mut max_abs = 0.0f32;
        let mut sum_sq_err = 0.0f64;
        let mut sum_sq_ref = 0.0f64;
        for (&a, &b) in got.iter().zip(want) {
            max_abs = max_abs.max((a - b).abs());
            sum_sq_err += ((a - b) as f64).powi(2);
            sum_sq_ref += (b as f64).powi(2);
        }
        let rel = (sum_sq_err / sum_sq_ref.max(1e-30)).sqrt();
        assert!(
            max_abs < tol,
            "{name}: max|diff| = {max_abs:.6} (tol {tol}), relative L2 = {rel:.6}. \
             This stage does not match the reference — fix it before looking at \
             any later stage."
        );
    }

    #[test]
    fn mel_front_end_matches_the_reference() {
        let Some(g) = fixtures() else {
            eprintln!("skipping: set CAR_PARAKEET_GOLDEN and install the checkpoint");
            return;
        };
        let samples = load_wav(&g.dir.join("local.wav")).expect("golden wav");
        let mel = compute_log_mel(&samples).expect("mel");
        mel.eval().unwrap();

        let (shape, want) = load_golden(&g, "mel");
        assert_eq!(
            mel.shape().iter().map(|&d| d as usize).collect::<Vec<_>>(),
            shape,
            "mel shape differs — framing or centring is wrong"
        );
        // Normalized log-mel is ~unit variance, so 1e-3 is a tight bound here.
        compare("mel", mel.as_slice(), &want, 1e-3);
    }

    #[test]
    fn encoder_and_transcript_match_the_reference() {
        let Some(g) = fixtures() else {
            eprintln!("skipping: set CAR_PARAKEET_GOLDEN and install the checkpoint");
            return;
        };
        let backend = ParakeetBackend::load(&g.model_dir).expect("the checkpoint must load");

        let samples = load_wav(&g.dir.join("local.wav")).expect("golden wav");
        let mel = compute_log_mel(&samples).expect("mel");
        let encoder_output = backend.encoder.forward(&mel).expect("encoder");
        encoder_output.eval().unwrap();

        let (shape, want) = load_golden(&g, "encoder");
        assert_eq!(
            encoder_output
                .shape()
                .iter()
                .map(|&d| d as usize)
                .collect::<Vec<_>>(),
            shape,
            "encoder shape differs — subsampling geometry is wrong"
        );
        // Encoder activations have std ~0.02, so this bound is ~10% of a
        // standard deviation: loose enough for f32 reassociation across 24
        // layers, far tighter than any real architectural mismatch survives.
        compare("encoder", encoder_output.as_slice(), &want, 2e-3);

        let (text, _) = backend
            .transcribe_samples_detailed(&samples)
            .expect("decode");
        let want_text = std::fs::read_to_string(g.dir.join("transcript.txt")).unwrap();
        assert_eq!(
            text.trim(),
            want_text.trim(),
            "the transcript must match the reference, not merely be non-empty"
        );
    }
}