//! Phase 13.9 — candle port of `facebook/wav2vec2-base-960h` for
//! word-level CTC forced alignment.
//!
//! Builds on the patterns from the Phase 13.8 emotion2vec port (`emotion2vec.rs`)
//! but with the wav2vec2-specific differences: POST-norm transformer
//! blocks, separate Q/K/V linears (not fused), GroupNorm on the first
//! feature-extractor layer only, single conv positional embedding
//! (kernel 128) instead of a 5-stack, and a 32-char CTC output head.
//!
//! Architecture (full notes in `docs/wav2vec2_port_notes.md`):
//! ```
//! audio 16 kHz (B, 1, T)
//!  → feature_extractor   7 × Conv1d, T → T/320
//!  → feature_projection  LayerNorm(512) + Linear(512 → 768)
//!  → conv_pos_embedding  Conv1d(768→768, kernel 128, groups 16) + GELU; bias
//!  → encoder.layers.0..11  12 × POST-norm transformer block
//!  → lm_head             Linear(768 → 32)
//!  → log_softmax → 32-char CTC log-probs (B, T', 32)
//! ```
//!
//! Slice plan:
//! - **Slice 1**: inspector + design notes ✅
//! - **Slice 2 (this module)**: full candle port with safetensors loader
//! - **Slice 3**: greedy CTC decode → ASR transcript smoke test
//! - **Slice 4**: Viterbi forced alignment given a known transcript

use crate::error::{CsmError, Result as CsmResult};
use candle_core::{Device, Module, Tensor};
use candle_nn::{
    Conv1d, Conv1dConfig, GroupNorm, LayerNorm, Linear, VarBuilder, conv1d, conv1d_no_bias,
    group_norm, layer_norm, linear,
};

/// Architecture hyperparameters from `config.json`. Only the fields
/// needed by the inference path are surfaced.
#[derive(Debug, Clone)]
pub struct Wav2Vec2Config {
    pub embed_dim: usize,
    pub num_heads: usize,
    pub mlp_dim: usize,
    pub num_layers: usize,
    pub vocab_size: usize,
    pub norm_eps: f64,
    pub feature_dim: usize,
    /// `(out_channels, kernel, stride)` per conv layer.
    pub conv_layers: Vec<(usize, usize, usize)>,
    /// Conv positional embedding kernel (128 for `_base`).
    pub conv_pos_kernel: usize,
    /// Conv positional embedding groups (16 for `_base`).
    pub conv_pos_groups: usize,
}

impl Wav2Vec2Config {
    pub fn base_960h() -> Self {
        Self {
            embed_dim: 768,
            num_heads: 12,
            mlp_dim: 3072,
            num_layers: 12,
            vocab_size: 32,
            norm_eps: 1e-5,
            feature_dim: 512,
            conv_layers: vec![
                (512, 10, 5),
                (512, 3, 2),
                (512, 3, 2),
                (512, 3, 2),
                (512, 3, 2),
                (512, 2, 2),
                (512, 2, 2),
            ],
            conv_pos_kernel: 128,
            conv_pos_groups: 16,
        }
    }

    pub fn conv_stride_total(&self) -> usize {
        self.conv_layers.iter().map(|(_, _, s)| *s).product()
    }
}

/// One conv layer in the feature extractor. Layer 0 in `_base` carries a
/// `GroupNorm` with `num_groups = num_channels` (= out_conv_dim, 512) —
/// per HF's wav2vec2 with `feat_extract_norm: "group"`. The state-dict
/// param name is `layer_norm.weight/bias` purely because of the wrapping
/// `Wav2Vec2GroupNormConvLayer.layer_norm` attribute name; the *operation*
/// is GroupNorm, not LayerNorm. Affine shape `[512]` matches both ops, so
/// loaders that picked LayerNorm here produced empty/garbage output.
/// Layers 1..6 are bare `Conv1d → GELU`.
#[derive(Debug)]
struct FeatureExtractorBlock {
    conv: Conv1d,
    /// Only present on layer 0 for `_base`. GroupNorm with
    /// num_groups = num_channels = InstanceNorm semantics.
    norm: Option<GroupNorm>,
}

impl FeatureExtractorBlock {
    fn new(
        in_channels: usize,
        out_channels: usize,
        kernel_size: usize,
        stride: usize,
        norm_eps: f64,
        with_norm: bool,
        vb: VarBuilder,
    ) -> CsmResult<Self> {
        let cfg = Conv1dConfig {
            stride,
            ..Default::default()
        };
        let conv = conv1d_no_bias(in_channels, out_channels, kernel_size, cfg, vb.pp("conv"))
            .map_err(|e| CsmError::Config(format!("feat ext conv: {e}")))?;
        let norm = if with_norm {
            // num_groups = num_channels per HF wav2vec2 group-norm config.
            Some(
                group_norm(out_channels, out_channels, norm_eps, vb.pp("layer_norm"))
                    .map_err(|e| CsmError::Config(format!("feat ext group_norm: {e}")))?,
            )
        } else {
            None
        };
        Ok(Self { conv, norm })
    }
}

impl Module for FeatureExtractorBlock {
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        // (B, C_in, T_in) → (B, C_out, T_out)
        let xs = self.conv.forward(xs)?;
        let xs = if let Some(n) = &self.norm {
            // GroupNorm operates on (B, C, T) directly — no transpose.
            n.forward(&xs)?
        } else {
            xs
        };
        xs.gelu()
    }
}

/// 7-layer feature extractor: raw 16 kHz mono → `(B, 512, T/320)`.
#[derive(Debug)]
pub struct FeatureExtractor {
    blocks: Vec<FeatureExtractorBlock>,
}

impl FeatureExtractor {
    pub fn new(cfg: &Wav2Vec2Config, vb: VarBuilder) -> CsmResult<Self> {
        let mut blocks = Vec::with_capacity(cfg.conv_layers.len());
        let mut in_ch = 1usize;
        for (i, (out_ch, k, s)) in cfg.conv_layers.iter().enumerate() {
            let with_norm = i == 0; // `_base` only norms layer 0.
            let block = FeatureExtractorBlock::new(
                in_ch,
                *out_ch,
                *k,
                *s,
                cfg.norm_eps,
                with_norm,
                vb.pp("conv_layers").pp(i.to_string()),
            )?;
            blocks.push(block);
            in_ch = *out_ch;
        }
        Ok(Self { blocks })
    }
}

impl Module for FeatureExtractor {
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        let mut xs = xs.clone();
        for block in self.blocks.iter() {
            xs = block.forward(&xs)?;
        }
        Ok(xs)
    }
}

/// `LayerNorm(feature_dim) + Linear(feature_dim → embed_dim)` projecting
/// the feature extractor output up to the transformer dim.
#[derive(Debug)]
pub struct FeatureProjection {
    norm: LayerNorm,
    proj: Linear,
}

impl FeatureProjection {
    pub fn new(in_dim: usize, out_dim: usize, norm_eps: f64, vb: VarBuilder) -> CsmResult<Self> {
        let norm = layer_norm(in_dim, norm_eps, vb.pp("layer_norm"))
            .map_err(|e| CsmError::Config(format!("feat proj norm: {e}")))?;
        let proj = linear(in_dim, out_dim, vb.pp("projection"))
            .map_err(|e| CsmError::Config(format!("feat proj linear: {e}")))?;
        Ok(Self { norm, proj })
    }
}

impl Module for FeatureProjection {
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        // xs: (B, T, in_dim)  — channel-last
        let xs = self.norm.forward(xs)?;
        self.proj.forward(&xs)
    }
}

/// Conv-based positional embedding: a single grouped Conv1d (kernel 128,
/// groups 16) with same-padding, GELU activation, output added to the
/// input as a positional bias. The pickle key prefix is
/// `encoder.pos_conv_embed.conv` (Conv1d only; the GELU has no params).
///
/// Note: for an EVEN kernel (128), exact same-padding requires asymmetric
/// padding. Standard transformers-style impl pads both sides by 64 and
/// drops the last frame to keep T unchanged.
#[derive(Debug)]
pub struct ConvPosEmbedding {
    conv: Conv1d,
    drop_last: bool,
}

impl ConvPosEmbedding {
    /// Build with random init (used by tests). Production code path is
    /// [`Self::from_weight_normed`] because the upstream pickle stores
    /// `weight_g` + `weight_v` instead of a materialized `weight`.
    pub fn new(embed_dim: usize, kernel: usize, groups: usize, vb: VarBuilder) -> CsmResult<Self> {
        let (pad, drop_last) = pad_for(kernel);
        let cfg = Conv1dConfig {
            padding: pad,
            stride: 1,
            dilation: 1,
            groups,
            cudnn_fwd_algo: None,
        };
        let conv = conv1d(embed_dim, embed_dim, kernel, cfg, vb.pp("conv"))
            .map_err(|e| CsmError::Config(format!("pos conv: {e}")))?;
        Ok(Self { conv, drop_last })
    }

    /// Materialize a weight-normed Conv1d from upstream's `weight_g` +
    /// `weight_v` storage. fairseq's wav2vec2 applies
    /// `nn.utils.weight_norm(self.pos_conv, name="weight", dim=2)` so the
    /// stored shapes are:
    ///   weight_v: (out, in/groups, kernel)
    ///   weight_g: (1, 1, kernel)         ← per-kernel-position gain
    /// And the materialized weight is `weight_g * weight_v / ||weight_v||_2`
    /// where the L2 norm is along axis 2 (kernel), keepdim=true so it
    /// broadcasts back to `weight_v`'s shape.
    pub fn from_weight_normed(
        embed_dim: usize,
        kernel: usize,
        groups: usize,
        vb: VarBuilder,
    ) -> CsmResult<Self> {
        let (pad, drop_last) = pad_for(kernel);
        let cfg = Conv1dConfig {
            padding: pad,
            stride: 1,
            dilation: 1,
            groups,
            cudnn_fwd_algo: None,
        };
        let v_shape = (embed_dim, embed_dim / groups, kernel);
        let weight_v = vb
            .pp("conv")
            .get(v_shape, "weight_v")
            .map_err(|e| CsmError::Config(format!("pos conv weight_v: {e}")))?;
        let weight_g = vb
            .pp("conv")
            .get((1, 1, kernel), "weight_g")
            .map_err(|e| CsmError::Config(format!("pos conv weight_g: {e}")))?;
        let bias = vb
            .pp("conv")
            .get(embed_dim, "bias")
            .map_err(|e| CsmError::Config(format!("pos conv bias: {e}")))?;
        // Materialize: weight = weight_g * weight_v / (||weight_v||_dim=2 + eps)
        // Eps mirrors PyTorch's weight_norm numerical guard and keeps
        // random-init tests numerically stable.
        let norm = weight_v
            .sqr()
            .and_then(|t| t.sum_keepdim(2))
            .and_then(|t| t.sqrt())
            .and_then(|t| (t + 1e-12)?.broadcast_as(weight_v.shape()))
            .map_err(|e| CsmError::Config(format!("pos conv weight norm: {e}")))?;
        let weight = (weight_v / &norm)
            .and_then(|t| t.broadcast_mul(&weight_g))
            .map_err(|e| CsmError::Config(format!("pos conv materialize: {e}")))?;
        let conv = Conv1d::new(weight, Some(bias), cfg);
        Ok(Self { conv, drop_last })
    }
}

/// Compute `(padding, drop_last_frame)` for same-padding given a kernel.
/// Even kernels need asymmetric handling: pad symmetrically by k/2 and
/// trim the trailing extra frame.
fn pad_for(kernel: usize) -> (usize, bool) {
    if kernel.is_multiple_of(2) {
        (kernel / 2, true)
    } else {
        ((kernel - 1) / 2, false)
    }
}

impl Module for ConvPosEmbedding {
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        // xs: (B, T, C). Conv1d operates on (B, C, T).
        let h = xs.transpose(1, 2)?.contiguous()?;
        let h = self.conv.forward(&h)?;
        let h = if self.drop_last {
            // Trim trailing frame from even-kernel padding.
            let t = h.dim(2)?;
            h.narrow(2, 0, t - 1)?
        } else {
            h
        };
        let h = h.gelu()?;
        let bias = h.transpose(1, 2)?.contiguous()?;
        xs + bias
    }
}

/// One POST-norm transformer block. Pickle keys per block (16 tensors):
/// `attention.{q,k,v,out}_proj.{weight,bias}`,
/// `layer_norm.{weight,bias}`,
/// `feed_forward.{intermediate,output}_dense.{weight,bias}`,
/// `final_layer_norm.{weight,bias}`.
#[derive(Debug)]
pub struct Block {
    q: Linear,
    k: Linear,
    v: Linear,
    out: Linear,
    layer_norm: LayerNorm,
    fc1: Linear,
    fc2: Linear,
    final_layer_norm: LayerNorm,
    num_heads: usize,
    head_dim: usize,
    scale: f64,
}

impl Block {
    pub fn new(
        embed_dim: usize,
        num_heads: usize,
        mlp_dim: usize,
        norm_eps: f64,
        vb: VarBuilder,
    ) -> CsmResult<Self> {
        let head_dim = embed_dim / num_heads;
        let scale = 1.0 / (head_dim as f64).sqrt();
        let q = linear(embed_dim, embed_dim, vb.pp("attention").pp("q_proj"))
            .map_err(|e| CsmError::Config(format!("block q: {e}")))?;
        let k = linear(embed_dim, embed_dim, vb.pp("attention").pp("k_proj"))
            .map_err(|e| CsmError::Config(format!("block k: {e}")))?;
        let v = linear(embed_dim, embed_dim, vb.pp("attention").pp("v_proj"))
            .map_err(|e| CsmError::Config(format!("block v: {e}")))?;
        let out = linear(embed_dim, embed_dim, vb.pp("attention").pp("out_proj"))
            .map_err(|e| CsmError::Config(format!("block out: {e}")))?;
        let mid_norm = layer_norm(embed_dim, norm_eps, vb.pp("layer_norm"))
            .map_err(|e| CsmError::Config(format!("block ln: {e}")))?;
        let fc1 = linear(
            embed_dim,
            mlp_dim,
            vb.pp("feed_forward").pp("intermediate_dense"),
        )
        .map_err(|e| CsmError::Config(format!("block fc1: {e}")))?;
        let fc2 = linear(mlp_dim, embed_dim, vb.pp("feed_forward").pp("output_dense"))
            .map_err(|e| CsmError::Config(format!("block fc2: {e}")))?;
        let final_norm = layer_norm(embed_dim, norm_eps, vb.pp("final_layer_norm"))
            .map_err(|e| CsmError::Config(format!("block fln: {e}")))?;
        Ok(Self {
            q,
            k,
            v,
            out,
            layer_norm: mid_norm,
            fc1,
            fc2,
            final_layer_norm: final_norm,
            num_heads,
            head_dim,
            scale,
        })
    }

    fn attention(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        // xs: (B, T, D). Standard separate-QKV multi-head attention.
        let (b, t, _d) = xs.dims3()?;
        let q = self
            .q
            .forward(xs)?
            .reshape((b, t, self.num_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?;
        let k = self
            .k
            .forward(xs)?
            .reshape((b, t, self.num_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?;
        let v = self
            .v
            .forward(xs)?
            .reshape((b, t, self.num_heads, self.head_dim))?
            .transpose(1, 2)?
            .contiguous()?;
        // Collapse (B, H) for the Metal-friendly 3D matmul (Phase 8.8 Moonshine
        // workaround for candle's 4D matmul shape bug).
        let bh = b * self.num_heads;
        let q3 = q.reshape((bh, t, self.head_dim))?;
        let k3 = k.reshape((bh, t, self.head_dim))?;
        let v3 = v.reshape((bh, t, self.head_dim))?;
        let scores = (q3.matmul(&k3.transpose(1, 2)?)? * self.scale)?;
        let attn = candle_nn::ops::softmax_last_dim(&scores)?;
        let out = attn.matmul(&v3)?;
        let out = out
            .reshape((b, self.num_heads, t, self.head_dim))?
            .transpose(1, 2)?
            .reshape((b, t, self.num_heads * self.head_dim))?
            .contiguous()?;
        self.out.forward(&out)
    }

    fn ffn(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        let xs = self.fc1.forward(xs)?;
        let xs = xs.gelu()?;
        self.fc2.forward(&xs)
    }
}

impl Module for Block {
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        // POST-norm: x = layer_norm(x + attn(x)); x = final_layer_norm(x + ffn(x))
        let h = self.attention(xs)?;
        let xs = (xs + h)?;
        let xs = self.layer_norm.forward(&xs)?;
        let h = self.ffn(&xs)?;
        let xs = (xs + h)?;
        self.final_layer_norm.forward(&xs)
    }
}

/// 12-block encoder + the ConvPosEmbedding bias added before block 0.
/// State-dict keys live under `wav2vec2.encoder.*`:
///   `encoder.pos_conv_embed.conv.{weight,bias}`
///   `encoder.layer_norm.{weight,bias}`  ← initial LN before blocks
///   `encoder.layers.{0..11}.*`
#[derive(Debug)]
pub struct Encoder {
    pos: ConvPosEmbedding,
    init_norm: LayerNorm,
    blocks: Vec<Block>,
}

impl Encoder {
    pub fn new(cfg: &Wav2Vec2Config, vb: VarBuilder) -> CsmResult<Self> {
        let pos = ConvPosEmbedding::from_weight_normed(
            cfg.embed_dim,
            cfg.conv_pos_kernel,
            cfg.conv_pos_groups,
            vb.pp("pos_conv_embed"),
        )?;
        let init_norm = layer_norm(cfg.embed_dim, cfg.norm_eps, vb.pp("layer_norm"))
            .map_err(|e| CsmError::Config(format!("encoder init_norm: {e}")))?;
        let mut blocks = Vec::with_capacity(cfg.num_layers);
        let vb_layers = vb.pp("layers");
        for i in 0..cfg.num_layers {
            blocks.push(Block::new(
                cfg.embed_dim,
                cfg.num_heads,
                cfg.mlp_dim,
                cfg.norm_eps,
                vb_layers.pp(i.to_string()),
            )?);
        }
        Ok(Self {
            pos,
            init_norm,
            blocks,
        })
    }
}

impl Module for Encoder {
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        // xs: (B, T, D)
        let xs = self.pos.forward(xs)?;
        let mut xs = self.init_norm.forward(&xs)?;
        for block in self.blocks.iter() {
            xs = block.forward(&xs)?;
        }
        Ok(xs)
    }
}

/// Top-level wav2vec2 + CTC head. State-dict prefix: everything below
/// `wav2vec2.*` (feature_extractor, feature_projection, encoder), plus
/// `lm_head.{weight,bias}` at the root.
#[derive(Debug)]
pub struct Wav2Vec2 {
    cfg: Wav2Vec2Config,
    feature_extractor: FeatureExtractor,
    feature_projection: FeatureProjection,
    encoder: Encoder,
    lm_head: Linear,
    device: Device,
}

impl Wav2Vec2 {
    pub fn new(cfg: Wav2Vec2Config, vb: VarBuilder, device: Device) -> CsmResult<Self> {
        let vb_w = vb.pp("wav2vec2");
        let feature_extractor = FeatureExtractor::new(&cfg, vb_w.pp("feature_extractor"))?;
        let feature_projection = FeatureProjection::new(
            cfg.feature_dim,
            cfg.embed_dim,
            cfg.norm_eps,
            vb_w.pp("feature_projection"),
        )?;
        let encoder = Encoder::new(&cfg, vb_w.pp("encoder"))?;
        let lm_head = linear(cfg.embed_dim, cfg.vocab_size, vb.pp("lm_head"))
            .map_err(|e| CsmError::Config(format!("lm_head: {e}")))?;
        Ok(Self {
            cfg,
            feature_extractor,
            feature_projection,
            encoder,
            lm_head,
            device,
        })
    }

    /// Load from `model.safetensors` via mmap.
    pub fn load_from_safetensors<P: AsRef<std::path::Path>>(
        path: P,
        device: &Device,
    ) -> CsmResult<Self> {
        let vb = unsafe {
            candle_nn::VarBuilder::from_mmaped_safetensors(
                &[path.as_ref()],
                candle_core::DType::F32,
                device,
            )
        }
        .map_err(|e| CsmError::Config(format!("wav2vec2 safetensors load: {e}")))?;
        Self::new(Wav2Vec2Config::base_960h(), vb, device.clone())
    }

    /// Forward pass: 16 kHz audio (B, 1, T) → CTC logits (B, T', 32).
    /// Caller applies `log_softmax` on the last dim for proper CTC math.
    pub fn forward(&self, audio_16k: &Tensor) -> candle_core::Result<Tensor> {
        let feats = self.feature_extractor.forward(audio_16k)?;
        // (B, 512, T') → (B, T', 512)
        let feats = feats.transpose(1, 2)?.contiguous()?;
        let h = self.feature_projection.forward(&feats)?;
        let h = self.encoder.forward(&h)?;
        self.lm_head.forward(&h)
    }

    pub fn config(&self) -> &Wav2Vec2Config {
        &self.cfg
    }

    pub fn device(&self) -> &Device {
        &self.device
    }
}

/// Default CTC vocab for `wav2vec2-base-960h` per the upstream `vocab.json`.
/// Used by greedy + Viterbi decoders when the user doesn't supply a custom
/// tokenizer. Index = token id; value = char (`'|'` is the word separator).
pub const VOCAB_960H: &[&str] = &[
    "<pad>", "<s>", "</s>", "<unk>", "|", "E", "T", "A", "O", "N", "I", "H", "S", "R", "D", "L",
    "U", "M", "W", "C", "F", "G", "Y", "P", "B", "V", "K", "'", "X", "J", "Q", "Z",
];

/// CTC blank token id (`<pad>` = 0 by upstream convention).
pub const CTC_BLANK_ID: usize = 0;

/// Greedy CTC decode: per frame argmax → collapse repeats → drop blanks.
/// Returns the decoded transcript as a String. `vocab` indexes are looked
/// up from [`VOCAB_960H`] by default; supply a custom slice for non-960h
/// checkpoints.
pub fn ctc_greedy_decode(logits: &Tensor, vocab: &[&str]) -> CsmResult<String> {
    // logits: (1, T', V) or (T', V); reduce to (T', V).
    let logits = if logits.dims().len() == 3 {
        logits
            .i((0, .., ..))
            .map_err(|e| CsmError::Config(format!("ctc decode squeeze: {e}")))?
    } else {
        logits.clone()
    };
    let argmax: Vec<u32> = logits
        .argmax(candle_core::D::Minus1)
        .and_then(|t| t.to_vec1::<u32>())
        .map_err(|e| CsmError::Config(format!("ctc argmax: {e}")))?;
    // CTC collapse: skip if same as previous; drop blanks.
    let mut out = String::new();
    let mut prev: i64 = -1;
    for &id in argmax.iter() {
        let id_i = id as i64;
        if id_i == prev {
            continue;
        }
        if (id as usize) != CTC_BLANK_ID
            && let Some(tok) = vocab.get(id as usize)
        {
            if *tok == "|" {
                out.push(' ');
            } else if !tok.starts_with('<') {
                out.push_str(tok);
            }
        }
        prev = id_i;
    }
    Ok(out)
}

use candle_core::IndexOp;

/// One aligned token — output of [`viterbi_align`]. `token` is the raw
/// CTC unit (single character for the default `_base-960h` vocab).
/// `frame_start` / `frame_end` index 50 Hz feature frames; multiply by
/// 20 ms (1000 ms / 50 Hz) to convert to milliseconds.
#[derive(Debug, Clone, serde::Serialize)]
pub struct AlignedToken {
    pub token: String,
    pub frame_start: usize,
    pub frame_end: usize,
}

/// Group consecutive aligned tokens into words. Word-separator handling:
/// the upstream `_base-960h` vocab uses `|` as the inter-word boundary,
/// which ctc_greedy_decode renders as a space. After [`viterbi_align`]
/// the `|` shows up as its own token; this helper folds adjacent
/// non-separator tokens into a `Word { text, frame_start, frame_end }`.
#[derive(Debug, Clone, serde::Serialize)]
pub struct AlignedWord {
    pub word: String,
    pub frame_start: usize,
    pub frame_end: usize,
}

/// Convert a transcript string into the CTC token id sequence over the
/// supplied vocab. Spaces in the transcript become the word-separator
/// token (`|` for `_base-960h`); other characters are looked up directly.
/// Unknown characters are replaced with the `<unk>` token.
pub fn transcript_to_token_ids(transcript: &str, vocab: &[&str]) -> CsmResult<Vec<u32>> {
    let mut id_for: std::collections::HashMap<&str, u32> = Default::default();
    for (i, t) in vocab.iter().enumerate() {
        id_for.insert(*t, i as u32);
    }
    let unk_id = *id_for
        .get("<unk>")
        .ok_or_else(|| CsmError::Config("vocab missing <unk>".into()))?;
    let sep_id = *id_for
        .get("|")
        .ok_or_else(|| CsmError::Config("vocab missing | (word sep)".into()))?;
    let mut out = Vec::new();
    let upper = transcript.to_uppercase();
    let chars: Vec<char> = upper.chars().collect();
    let mut i = 0;
    while i < chars.len() {
        let c = chars[i];
        if c == ' ' {
            // Collapse runs of spaces into a single separator.
            out.push(sep_id);
            while i < chars.len() && chars[i] == ' ' {
                i += 1;
            }
            continue;
        }
        let s = c.to_string();
        let id = id_for.get(s.as_str()).copied().unwrap_or(unk_id);
        out.push(id);
        i += 1;
    }
    Ok(out)
}

/// CTC Viterbi forced alignment.
///
/// Given per-frame log-probabilities `log_probs` of shape `(T, V)` (caller
/// supplies the log-softmax) and a target token sequence `tokens`, find
/// the maximum-likelihood alignment of `tokens` against the frames using
/// the standard CTC transition lattice (interleave each token with a
/// blank, allow self-loop or advance-to-next at every step).
///
/// Returns one [`AlignedToken`] per token in the input sequence with
/// inclusive `frame_start..=frame_end` indices into the 50 Hz feature
/// frame grid.
///
/// Error if `T < tokens.len()` (not enough frames) or if `tokens` is empty.
pub fn viterbi_align(
    log_probs: &Tensor,
    tokens: &[u32],
    blank_id: usize,
    vocab: &[&str],
) -> CsmResult<Vec<AlignedToken>> {
    if tokens.is_empty() {
        return Err(CsmError::Config("viterbi_align: tokens is empty".into()));
    }
    let lp = if log_probs.dims().len() == 3 {
        log_probs
            .i((0, .., ..))
            .map_err(|e| CsmError::Config(format!("viterbi squeeze: {e}")))?
    } else {
        log_probs.clone()
    };
    let lp_vec = lp
        .to_vec2::<f32>()
        .map_err(|e| CsmError::Config(format!("viterbi to_vec2: {e}")))?;
    let big_t = lp_vec.len();
    let v = if big_t > 0 { lp_vec[0].len() } else { 0 };
    if big_t == 0 || v == 0 {
        return Err(CsmError::Config("viterbi_align: empty log_probs".into()));
    }

    // CTC trellis state sequence: [blank, t0, blank, t1, blank, ..., tN, blank]
    // length S = 2 * tokens.len() + 1
    let n = tokens.len();
    let big_s = 2 * n + 1;
    if big_t < n {
        return Err(CsmError::Config(format!(
            "viterbi_align: T={big_t} < tokens.len()={n}, can't align"
        )));
    }

    // Build the state token-id table.
    let mut state_id: Vec<usize> = Vec::with_capacity(big_s);
    for (i, _) in (0..big_s).enumerate() {
        if i % 2 == 0 {
            state_id.push(blank_id);
        } else {
            state_id.push(tokens[i / 2] as usize);
        }
    }

    let neg_inf = f32::NEG_INFINITY;
    // dp[t][s] = best log-prob for reaching state s at frame t
    // backptr[t][s] = state at t-1 we came from (for path recovery)
    let mut dp = vec![vec![neg_inf; big_s]; big_t];
    let mut backptr = vec![vec![0usize; big_s]; big_t];

    // Initialization: at t=0 we can be in state 0 (initial blank) or
    // state 1 (first token).
    dp[0][0] = lp_vec[0][state_id[0]];
    if big_s > 1 {
        dp[0][1] = lp_vec[0][state_id[1]];
    }

    for t in 1..big_t {
        for s in 0..big_s {
            // Stay in s, or advance from s-1, or skip from s-2 (only if
            // s is a non-blank state AND state s-2 is a different token —
            // CTC's "ε-skip" rule allows skipping blank between two
            // *different* tokens but NOT between two same tokens).
            let mut best = dp[t - 1][s];
            let mut best_prev = s;

            if s >= 1 && dp[t - 1][s - 1] > best {
                best = dp[t - 1][s - 1];
                best_prev = s - 1;
            }

            // ε-skip: only when current state is a non-blank token AND
            // the token at s-2 is different (so we can skip the
            // intermediate blank).
            if s >= 2 && s % 2 == 1 {
                let cur_tok = state_id[s];
                let prev_tok = state_id[s - 2];
                if cur_tok != prev_tok && dp[t - 1][s - 2] > best {
                    best = dp[t - 1][s - 2];
                    best_prev = s - 2;
                }
            }
            // Add the emission cost for being in state s at frame t.
            if best > neg_inf {
                dp[t][s] = best + lp_vec[t][state_id[s]];
                backptr[t][s] = best_prev;
            }
        }
    }

    // Final state: must end in either the last token (s = 2N-1) or the
    // trailing blank (s = 2N). Pick whichever has the higher log-prob.
    let last_token_state = big_s - 2;
    let last_blank_state = big_s - 1;
    let (mut s, _) = if dp[big_t - 1][last_token_state] >= dp[big_t - 1][last_blank_state] {
        (last_token_state, dp[big_t - 1][last_token_state])
    } else {
        (last_blank_state, dp[big_t - 1][last_blank_state])
    };

    // Recover the per-frame state path (reverse).
    let mut path = vec![0usize; big_t];
    path[big_t - 1] = s;
    for t in (1..big_t).rev() {
        s = backptr[t][s];
        path[t - 1] = s;
    }

    // Convert to per-token spans by finding the first and last frames
    // each non-blank state index (1, 3, 5, ..., 2N-1) appears in path.
    let mut out: Vec<AlignedToken> = Vec::with_capacity(n);
    for (i, &tok_id) in tokens.iter().enumerate() {
        let target_state = 2 * i + 1;
        let (mut start, mut end) = (None, None);
        for (t, &p) in path.iter().enumerate() {
            if p == target_state {
                if start.is_none() {
                    start = Some(t);
                }
                end = Some(t);
            }
        }
        // Fallback: if a token is degenerate / never visited, place it
        // at the path's nearest blank-bracket. This shouldn't happen
        // for legitimate inputs but keeps the function total.
        let (frame_start, frame_end) = match (start, end) {
            (Some(s0), Some(e0)) => (s0, e0),
            _ => (0, 0),
        };
        let token = vocab
            .get(tok_id as usize)
            .map(|s| s.to_string())
            .unwrap_or_else(|| format!("<id_{tok_id}>"));
        out.push(AlignedToken {
            token,
            frame_start,
            frame_end,
        });
    }
    Ok(out)
}

/// Group [`AlignedToken`]s into words at `|` boundaries.
pub fn group_into_words(tokens: &[AlignedToken]) -> Vec<AlignedWord> {
    let mut out = Vec::new();
    let mut buf = String::new();
    let mut start: Option<usize> = None;
    let mut end: usize = 0;
    for tok in tokens.iter() {
        if tok.token == "|" {
            if !buf.is_empty() {
                out.push(AlignedWord {
                    word: std::mem::take(&mut buf),
                    frame_start: start.unwrap_or(tok.frame_start),
                    frame_end: end,
                });
                start = None;
            }
            continue;
        }
        if start.is_none() {
            start = Some(tok.frame_start);
        }
        buf.push_str(&tok.token);
        end = tok.frame_end;
    }
    if !buf.is_empty() {
        out.push(AlignedWord {
            word: buf,
            frame_start: start.unwrap_or(0),
            frame_end: end,
        });
    }
    out
}

/// Convert a frame index (50 Hz, conv stride 320 at 16 kHz) to milliseconds.
pub fn frame_to_ms(frame: usize) -> f32 {
    frame as f32 * 20.0
}

#[cfg(test)]
mod tests {
    use super::*;
    use candle_core::DType;
    use candle_nn::VarMap;

    #[test]
    fn config_stride_product_matches_320() {
        let cfg = Wav2Vec2Config::base_960h();
        assert_eq!(cfg.conv_stride_total(), 320);
        assert_eq!(cfg.conv_layers.len(), 7);
        assert_eq!(cfg.num_layers, 12);
        assert_eq!(cfg.vocab_size, 32);
    }

    #[test]
    fn vocab_960h_size_matches_config() {
        assert_eq!(VOCAB_960H.len(), Wav2Vec2Config::base_960h().vocab_size);
        assert_eq!(VOCAB_960H[CTC_BLANK_ID], "<pad>");
        assert_eq!(VOCAB_960H[4], "|");
    }

    #[test]
    fn feature_extractor_random_init_shape() {
        let dev = Device::Cpu;
        let cfg = Wav2Vec2Config::base_960h();
        let vm = VarMap::new();
        let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
        let fe = FeatureExtractor::new(&cfg, vb).expect("build");
        let xs = Tensor::zeros((1, 1, 16_000), DType::F32, &dev).unwrap();
        let ys = <FeatureExtractor as Module>::forward(&fe, &xs).expect("forward");
        let dims = ys.dims();
        assert_eq!(dims.len(), 3);
        assert_eq!(dims[1], 512);
        // Same conv arithmetic as emotion2vec (identical conv stack).
        let ideal = 16_000 / cfg.conv_stride_total();
        assert!(dims[2] <= ideal && dims[2] >= ideal.saturating_sub(2));
    }

    #[test]
    fn conv_pos_embedding_preserves_t_with_even_kernel() {
        let dev = Device::Cpu;
        let vm = VarMap::new();
        let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
        // Kernel 128, groups 16 — exact `_base` config.
        let pos = ConvPosEmbedding::new(768, 128, 16, vb).expect("build");
        let xs = Tensor::zeros((1, 50, 768), DType::F32, &dev).unwrap();
        let ys = <ConvPosEmbedding as Module>::forward(&pos, &xs).expect("forward");
        // Same-padding via even-kernel-trim must keep T=50 exactly.
        assert_eq!(ys.dims(), &[1, 50, 768]);
    }

    #[test]
    fn block_post_norm_random_init_shape() {
        let dev = Device::Cpu;
        let cfg = Wav2Vec2Config::base_960h();
        let vm = VarMap::new();
        let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
        let blk =
            Block::new(cfg.embed_dim, cfg.num_heads, cfg.mlp_dim, cfg.norm_eps, vb).expect("build");
        let xs = Tensor::randn(0f32, 1.0, (1, 16, cfg.embed_dim), &dev).unwrap();
        let ys = <Block as Module>::forward(&blk, &xs).expect("forward");
        assert_eq!(ys.dims(), &[1, 16, cfg.embed_dim]);
        let v = ys.flatten_all().unwrap().to_vec1::<f32>().unwrap();
        assert!(v.iter().all(|x| x.is_finite()));
    }

    #[test]
    fn wav2vec2_random_init_end_to_end_shape() {
        let dev = Device::Cpu;
        let cfg = Wav2Vec2Config::base_960h();
        let vm = VarMap::new();
        let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
        let model = Wav2Vec2::new(cfg.clone(), vb, dev.clone()).expect("build");
        // 1 s of 16 kHz audio → expected output (1, ~50, 32).
        let audio = Tensor::randn(0f32, 0.1, (1, 1, 16_000), &dev).unwrap();
        let logits = model.forward(&audio).expect("forward");
        let dims = logits.dims();
        assert_eq!(dims.len(), 3);
        assert_eq!(dims[0], 1);
        assert_eq!(dims[2], cfg.vocab_size);
        assert!(dims[1] >= 48 && dims[1] <= 52);
        let v = logits.flatten_all().unwrap().to_vec1::<f32>().unwrap();
        assert!(v.iter().all(|x| x.is_finite()));
    }

    #[test]
    fn transcript_to_token_ids_handles_spaces_and_unknowns() {
        let ids = transcript_to_token_ids("HE LL", VOCAB_960H).unwrap();
        // H=11 E=5 |=4 L=15 L=15
        assert_eq!(ids, vec![11, 5, 4, 15, 15]);

        // Unknown char → <unk>=3
        let ids = transcript_to_token_ids("@", VOCAB_960H).unwrap();
        assert_eq!(ids, vec![3]);

        // Multiple spaces collapse to a single separator.
        let ids = transcript_to_token_ids("A   B", VOCAB_960H).unwrap();
        assert_eq!(ids, vec![7, 4, 24]);
    }

    #[test]
    fn viterbi_align_recovers_obvious_alignment() {
        let dev = Device::Cpu;
        let v = 32usize;
        // Synthesize log-probs so that the optimal alignment of "HE"
        // (tokens [11, 5]) is unambiguous. Frames: [pad pad H H E E pad].
        // We give a strong score to the intended state at each frame and
        // -10 elsewhere — small enough to keep blank fallback feasible.
        let layout: [u32; 7] = [0, 0, 11, 11, 5, 5, 0];
        let mut data = vec![-10.0f32; layout.len() * v];
        for (t, id) in layout.iter().enumerate() {
            data[t * v + *id as usize] = 0.0;
        }
        let lp = Tensor::from_vec(data, (layout.len(), v), &dev).unwrap();
        let aligned = viterbi_align(&lp, &[11, 5], CTC_BLANK_ID, VOCAB_960H).unwrap();
        assert_eq!(aligned.len(), 2);
        assert_eq!(aligned[0].token, "H");
        assert_eq!(aligned[0].frame_start, 2);
        assert_eq!(aligned[0].frame_end, 3);
        assert_eq!(aligned[1].token, "E");
        assert_eq!(aligned[1].frame_start, 4);
        assert_eq!(aligned[1].frame_end, 5);
    }

    #[test]
    fn group_into_words_splits_on_separator() {
        let toks = vec![
            AlignedToken {
                token: "H".into(),
                frame_start: 0,
                frame_end: 1,
            },
            AlignedToken {
                token: "I".into(),
                frame_start: 2,
                frame_end: 3,
            },
            AlignedToken {
                token: "|".into(),
                frame_start: 4,
                frame_end: 5,
            },
            AlignedToken {
                token: "Y".into(),
                frame_start: 6,
                frame_end: 6,
            },
            AlignedToken {
                token: "O".into(),
                frame_start: 7,
                frame_end: 8,
            },
        ];
        let words = group_into_words(&toks);
        assert_eq!(words.len(), 2);
        assert_eq!(words[0].word, "HI");
        assert_eq!(words[0].frame_start, 0);
        assert_eq!(words[0].frame_end, 3);
        assert_eq!(words[1].word, "YO");
        assert_eq!(words[1].frame_start, 6);
        assert_eq!(words[1].frame_end, 8);
    }

    #[test]
    fn frame_to_ms_50hz_grid() {
        // 50 Hz = 20 ms per frame.
        assert_eq!(frame_to_ms(0), 0.0);
        assert_eq!(frame_to_ms(50), 1000.0);
    }

    #[test]
    fn ctc_greedy_decode_collapses_repeats_and_drops_blanks() {
        let dev = Device::Cpu;
        // Synthesize logits such that argmax sequence is:
        // [pad pad H H E L L L O pad O pad pad |]
        // (pad=0, H=11, E=5, L=15, O=8, |=4)
        //
        // Standard CTC collapse:
        //   - blanks (pad) drop out
        //   - consecutive same-id with NO blank between → collapsed
        //   - same-id with a blank between → both emitted
        //
        // So:
        //   H H → H ; E ; L L L → L ; O ; (blank) ; O ; (blanks) ; | → space
        //   final: "HELOO "  (one H, one E, one L, two O's separated by
        //                     blank, trailing space from | word-separator)
        let argmax_seq: [u32; 14] = [0, 0, 11, 11, 5, 15, 15, 15, 8, 0, 8, 0, 0, 4];
        let v = 32usize;
        let mut data = vec![0.0f32; argmax_seq.len() * v];
        for (t, id) in argmax_seq.iter().enumerate() {
            data[t * v + *id as usize] = 1.0;
        }
        let logits = Tensor::from_vec(data, (1, argmax_seq.len(), v), &dev).unwrap();
        let text = ctc_greedy_decode(&logits, VOCAB_960H).unwrap();
        assert_eq!(text, "HELOO ");
    }
}