//! WavLM-Base+ Speaker Verification — Rust port (scaffold).
//!
//! Architectural port of `microsoft/wavlm-base-plus-sv` against candle 0.9.
//! WavLM-Base+ is 12-layer / 768-dim transformer over 16 kHz waveforms;
//! the `-sv` variant adds an X-vector head producing 512-d speaker
//! embeddings. Cosine similarity between embeddings → speaker similarity
//! score (drop-in replacement for the weak `SpectralCentroidSimilarity`
//! in `speaker_sim.rs`).
//!
//! ## Pipeline (input → embedding)
//!
//! ```text
//!   waveform [B, T] (16 kHz)
//!       │
//!       ▼
//!   FeatureExtractor (7 conv layers, 320× downsample) → [B, 512, T/320]
//!       │
//!   transpose ▼
//!   FeatureProjection: LN → Linear 512→768                → [B, T/320, 768]
//!       │
//!       ▼
//!   PosConv (Conv1d 768→768, k=128, groups=16, weight_norm) + GELU
//!       │
//!   add ▼ residual
//!   LayerNorm → 12 × WavLMEncoderLayer (with rel_pos_bias)  → 13 hidden states
//!       │
//!       ▼
//!   Weighted sum (softmax over 13 layer_weights)             → [B, T/320, 768]
//!       │
//!       ▼
//!   X-Vector head:
//!     - projector Linear 768→512
//!     - 5 TDNN layers with kernels [5,3,3,1,1] dilations [1,2,3,1,1]
//!     - ReLU after each
//!     - Statistics pooling [mean | std] over time          → [B, 3000]
//!     - feature_extractor Linear 3000→512                  → [B, 512]
//! ```
//!
//! ## Status: PARTIAL SCAFFOLD
//!
//! What this module ships:
//!   - All struct definitions for the full pipeline
//!   - `FeatureExtractor` (7-layer conv, GroupNorm at layer 0, GELU)
//!   - `FeatureProjection` (LayerNorm + Linear 512→768)
//!   - `XVectorHead` (5 TDNN + statistics pool + 3000→512 projection)
//!   - Stubs for `PosConv`, `WavLMEncoderLayer`, `Encoder`
//!   - Shape-correctness tests for the implemented blocks
//!
//! Deferred to Phase 5b:
//!   - Bucketed relative-position bias (T5-style, 320 buckets)
//!   - Gated relative-position bias (per-layer 1×12×1×1 const + 64→8 linear)
//!   - The full `WavLMEncoderLayer` forward (attention + FFN)
//!
//! Deferred to Phase 5c:
//!   - PyTorch `pytorch_model.bin` → safetensors conversion (handles weight_norm,
//!     `layer_weights` softmax-weights, TDNN `kernel.weight` reshape)
//!   - HF Hub asset resolution
//!
//! Deferred to Phase 5d:
//!   - Numerical parity vs the HF reference (cosine similarity within 1e-4 on
//!     a paired-utterance test set)

use crate::error::{CsmError, Result};
use candle_core::{D, DType, Device, IndexOp, Module, Tensor};
use candle_nn::{
    Conv1d, Conv1dConfig, GroupNorm, LayerNorm, LayerNormConfig, Linear, VarBuilder, conv1d,
    group_norm, layer_norm, linear, ops,
};
use std::path::Path;

pub const SAMPLE_RATE: u32 = 16_000;
pub const HIDDEN_DIM: usize = 768;
pub const NUM_HEADS: usize = 12;
pub const HEAD_DIM: usize = HIDDEN_DIM / NUM_HEADS; // 64
pub const FFN_DIM: usize = 3072;
pub const NUM_LAYERS: usize = 12;
pub const FEATURE_DIM: usize = 512;
pub const HOP_LENGTH: usize = 320; // total CNN downsampling
pub const REL_NUM_BUCKETS: usize = 320;
pub const REL_MAX_DISTANCE: usize = 800;
pub const EMBEDDING_DIM: usize = 512;
pub const STAT_POOL_DIM: usize = 3000; // 1500 * 2 (mean+std)

// -- 1. Feature extractor (7-layer Conv1d, 320× downsampling) --------------

#[derive(Debug, Clone)]
struct ConvLayer {
    conv: Conv1d,
    norm: Option<GroupNorm>,
    #[allow(dead_code)]
    kernel: usize,
    #[allow(dead_code)]
    stride: usize,
}

impl ConvLayer {
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        let h = xs.apply(&self.conv)?;
        let h = match self.norm.as_ref() {
            Some(gn) => h.apply(gn)?,
            None => h,
        };
        h.gelu_erf()
    }
}

#[derive(Debug, Clone)]
pub struct FeatureExtractor {
    layers: Vec<ConvLayer>,
}

impl FeatureExtractor {
    /// vb is rooted at `wavlm.feature_extractor`. Each layer's path is
    /// `conv_layers.{i}.conv.weight`. Layer 0 also has
    /// `conv_layers.0.layer_norm.{weight,bias}` (WavLM uses GroupNorm with
    /// num_groups = num_channels = 512, despite the parameter name).
    pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
        // (out_channels, kernel, stride). Input is mono so first in=1.
        let specs: [(usize, usize, usize); 7] = [
            (FEATURE_DIM, 10, 5),
            (FEATURE_DIM, 3, 2),
            (FEATURE_DIM, 3, 2),
            (FEATURE_DIM, 3, 2),
            (FEATURE_DIM, 3, 2),
            (FEATURE_DIM, 2, 2),
            (FEATURE_DIM, 2, 2),
        ];
        let mut layers = Vec::with_capacity(7);
        let mut in_ch = 1usize;
        let vb_layers = vb.pp("conv_layers");
        for (i, (out_ch, k, s)) in specs.iter().enumerate() {
            let cfg = Conv1dConfig {
                stride: *s,
                ..Default::default()
            };
            let conv = candle_nn::conv1d_no_bias(
                in_ch,
                *out_ch,
                *k,
                cfg,
                vb_layers.pp(i.to_string()).pp("conv"),
            )?;
            let norm = if i == 0 {
                // WavLM/wav2vec2 GroupNorm layer 0: num_groups = num_channels.
                Some(group_norm(
                    *out_ch,
                    *out_ch,
                    1e-5,
                    vb_layers.pp(i.to_string()).pp("layer_norm"),
                )?)
            } else {
                None
            };
            layers.push(ConvLayer {
                conv,
                norm,
                kernel: *k,
                stride: *s,
            });
            in_ch = *out_ch;
        }
        Ok(Self { layers })
    }
}

impl Module for FeatureExtractor {
    /// Input `(B, 1, T)`; output `(B, FEATURE_DIM=512, T/320)`.
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        let mut h = xs.clone();
        for layer in &self.layers {
            h = layer.forward(&h)?;
        }
        Ok(h)
    }
}

// -- 2. Feature projection (LN + Linear 512 → 768) ------------------------

#[derive(Debug, Clone)]
pub struct FeatureProjection {
    layer_norm: LayerNorm,
    projection: Linear,
}

impl FeatureProjection {
    /// vb is rooted at `wavlm.feature_projection`.
    pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
        let layer_norm = layer_norm(
            FEATURE_DIM,
            LayerNormConfig {
                eps: 1e-5,
                ..Default::default()
            },
            vb.pp("layer_norm"),
        )?;
        let projection = linear(FEATURE_DIM, HIDDEN_DIM, vb.pp("projection"))?;
        Ok(Self {
            layer_norm,
            projection,
        })
    }
}

impl Module for FeatureProjection {
    /// Input `(B, T/320, 512)`; output `(B, T/320, 768)`.
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        xs.apply(&self.layer_norm)?.apply(&self.projection)
    }
}

// -- 3. Positional convolution (depthwise Conv1d with weight_norm) --------

/// `Conv1d(768, 768, kernel=128, padding=64, groups=16)` with `weight_norm`.
/// Output is GELU-activated and added to the input as a residual.
///
/// Note: weight_norm storage requires a converter pass (g/v → weight),
/// matching the AudioSeal converter pattern. Currently builds a plain
/// Conv1d; safetensors loader must merge before reading.
#[derive(Debug, Clone)]
pub struct PosConv {
    conv: Conv1d,
    kernel: usize,
}

impl PosConv {
    pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
        let cfg = Conv1dConfig {
            padding: 64,
            groups: 16,
            ..Default::default()
        };
        let conv = conv1d(
            HIDDEN_DIM,
            HIDDEN_DIM,
            128,
            cfg,
            vb.pp("pos_conv_embed.conv"),
        )?;
        Ok(Self { conv, kernel: 128 })
    }
}

impl Module for PosConv {
    /// Input/output: `(B, T/320, 768)` (treats time axis -2 as conv axis).
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        // (B, T, C) → (B, C, T) → conv → (B, C, T) → drop trailing pad → (B, T, C)
        let in_len = xs.dim(D::Minus2)?;
        let xs_bct = xs.transpose(1, 2)?.contiguous()?;
        let h = xs_bct.apply(&self.conv)?;
        // SamePad strips 1 trailing frame because kernel=128 is even.
        let out_len = h.dim(D::Minus1)?;
        let h = if self.kernel.is_multiple_of(2) && out_len > in_len {
            h.narrow(D::Minus1, 0, in_len)?
        } else {
            h
        };
        let h = h.gelu_erf()?;
        // Back to (B, T, C).
        h.transpose(1, 2)?.contiguous()
    }
}

// -- 4. Transformer encoder layer with gated relative-position attention --

/// T5-style relative-position bucket. Maps a relative offset (k - q) into
/// `[0, num_buckets)`. Half the buckets cover negative offsets, half cover
/// positive; within each half, the first `max_exact = num_buckets/4` are
/// linear, the rest log-spaced up to `max_distance`.
fn relative_position_bucket(rel_pos: i64, num_buckets: usize, max_distance: usize) -> u32 {
    let half = num_buckets / 2;
    let mut bucket = if rel_pos > 0 { half } else { 0 };
    let abs_pos = rel_pos.unsigned_abs() as usize;
    let max_exact = half / 2;
    if abs_pos < max_exact {
        bucket += abs_pos;
    } else {
        let log_ratio = (abs_pos as f64 / max_exact as f64).ln();
        let log_factor = (max_distance as f64 / max_exact as f64).ln();
        let log_bucket = (log_ratio / log_factor) * (half - max_exact) as f64;
        let large = max_exact + log_bucket as usize;
        bucket += large.min(half - 1);
    }
    bucket as u32
}

/// One WavLM encoder layer. Post-norm: residual+attn → LN → FFN-residual → final-LN.
#[derive(Debug, Clone)]
pub struct WavLmEncoderLayer {
    // Phase 5b — these will be populated:
    pub q_proj: Linear,
    pub k_proj: Linear,
    pub v_proj: Linear,
    pub out_proj: Linear,
    pub attn_norm: LayerNorm,
    pub fc1: Linear,
    pub fc2: Linear,
    pub final_norm: LayerNorm,
    /// Per-layer gated rel-pos bias parameters (1, 12, 1, 1).
    pub gru_rel_pos_const: Tensor,
    /// Linear from head_dim=64 → 8 for the 2-gate split.
    pub gru_rel_pos_linear: Linear,
    /// Only set for layer 0 — bucketed embedding (320, 12).
    pub rel_attn_embed: Option<candle_nn::Embedding>,
}

impl WavLmEncoderLayer {
    pub fn new(layer_idx: usize, vb: VarBuilder) -> candle_core::Result<Self> {
        let attn = vb.pp("attention");
        let q_proj = linear(HIDDEN_DIM, HIDDEN_DIM, attn.pp("q_proj"))?;
        let k_proj = linear(HIDDEN_DIM, HIDDEN_DIM, attn.pp("k_proj"))?;
        let v_proj = linear(HIDDEN_DIM, HIDDEN_DIM, attn.pp("v_proj"))?;
        let out_proj = linear(HIDDEN_DIM, HIDDEN_DIM, attn.pp("out_proj"))?;
        let attn_norm = layer_norm(HIDDEN_DIM, 1e-5, vb.pp("layer_norm"))?;
        let ff = vb.pp("feed_forward");
        let fc1 = linear(HIDDEN_DIM, FFN_DIM, ff.pp("intermediate_dense"))?;
        let fc2 = linear(FFN_DIM, HIDDEN_DIM, ff.pp("output_dense"))?;
        let final_norm = layer_norm(HIDDEN_DIM, 1e-5, vb.pp("final_layer_norm"))?;
        let gru_rel_pos_const = attn.get((1, NUM_HEADS, 1, 1), "gru_rel_pos_const")?;
        let gru_rel_pos_linear = linear(HEAD_DIM, 8, attn.pp("gru_rel_pos_linear"))?;
        let rel_attn_embed = if layer_idx == 0 {
            Some(candle_nn::embedding(
                REL_NUM_BUCKETS,
                NUM_HEADS,
                attn.pp("rel_attn_embed"),
            )?)
        } else {
            None
        };
        Ok(Self {
            q_proj,
            k_proj,
            v_proj,
            out_proj,
            attn_norm,
            fc1,
            fc2,
            final_norm,
            gru_rel_pos_const,
            gru_rel_pos_linear,
            rel_attn_embed,
        })
    }

    /// Compute the bucketed relative-position bias `(num_heads, T, T)` for
    /// a given query length. Only callable on layer 0 (the layer that owns
    /// `rel_attn_embed`). Result is reused by the rest of the stack.
    pub fn compute_position_bias(&self, t: usize) -> candle_core::Result<Tensor> {
        let embed = self
            .rel_attn_embed
            .as_ref()
            .ok_or_else(|| candle_core::Error::Msg(
                "compute_position_bias called on layer without rel_attn_embed (only layer 0 has one)".into(),
            ))?;
        // Build (T, T) bucket index tensor on the host.
        let mut buckets = Vec::with_capacity(t * t);
        for q in 0..t {
            for k in 0..t {
                let rel = k as i64 - q as i64;
                buckets.push(relative_position_bucket(
                    rel,
                    REL_NUM_BUCKETS,
                    REL_MAX_DISTANCE,
                ));
            }
        }
        let device = embed.embeddings().device();
        let idx = Tensor::from_vec(buckets, (t, t), device)?;
        // Embedding lookup: (T, T) → (T, T, num_heads). Permute to
        // (num_heads, T, T) to match the HF compute_bias output.
        let values = embed.forward(&idx)?;
        values.permute((2, 0, 1))?.contiguous()
    }

    /// Compute the gated bias to add to attention scores.
    /// `position_bias`: `(num_heads, T, T)`; `xs`: `(B, T, embed_dim)`.
    /// Output: `(B, num_heads, T, T)` — the per-batch, per-head gated bias.
    fn gated_position_bias(
        &self,
        position_bias: &Tensor,
        xs: &Tensor,
    ) -> candle_core::Result<Tensor> {
        let (b, t, _) = xs.dims3()?;
        // (B, T, embed_dim) → (B, T, num_heads, head_dim) → (B, num_heads, T, head_dim)
        let h = xs
            .reshape((b, t, NUM_HEADS, HEAD_DIM))?
            .permute((0, 2, 1, 3))?
            .contiguous()?;
        // Linear(head_dim → 8) → (B, num_heads, T, 8)
        let proj = h.apply(&self.gru_rel_pos_linear)?;
        // (B, num_heads, T, 2, 4) → sum(-1) → (B, num_heads, T, 2)
        let proj = proj.reshape((b, NUM_HEADS, t, 2, 4))?.sum(D::Minus1)?;
        let gates = ops::sigmoid(&proj)?; // (B, num_heads, T, 2)
        let gate_a = gates.narrow(D::Minus1, 0, 1)?; // (B, num_heads, T, 1)
        let gate_b = gates.narrow(D::Minus1, 1, 1)?; // (B, num_heads, T, 1)
        // gate_output = gate_a * (gate_b * const - 1) + 2
        let const_g = self
            .gru_rel_pos_const
            .broadcast_as((b, NUM_HEADS, t, 1))?
            .to_dtype(gate_b.dtype())?;
        let inner = ((gate_b * const_g)? - 1.0f64)?;
        let gate_out = ((gate_a * inner)? + 2.0f64)?; // (B, num_heads, T, 1)
        // Broadcast position_bias (num_heads, T, T) → (1, num_heads, T, T) → (B, num_heads, T, T)
        let bias = position_bias
            .unsqueeze(0)?
            .broadcast_as((b, NUM_HEADS, t, t))?
            .to_dtype(gate_out.dtype())?;
        gate_out.broadcast_mul(&bias)
    }

    /// Multi-head self-attention with the gated bias added to the score
    /// matrix before softmax.
    fn attention(&self, xs: &Tensor, gated_bias: &Tensor) -> candle_core::Result<Tensor> {
        let (b, t, _) = xs.dims3()?;
        let split_heads = |proj: Tensor| -> candle_core::Result<Tensor> {
            proj.reshape((b, t, NUM_HEADS, HEAD_DIM))?
                .permute((0, 2, 1, 3))?
                .contiguous()
        };
        let q = split_heads(xs.apply(&self.q_proj)?)?;
        let k = split_heads(xs.apply(&self.k_proj)?)?;
        let v = split_heads(xs.apply(&self.v_proj)?)?;
        let scale = (HEAD_DIM as f64).powf(-0.5);
        let scores = (q.matmul(&k.transpose(2, 3)?.contiguous()?)? * scale)?;
        let scores = (scores + gated_bias)?;
        let attn = ops::softmax_last_dim(&scores)?;
        let out = attn.matmul(&v)?; // (B, num_heads, T, head_dim)
        let out = out
            .permute((0, 2, 1, 3))?
            .contiguous()?
            .reshape((b, t, HIDDEN_DIM))?;
        out.apply(&self.out_proj)
    }

    /// Forward pass. Returns `(hidden_states, position_bias)` so the
    /// encoder can thread the bias through subsequent layers.
    ///
    /// `in_bias` is `Some` for layers 1..N — the bias computed by layer 0.
    /// `None` is acceptable on layer 0 (we'll compute it locally) and an
    /// error on any other layer (caller's bug).
    pub fn forward_with_bias(
        &self,
        xs: &Tensor,
        in_bias: Option<&Tensor>,
    ) -> candle_core::Result<(Tensor, Tensor)> {
        let owned;
        let bias = match in_bias {
            Some(b) => b,
            None => {
                let t = xs.dim(D::Minus2)?;
                owned = self.compute_position_bias(t)?;
                &owned
            }
        };
        let gated = self.gated_position_bias(bias, xs)?;
        let attn_out = self.attention(xs, &gated)?;
        let h = (xs + attn_out)?;
        let h = h.apply(&self.attn_norm)?;
        // FFN: 768 → 3072 → GELU → 768
        let ffn = h.apply(&self.fc1)?.gelu_erf()?.apply(&self.fc2)?;
        let h = (h + ffn)?;
        let h = h.apply(&self.final_norm)?;
        Ok((h, bias.clone()))
    }

    /// Convenience used by the smoke test where the caller doesn't carry a
    /// bias — equivalent to layer 0's `forward_with_bias(xs, None)`.
    pub fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        let (out, _bias) = self.forward_with_bias(xs, None)?;
        Ok(out)
    }
}

// -- 5. Encoder (12 stacked layers, returns 13 hidden states) -------------

#[derive(Debug, Clone)]
pub struct Encoder {
    pos_conv: PosConv,
    layer_norm: LayerNorm,
    layers: Vec<WavLmEncoderLayer>,
}

impl Encoder {
    pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
        let pos_conv = PosConv::new(vb.clone())?;
        let layer_norm = layer_norm(HIDDEN_DIM, 1e-5, vb.pp("layer_norm"))?;
        let mut layers = Vec::with_capacity(NUM_LAYERS);
        let vb_layers = vb.pp("layers");
        for i in 0..NUM_LAYERS {
            layers.push(WavLmEncoderLayer::new(i, vb_layers.pp(i.to_string()))?);
        }
        Ok(Self {
            pos_conv,
            layer_norm,
            layers,
        })
    }

    /// Input `(B, T, 768)`; output is `(B, T, 768)` per layer + initial,
    /// returned as a `Vec<Tensor>` of length `NUM_LAYERS + 1 = 13`.
    /// The first hidden state is the post-norm input embedding (matches
    /// the HF `output_hidden_states` convention used by `WavLMForXVector`).
    pub fn forward_all_layers(&self, xs: &Tensor) -> candle_core::Result<Vec<Tensor>> {
        let pos = self.pos_conv.forward(xs)?;
        let mut h = (xs + &pos)?;
        h = h.apply(&self.layer_norm)?;
        let mut hidden_states = Vec::with_capacity(NUM_LAYERS + 1);
        hidden_states.push(h.clone());
        // Layer 0 computes the position_bias (shared across all 12 layers).
        let mut position_bias: Option<Tensor> = None;
        for layer in &self.layers {
            let (next_h, bias) = layer.forward_with_bias(&h, position_bias.as_ref())?;
            h = next_h;
            position_bias = Some(bias);
            hidden_states.push(h.clone());
        }
        Ok(hidden_states)
    }
}

// -- 6. X-Vector head (TDNN + stats pool + projection) --------------------

/// One TDNN layer: implemented per the HF reference as `Linear(in*kernel, out)`
/// fed via a manually-strided unfold over the time axis with `dilation`.
#[derive(Debug, Clone)]
pub struct Tdnn {
    kernel_linear: Linear,
    #[allow(dead_code)]
    in_dim: usize,
    #[allow(dead_code)]
    out_dim: usize,
    kernel: usize,
    dilation: usize,
}

impl Tdnn {
    pub fn new(
        in_dim: usize,
        out_dim: usize,
        kernel: usize,
        dilation: usize,
        vb: VarBuilder,
    ) -> candle_core::Result<Self> {
        let kernel_linear = linear(in_dim * kernel, out_dim, vb.pp("kernel"))?;
        Ok(Self {
            kernel_linear,
            in_dim,
            out_dim,
            kernel,
            dilation,
        })
    }
}

impl Module for Tdnn {
    /// Input `(B, T, in_dim)`; output `(B, T_out, out_dim)` where
    /// `T_out = T - (kernel - 1) * dilation` (no padding, valid-only).
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        let (b, t, c) = xs.dims3()?;
        let span = (self.kernel - 1) * self.dilation;
        if t <= span {
            return Err(candle_core::Error::Msg(format!(
                "TDNN input length {t} too short for kernel {} dilation {}",
                self.kernel, self.dilation
            )));
        }
        let t_out = t - span;
        // Build a (B, T_out, kernel*in_dim) tensor by gathering kernel
        // dilated samples per output frame.
        let mut frames: Vec<Tensor> = Vec::with_capacity(self.kernel);
        for k in 0..self.kernel {
            let offset = k * self.dilation;
            // (B, T_out, in_dim)
            let slice = xs.narrow(1, offset, t_out)?;
            frames.push(slice);
        }
        let stacked = Tensor::cat(&frames, 2)?; // (B, T_out, kernel * in_dim)
        debug_assert_eq!(stacked.dims(), &[b, t_out, self.kernel * c]);
        // Linear → (B, T_out, out_dim) → ReLU
        stacked.apply(&self.kernel_linear)?.relu()
    }
}

#[derive(Debug, Clone)]
pub struct XVectorHead {
    /// Softmax over `layer_weights` of length `NUM_LAYERS + 1 = 13`.
    layer_weights: Tensor,
    projector: Linear,
    tdnn: Vec<Tdnn>,
    /// Linear 3000 → 512 — the "feature_extractor" key in the HF state_dict.
    /// We rename to avoid collision with the WavLM CNN feature extractor.
    embedding_proj: Linear,
}

impl XVectorHead {
    /// vb is rooted at the *top* of the WavLMForXVector state_dict (so we
    /// read `layer_weights`, `projector`, `tdnn.{0..4}`, `feature_extractor`).
    pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
        let layer_weights = vb.get(NUM_LAYERS + 1, "layer_weights")?;
        let projector = linear(HIDDEN_DIM, FEATURE_DIM, vb.pp("projector"))?;
        // TDNN specs: (in, out, kernel, dilation)
        let tdnn_specs: [(usize, usize, usize, usize); 5] = [
            (FEATURE_DIM, 512, 5, 1),
            (512, 512, 3, 2),
            (512, 512, 3, 3),
            (512, 512, 1, 1),
            (512, 1500, 1, 1),
        ];
        let mut tdnn = Vec::with_capacity(5);
        let vb_tdnn = vb.pp("tdnn");
        for (i, (in_dim, out_dim, k, d)) in tdnn_specs.iter().enumerate() {
            tdnn.push(Tdnn::new(
                *in_dim,
                *out_dim,
                *k,
                *d,
                vb_tdnn.pp(i.to_string()),
            )?);
        }
        let embedding_proj = linear(STAT_POOL_DIM, EMBEDDING_DIM, vb.pp("feature_extractor"))?;
        Ok(Self {
            layer_weights,
            projector,
            tdnn,
            embedding_proj,
        })
    }

    /// Input: `Vec<Tensor>` of length 13, each `(B, T, 768)`. Output: `(B, 512)`.
    pub fn forward(&self, hidden_states: &[Tensor]) -> candle_core::Result<Tensor> {
        if hidden_states.len() != NUM_LAYERS + 1 {
            return Err(candle_core::Error::Msg(format!(
                "expected {} hidden states, got {}",
                NUM_LAYERS + 1,
                hidden_states.len()
            )));
        }
        // Softmax-weighted sum over layers.
        let weights = ops::softmax(&self.layer_weights, 0)?; // (13,)
        let weights = weights.to_dtype(hidden_states[0].dtype())?;
        let mut sum: Option<Tensor> = None;
        for (i, h) in hidden_states.iter().enumerate() {
            let w = weights.i(i)?; // scalar
            let scaled = h.broadcast_mul(&w.reshape((1, 1, 1))?)?;
            sum = Some(match sum {
                Some(prev) => (prev + scaled)?,
                None => scaled,
            });
        }
        let h = sum.expect("at least one hidden state");

        // Projector 768 → 512.
        let mut h = h.apply(&self.projector)?;
        // 5 TDNN layers (reduce time dim each time).
        for layer in &self.tdnn {
            h = layer.forward(&h)?;
        }
        // Statistics pooling: mean & std over time axis (dim=1).
        let mean = h.mean(1)?; // (B, 1500)
        let var = h.var(1)?;
        let std = (var + 1e-9)?.sqrt()?;
        let stat = Tensor::cat(&[&mean, &std], D::Minus1)?; // (B, 3000)
        // Final projection 3000 → 512.
        stat.apply(&self.embedding_proj)
    }
}

// -- 7. Top-level model ----------------------------------------------------

#[derive(Debug, Clone)]
pub struct WavLmSv {
    feature_extractor: FeatureExtractor,
    feature_projection: FeatureProjection,
    encoder: Encoder,
    head: XVectorHead,
}

impl WavLmSv {
    pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
        let backbone = vb.pp("wavlm");
        let feature_extractor = FeatureExtractor::new(backbone.pp("feature_extractor"))?;
        let feature_projection = FeatureProjection::new(backbone.pp("feature_projection"))?;
        let encoder = Encoder::new(backbone.pp("encoder"))?;
        let head = XVectorHead::new(vb.clone())?;
        Ok(Self {
            feature_extractor,
            feature_projection,
            encoder,
            head,
        })
    }

    /// Per-utterance zero-mean/unit-variance normalization (matches
    /// HF Wav2Vec2FeatureExtractor with do_normalize=True).
    pub fn normalize_waveform(samples: &[f32]) -> Vec<f32> {
        if samples.is_empty() {
            return Vec::new();
        }
        let n = samples.len() as f32;
        let mean = samples.iter().sum::<f32>() / n;
        let var = samples.iter().map(|s| (s - mean).powi(2)).sum::<f32>() / n;
        let std = (var + 1e-7).sqrt();
        samples.iter().map(|s| (s - mean) / std).collect()
    }

    /// Compute a `(B, 512)` speaker embedding from a 16 kHz mono waveform
    /// tensor `(B, 1, T)`. Caller is responsible for normalization.
    pub fn embed(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        let features = self.feature_extractor.forward(xs)?; // (B, 512, T/320)
        // (B, 512, T') → (B, T', 512)
        let features = features.transpose(1, 2)?.contiguous()?;
        let projected = self.feature_projection.forward(&features)?; // (B, T', 768)
        let hidden_states = self.encoder.forward_all_layers(&projected)?; // 13 × (B, T', 768)
        self.head.forward(&hidden_states)
    }

    /// Convenience: embed a `Vec<f32>` and return the embedding as a Vec<f32>.
    pub fn embed_samples(&self, samples: &[f32], device: &Device) -> Result<Vec<f32>> {
        let normalized = Self::normalize_waveform(samples);
        let xs = Tensor::from_slice(&normalized, (1, 1, normalized.len()), device)
            .map_err(|e| CsmError::Config(format!("embed: tensor: {e}")))?;
        let emb = self
            .embed(&xs)
            .map_err(|e| CsmError::Config(format!("embed: forward: {e}")))?;
        emb.i(0)?
            .to_dtype(DType::F32)
            .and_then(|t| t.to_vec1::<f32>())
            .map_err(|e| CsmError::Config(format!("embed: to_vec: {e}")))
    }

    /// Cosine similarity between two embeddings, both expected as `(D,)`
    /// or `(B, D)` tensors of compatible shape.
    pub fn cosine_similarity(a: &Tensor, b: &Tensor) -> candle_core::Result<Tensor> {
        let a_norm = a.broadcast_div(&(a.sqr()?.sum_keepdim(D::Minus1)? + 1e-9)?.sqrt()?)?;
        let b_norm = b.broadcast_div(&(b.sqr()?.sum_keepdim(D::Minus1)? + 1e-9)?.sqrt()?)?;
        (a_norm * b_norm)?.sum(D::Minus1)
    }
}

/// Stub loader. Phase 5c will implement `pytorch_model.bin` → safetensors
/// conversion (weight_norm merge for pos_conv + TDNN kernel reshape) and
/// HF Hub asset resolution under `microsoft/wavlm-base-plus-sv`.
pub fn load_from_safetensors<P: AsRef<Path>>(safetensors: P, device: &Device) -> Result<WavLmSv> {
    let vb = unsafe {
        candle_nn::VarBuilder::from_mmaped_safetensors(&[safetensors.as_ref()], DType::F32, device)
    }
    .map_err(|e| CsmError::Config(format!("opening WavLM safetensors: {e}")))?;
    WavLmSv::new(vb).map_err(|e| CsmError::Config(format!("WavLmSv::new: {e}")))
}

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

    fn random_vb(device: &Device) -> (VarMap, VarBuilder<'static>) {
        let vm = VarMap::new();
        let vb = VarBuilder::from_varmap(&vm, DType::F32, device);
        (vm, vb)
    }

    #[test]
    fn feature_extractor_downsamples_by_320() {
        let device = Device::Cpu;
        let (_vm, vb) = random_vb(&device);
        let fe = FeatureExtractor::new(vb).unwrap();
        // 16000 samples = 1s @ 16 kHz → ~50 frames.
        let xs = Tensor::randn(0f32, 1f32, (1, 1, 16000), &device).unwrap();
        let out = fe.forward(&xs).unwrap();
        let frames = out.dim(D::Minus1).unwrap();
        assert_eq!(out.dim(0).unwrap(), 1);
        assert_eq!(out.dim(1).unwrap(), FEATURE_DIM);
        assert!(
            (frames as i64 - 49).abs() <= 2,
            "expected ~49 frames, got {frames}"
        );
    }

    #[test]
    fn feature_projection_lifts_512_to_768() {
        let device = Device::Cpu;
        let (_vm, vb) = random_vb(&device);
        let fp = FeatureProjection::new(vb).unwrap();
        let xs = Tensor::randn(0f32, 1f32, (2, 49, FEATURE_DIM), &device).unwrap();
        let out = fp.forward(&xs).unwrap();
        assert_eq!(out.dims(), &[2, 49, HIDDEN_DIM]);
    }

    #[test]
    fn pos_conv_preserves_shape() {
        let device = Device::Cpu;
        let (_vm, vb) = random_vb(&device);
        let pc = PosConv::new(vb).unwrap();
        let xs = Tensor::randn(0f32, 1f32, (1, 49, HIDDEN_DIM), &device).unwrap();
        let out = pc.forward(&xs).unwrap();
        assert_eq!(out.dims(), &[1, 49, HIDDEN_DIM]);
    }

    #[test]
    fn tdnn_reduces_time_axis() {
        let device = Device::Cpu;
        let (_vm, vb) = random_vb(&device);
        // kernel=5, dilation=1 → T_out = T - 4
        let t = Tdnn::new(FEATURE_DIM, 512, 5, 1, vb.clone()).unwrap();
        let xs = Tensor::randn(0f32, 1f32, (1, 49, FEATURE_DIM), &device).unwrap();
        let out = t.forward(&xs).unwrap();
        assert_eq!(out.dims(), &[1, 45, 512]);
    }

    #[test]
    fn tdnn_dilated_reduces_correctly() {
        let device = Device::Cpu;
        let (_vm, vb) = random_vb(&device);
        // kernel=3, dilation=2 → span = 4, T_out = T - 4
        let t = Tdnn::new(512, 512, 3, 2, vb).unwrap();
        let xs = Tensor::randn(0f32, 1f32, (1, 30, 512), &device).unwrap();
        let out = t.forward(&xs).unwrap();
        assert_eq!(out.dims(), &[1, 26, 512]);
    }

    #[test]
    fn xvector_head_produces_512d_embedding() {
        let device = Device::Cpu;
        let (_vm, vb) = random_vb(&device);
        let head = XVectorHead::new(vb).unwrap();
        // T must be at least 5+4+4 = 13 frames after the TDNN cascade.
        // Cascade reductions: 4, 4, 4, 0, 0 = 12 frames lost. Need T >= 13.
        let t = 49usize;
        let hidden_states: Vec<Tensor> = (0..NUM_LAYERS + 1)
            .map(|_| Tensor::randn(0f32, 1f32, (1, t, HIDDEN_DIM), &device).unwrap())
            .collect();
        let emb = head.forward(&hidden_states).unwrap();
        assert_eq!(emb.dims(), &[1, EMBEDDING_DIM]);
    }

    #[test]
    fn end_to_end_forward_smoke() {
        let device = Device::Cpu;
        let (_vm, vb) = random_vb(&device);
        let model = WavLmSv::new(vb).unwrap();
        // 1s @ 16 kHz = 16000 samples.
        let xs = Tensor::randn(0f32, 1f32, (1, 1, 16000), &device).unwrap();
        let emb = model.embed(&xs).unwrap();
        assert_eq!(emb.dims(), &[1, EMBEDDING_DIM]);
    }

    #[test]
    fn relative_position_bucket_basics() {
        // Distance 0 → bucket 0 (small, left side).
        assert_eq!(relative_position_bucket(0, 320, 800), 0);
        // Tiny positive → goes to right half (bucket >= 160).
        let b1 = relative_position_bucket(1, 320, 800);
        assert!(b1 >= 160 && b1 < 320);
        // Large positive → still on right half, capped to half - 1.
        let b_large = relative_position_bucket(10_000, 320, 800);
        assert_eq!(b_large, 320 - 1);
        // Negative on the left half.
        let b_neg = relative_position_bucket(-1, 320, 800);
        assert!(b_neg < 160);
    }

    #[test]
    fn encoder_layer_runs_with_layer0_bias() {
        let device = Device::Cpu;
        let (_vm, vb) = random_vb(&device);
        let layer = WavLmEncoderLayer::new(0, vb).unwrap();
        let xs = Tensor::randn(0f32, 1f32, (1, 24, HIDDEN_DIM), &device).unwrap();
        let (out, bias) = layer.forward_with_bias(&xs, None).unwrap();
        assert_eq!(out.dims(), &[1, 24, HIDDEN_DIM]);
        assert_eq!(bias.dims(), &[NUM_HEADS, 24, 24]);
        // Verify the layer actually transformed the input (not a no-op).
        let diff = (&out - &xs).unwrap();
        let l2 = diff.sqr().unwrap().sum_all().unwrap();
        let l2: f32 = l2.to_dtype(DType::F32).unwrap().to_scalar().unwrap();
        assert!(l2 > 1e-6, "encoder layer is a no-op (l2 = {l2})");
    }

    #[test]
    fn encoder_threads_position_bias_through_stack() {
        let device = Device::Cpu;
        let (_vm, vb) = random_vb(&device);
        let enc = Encoder::new(vb).unwrap();
        let xs = Tensor::randn(0f32, 1f32, (1, 24, HIDDEN_DIM), &device).unwrap();
        let states = enc.forward_all_layers(&xs).unwrap();
        assert_eq!(states.len(), NUM_LAYERS + 1);
        for s in &states {
            assert_eq!(s.dims(), &[1, 24, HIDDEN_DIM]);
        }
    }

    #[test]
    fn cosine_similarity_self_is_one() {
        let device = Device::Cpu;
        let v = Tensor::from_slice(&[1.0f32, 2.0, 3.0, 4.0], (4,), &device).unwrap();
        let s = WavLmSv::cosine_similarity(&v, &v).unwrap();
        let v: f32 = s.to_dtype(DType::F32).unwrap().to_scalar().unwrap();
        assert!((v - 1.0).abs() < 1e-5, "expected 1.0, got {v}");
    }

    #[test]
    fn normalize_waveform_zero_mean() {
        let s = WavLmSv::normalize_waveform(&[1.0, 2.0, 3.0, 4.0, 5.0]);
        let mean: f32 = s.iter().sum::<f32>() / s.len() as f32;
        assert!(mean.abs() < 1e-5);
        let var = s.iter().map(|x| x * x).sum::<f32>() / s.len() as f32;
        assert!(
            (var - 1.0).abs() < 1e-3,
            "expected unit variance, got {var}"
        );
    }
}