//! Phase 13.8 — emotion2vec_plus_base candle port.
//!
//! Multi-slice port of the FunASR/iic emotion2vec_plus_base model
//! (`emotion2vec/emotion2vec_plus_base` on HF Hub). Final goal: a real
//! [`crate::ser::EmotionDetector`] implementation that replaces the
//! prosody-rule placeholder used by `audio_to_manifest --auto-emotion-tag`
//! and `converse_server --reactive-emotion`.
//!
//! Architecture confirmed via `examples/emotion2vec_inspect`:
//!   audio 16 kHz (1, T)
//!    → local_encoder              7 × Conv1d (1→512 chans, T → T/320)
//!    → project_features           Linear 512 → 768
//!    → relative_positional_encoder 5 × Conv1d (768→48, kernel=19)
//!    → context_encoder            4-layer transformer prenet
//!    → blocks.0..7                8-layer main transformer
//!    → norm                       LayerNorm(768)
//!    → mean-pool over time
//!    → proj                       Linear 768 → 9
//!    → softmax → 9-class probability
//!
//! Slicing:
//! - **Slice 1 (Phase 13.8)**: inspector + design notes ✅
//! - **Slice 2 (this module, growing)**: candle module scaffolding
//! - **Slice 3**: forward pass + shape verification on synthetic audio
//! - **Slice 4**: `EmotionDetector` impl + integration
//!
//! Full design notes live in `docs/emotion2vec_port_notes.md`.

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

/// Top-level architectural hyperparameters from `config.yaml`. Only the
/// fields needed by the inference path are surfaced; training-time fields
/// (mask probabilities, EMA decay, drop-path schedules, etc.) are
/// deliberately omitted because they have no effect on `forward()`.
#[derive(Debug, Clone)]
pub struct Emotion2VecConfig {
    /// Hidden dim of every transformer block (768 for `_base`).
    pub embed_dim: usize,
    /// Heads in every transformer block (12 for `_base`).
    pub num_heads: usize,
    /// MLP expansion ratio (4.0 for `_base` → 768→3072→768).
    pub mlp_ratio: f32,
    /// Main encoder depth — `blocks.0..depth-1` (8 for `_base`).
    pub depth: usize,
    /// Prenet (`context_encoder`) depth (4 for `_base`).
    pub prenet_depth: usize,
    /// Number of output emotion classes (9 for `_base`):
    /// angry/disgusted/fearful/happy/neutral/other/sad/surprised/`<unk>`.
    pub num_classes: usize,
    /// LayerNorm epsilon (1e-5).
    pub norm_eps: f64,
    /// Output channels of the local conv encoder (512 for `_base`).
    pub feature_dim: usize,
    /// Conv stack spec: `(out_channels, kernel_size, stride)` per layer.
    /// For `_base`: `[(512,10,5), (512,3,2), (512,3,2), (512,3,2),
    /// (512,3,2), (512,2,2), (512,2,2)]` — total stride product = 320,
    /// so 16 kHz → 50 Hz feature rate (one frame per 20 ms).
    pub conv_layers: Vec<(usize, usize, usize)>,
}

impl Emotion2VecConfig {
    pub fn plus_base() -> Self {
        Self {
            embed_dim: 768,
            num_heads: 12,
            mlp_ratio: 4.0,
            depth: 8,
            prenet_depth: 4,
            num_classes: 9,
            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),
            ],
        }
    }

    /// Total stride product of the conv stack — the audio-rate divisor.
    /// 320 for `_base`: 16 kHz → 50 Hz feature frames.
    pub fn conv_stride_total(&self) -> usize {
        self.conv_layers.iter().map(|(_, _, s)| *s).product()
    }
}

/// One Conv1d → LayerNorm → GELU block in the local encoder.
///
/// In the upstream PyTorch, each block is
///   `Sequential(Conv1d, Dropout, Sequential(TransposeLast, LayerNorm,
///   TransposeLast), GELU)`
/// with state-dict keys `.0.weight` (Conv1d, no bias), `.2.1.weight/bias`
/// (the LayerNorm). The Dropout has no params and is inactive at
/// inference, so we skip it. The TransposeLast wrappers move the channel
/// axis to last for LayerNorm and back; in candle we call `.transpose(1,2)`
/// before/after the norm.
#[derive(Debug)]
pub struct LocalConvBlock {
    conv: Conv1d,
    norm: LayerNorm,
}

impl LocalConvBlock {
    /// Build from a VarBuilder rooted at e.g.
    /// `d2v_model.modality_encoders.AUDIO.local_encoder.conv_layers.{i}`.
    /// `vb.pp("0")` reaches the Conv1d weight; `vb.pp("2").pp("1")` reaches
    /// the LayerNorm weight/bias. Mirrors the upstream nesting exactly.
    pub fn new(
        in_channels: usize,
        out_channels: usize,
        kernel_size: usize,
        stride: usize,
        norm_eps: f64,
        vb: VarBuilder,
    ) -> CsmResult<Self> {
        let cfg = Conv1dConfig {
            stride,
            ..Default::default()
        };
        let conv = conv1d_no_bias(in_channels, out_channels, kernel_size, cfg, vb.pp("0"))
            .map_err(|e| crate::CsmError::Config(format!("local conv: {e}")))?;
        let norm = layer_norm(out_channels, norm_eps, vb.pp("2").pp("1"))
            .map_err(|e| crate::CsmError::Config(format!("local norm: {e}")))?;
        Ok(Self { conv, norm })
    }
}

impl Module for LocalConvBlock {
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        // (B, C_in, T_in)
        let xs = self.conv.forward(xs)?;
        // → (B, C_out, T_out). LayerNorm applies over the last dim, so
        //   move channels to the last axis, normalize, swap back.
        let xs = xs.transpose(1, 2)?.contiguous()?; // (B, T, C)
        let xs = self.norm.forward(&xs)?;
        let xs = xs.transpose(1, 2)?.contiguous()?; // (B, C, T)
        xs.gelu()
    }
}

/// The 7-block local conv feature extractor.
///
/// Input: 16 kHz mono PCM tensor of shape `(B, 1, T)`.
/// Output: `(B, feature_dim, T / conv_stride_total())` — 512 channels,
/// 50 Hz frame rate for `_base` (T → T/320).
#[derive(Debug)]
pub struct LocalEncoder {
    blocks: Vec<LocalConvBlock>,
}

impl LocalEncoder {
    /// Build from a VarBuilder rooted at
    /// `d2v_model.modality_encoders.AUDIO.local_encoder.conv_layers`.
    pub fn new(cfg: &Emotion2VecConfig, vb: VarBuilder) -> CsmResult<Self> {
        let mut blocks = Vec::with_capacity(cfg.conv_layers.len());
        let mut in_ch = 1usize; // raw audio is mono
        for (i, (out_ch, k, s)) in cfg.conv_layers.iter().enumerate() {
            let block =
                LocalConvBlock::new(in_ch, *out_ch, *k, *s, cfg.norm_eps, vb.pp(i.to_string()))?;
            blocks.push(block);
            in_ch = *out_ch;
        }
        Ok(Self { blocks })
    }
}

impl Module for LocalEncoder {
    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)
    }
}

/// Linear `feature_dim → embed_dim` (512 → 768 for `_base`) — projects the
/// LocalEncoder's per-frame output into the transformer's working
/// dimension. State-dict keys: `project_features.2.weight/bias` (path
/// includes a LayerNorm at index 1 that we surface separately).
///
/// Per the inspector output the upstream stores:
///   `project_features.1.weight/bias`  → LayerNorm(512)
///   `project_features.2.weight/bias`  → Linear(512 → 768)
///
/// So the projection module is `LayerNorm → Linear`, both at the feature
/// dim (channel-last after the LocalEncoder transposes back to (B, T, C)
/// for the transformer).
#[derive(Debug)]
pub struct ProjectFeatures {
    norm: LayerNorm,
    proj: Linear,
}

impl ProjectFeatures {
    /// Build from a VarBuilder rooted at
    /// `d2v_model.modality_encoders.AUDIO.project_features`.
    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("1"))
            .map_err(|e| crate::CsmError::Config(format!("project_features norm: {e}")))?;
        let proj = linear(in_dim, out_dim, vb.pp("2"))
            .map_err(|e| crate::CsmError::Config(format!("project_features linear: {e}")))?;
        Ok(Self { norm, proj })
    }
}

impl Module for ProjectFeatures {
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        // xs: (B, T, in_dim) — LocalEncoder output already transposed
        // to channel-last. Both ops are channel-last.
        let xs = self.norm.forward(xs)?;
        self.proj.forward(&xs)
    }
}

/// One transformer block — pre-norm, fused QKV attention, MLP residual.
///
/// State-dict layout (per `blocks.{i}` and `context_encoder.blocks.{i}`):
///   `norm1.weight/bias`           LayerNorm(embed_dim)
///   `attn.qkv.weight/bias`        Linear(embed_dim → 3 × embed_dim)
///   `attn.proj.weight/bias`       Linear(embed_dim → embed_dim)
///   `norm2.weight/bias`           LayerNorm(embed_dim)
///   `mlp.fc1.weight/bias`         Linear(embed_dim → mlp_dim)
///   `mlp.fc2.weight/bias`         Linear(mlp_dim → embed_dim)
///
/// MLP activation is GELU. No positional encoding is added inside the
/// block — the upstream applies a conv-based positional bias at the
/// encoder boundary (slice 2c).
#[derive(Debug)]
pub struct Block {
    norm1: LayerNorm,
    qkv: Linear,
    attn_proj: Linear,
    norm2: LayerNorm,
    fc1: Linear,
    fc2: Linear,
    num_heads: usize,
    head_dim: usize,
    scale: f64,
}

impl Block {
    /// Build from a VarBuilder rooted at e.g. `d2v_model.blocks.{i}` or
    /// `d2v_model.modality_encoders.AUDIO.context_encoder.blocks.{i}`.
    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 norm1 = layer_norm(embed_dim, norm_eps, vb.pp("norm1"))
            .map_err(|e| crate::CsmError::Config(format!("block norm1: {e}")))?;
        let qkv = linear(embed_dim, embed_dim * 3, vb.pp("attn").pp("qkv"))
            .map_err(|e| crate::CsmError::Config(format!("block qkv: {e}")))?;
        let attn_proj = linear(embed_dim, embed_dim, vb.pp("attn").pp("proj"))
            .map_err(|e| crate::CsmError::Config(format!("block attn proj: {e}")))?;
        let norm2 = layer_norm(embed_dim, norm_eps, vb.pp("norm2"))
            .map_err(|e| crate::CsmError::Config(format!("block norm2: {e}")))?;
        let fc1 = linear(embed_dim, mlp_dim, vb.pp("mlp").pp("fc1"))
            .map_err(|e| crate::CsmError::Config(format!("block fc1: {e}")))?;
        let fc2 = linear(mlp_dim, embed_dim, vb.pp("mlp").pp("fc2"))
            .map_err(|e| crate::CsmError::Config(format!("block fc2: {e}")))?;
        Ok(Self {
            norm1,
            qkv,
            attn_proj,
            norm2,
            fc1,
            fc2,
            num_heads,
            head_dim,
            scale,
        })
    }

    fn attention(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        // xs: (B, T, D)
        let (b, t, d) = xs.dims3()?;
        let qkv = self.qkv.forward(xs)?; // (B, T, 3D)
        // Reshape to (B, T, 3, H, head_dim), permute to (3, B, H, T, hd)
        let qkv = qkv
            .reshape((b, t, 3, self.num_heads, self.head_dim))?
            .permute((2, 0, 3, 1, 4))?
            .contiguous()?;
        // Split into Q, K, V along the leading 3-axis.
        let q = qkv.narrow(0, 0, 1)?.squeeze(0)?; // (B, H, T, hd)
        let k = qkv.narrow(0, 1, 1)?.squeeze(0)?;
        let v = qkv.narrow(0, 2, 1)?.squeeze(0)?;
        // Attention scores. For Metal compatibility (the 4D matmul issue
        // we hit on Moonshine, Phase 8.8), collapse (B, H) to a single
        // batch dim before matmul.
        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)?; // (B*H, T, hd)
        let out = out
            .reshape((b, self.num_heads, t, self.head_dim))?
            .transpose(1, 2)?
            .reshape((b, t, d))?
            .contiguous()?;
        self.attn_proj.forward(&out)
    }

    fn mlp(&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> {
        // Pre-norm: x = x + attn(norm1(x)); x = x + mlp(norm2(x))
        let h = self.norm1.forward(xs)?;
        let h = self.attention(&h)?;
        let xs = (xs + h)?;
        let h = self.norm2.forward(&xs)?;
        let h = self.mlp(&h)?;
        xs + h
    }
}

/// 4-layer transformer prenet sitting between `ProjectFeatures` (or
/// rather: features + positional bias) and the main 8-layer encoder.
/// Has its own final LayerNorm — pickle keys
/// `context_encoder.blocks.0..3.*` + `context_encoder.norm.weight/bias`.
#[derive(Debug)]
pub struct ContextEncoder {
    blocks: Vec<Block>,
    norm: LayerNorm,
}

impl ContextEncoder {
    /// Build from a VarBuilder rooted at
    /// `d2v_model.modality_encoders.AUDIO.context_encoder`.
    pub fn new(cfg: &Emotion2VecConfig, vb: VarBuilder) -> CsmResult<Self> {
        let mlp_dim = (cfg.embed_dim as f32 * cfg.mlp_ratio) as usize;
        let mut blocks = Vec::with_capacity(cfg.prenet_depth);
        let vb_blocks = vb.pp("blocks");
        for i in 0..cfg.prenet_depth {
            blocks.push(Block::new(
                cfg.embed_dim,
                cfg.num_heads,
                mlp_dim,
                cfg.norm_eps,
                vb_blocks.pp(i.to_string()),
            )?);
        }
        let norm = layer_norm(cfg.embed_dim, cfg.norm_eps, vb.pp("norm"))
            .map_err(|e| crate::CsmError::Config(format!("context_encoder norm: {e}")))?;
        Ok(Self { blocks, norm })
    }
}

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

/// 8-layer main transformer encoder (`d2v_model.blocks.0..7`). No final
/// LayerNorm — the upstream pickle has none, and the pre-norm pattern's
/// per-block `norm2` keeps residuals well-conditioned without one.
#[derive(Debug)]
pub struct MainEncoder {
    blocks: Vec<Block>,
}

impl MainEncoder {
    /// Build from a VarBuilder rooted at `d2v_model.blocks`.
    pub fn new(cfg: &Emotion2VecConfig, vb: VarBuilder) -> CsmResult<Self> {
        let mlp_dim = (cfg.embed_dim as f32 * cfg.mlp_ratio) as usize;
        let mut blocks = Vec::with_capacity(cfg.depth);
        for i in 0..cfg.depth {
            blocks.push(Block::new(
                cfg.embed_dim,
                cfg.num_heads,
                mlp_dim,
                cfg.norm_eps,
                vb.pp(i.to_string()),
            )?);
        }
        Ok(Self { blocks })
    }
}

impl Module for MainEncoder {
    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)
    }
}

/// Final 9-class emotion classifier — a single Linear layer applied to
/// the time-mean-pooled hidden state. No LayerNorm; pickle keys are just
/// `proj.weight` (`[9, 768]`) and `proj.bias` (`[9]`) at the TOP level
/// of the state dict (NOT under the `d2v_model.` prefix).
///
/// Output classes (per `tokens.txt`):
/// 0=angry 1=disgusted 2=fearful 3=happy 4=neutral 5=other 6=sad
/// 7=surprised 8=`<unk>`. Map to [`crate::ser::EmotionLabel`] via
/// [`Self::tag_for_class`].
#[derive(Debug)]
pub struct Classifier {
    proj: Linear,
}

impl Classifier {
    /// Build from a VarBuilder rooted at the *root* of the model state
    /// dict (NOT at `d2v_model.`) — `proj.weight/bias` live there.
    pub fn new(embed_dim: usize, num_classes: usize, vb: VarBuilder) -> CsmResult<Self> {
        let proj = linear(embed_dim, num_classes, vb.pp("proj"))
            .map_err(|e| crate::CsmError::Config(format!("classifier proj: {e}")))?;
        Ok(Self { proj })
    }

    /// Map a 0..=8 emotion2vec class index directly to one of the 9 raw
    /// [`crate::ser::EmotionLabel`] variants — NO fold.
    ///
    /// The earlier 9→5 fold (`happy`/`surprised`/`other` → Excited;
    /// `disgusted`/`fearful` → Sad) collapsed emotion2vec's resolution
    /// to a single tag in practice: motivational speech, technical
    /// talks, and movie dialogue all bucketed to `[excited]`. With
    /// raw 9-class tags the LoRA trainer sees real emotional variety
    /// in its prompts.
    pub fn tag_for_class(idx: u32) -> crate::ser::EmotionLabel {
        match idx {
            0 => crate::ser::EmotionLabel::Angry,
            1 => crate::ser::EmotionLabel::Disgusted,
            2 => crate::ser::EmotionLabel::Fearful,
            3 => crate::ser::EmotionLabel::Happy,
            4 => crate::ser::EmotionLabel::Neutral,
            5 => crate::ser::EmotionLabel::Excited, // emotion2vec's "other"
            6 => crate::ser::EmotionLabel::Sad,
            7 => crate::ser::EmotionLabel::Surprised,
            _ => crate::ser::EmotionLabel::Unk, // 8 = <unk> + any unknown
        }
    }
}

impl Module for Classifier {
    fn forward(&self, pooled: &Tensor) -> candle_core::Result<Tensor> {
        // pooled: (B, embed_dim) — the time-mean-pooled hidden state.
        // Output: (B, num_classes) raw logits; caller applies softmax.
        self.proj.forward(pooled)
    }
}

/// Conv-based relative positional encoder. 5 grouped Conv1d layers with
/// kernel 19 and groups 16 — produces a positional bias of the same
/// shape as the input which is added back into the sequence.
///
/// State-dict keys: `relative_positional_encoder.{1..=5}.0.weight/bias`.
/// Indices are 1-based (no `.0.*` exists in the upstream pickle); the
/// inner `.0` is the Conv1d in a Sequential wrapper that also includes
/// a GELU at `.1` (no params).
///
/// The encoder uses **same-padding**: kernel 19 with `(19-1)/2 = 9`
/// padding on each side preserves the time dim through every layer.
#[derive(Debug)]
pub struct RelativePositionalEncoder {
    convs: Vec<Conv1d>,
}

impl RelativePositionalEncoder {
    /// Build from a VarBuilder rooted at
    /// `d2v_model.modality_encoders.AUDIO.relative_positional_encoder`.
    pub fn new(
        embed_dim: usize,
        depth: usize,
        kernel: usize,
        groups: usize,
        vb: VarBuilder,
    ) -> CsmResult<Self> {
        if kernel.is_multiple_of(2) {
            return Err(crate::CsmError::Config(format!(
                "relative_positional_encoder kernel must be odd for same-padding, got {kernel}"
            )));
        }
        let cfg = Conv1dConfig {
            padding: (kernel - 1) / 2,
            stride: 1,
            dilation: 1,
            groups,
            cudnn_fwd_algo: None,
        };
        let mut convs = Vec::with_capacity(depth);
        // Upstream indexing: 1..=depth (no `.0`).
        for i in 1..=depth {
            let conv = candle_nn::conv1d(
                embed_dim,
                embed_dim,
                kernel,
                cfg,
                vb.pp(i.to_string()).pp("0"),
            )
            .map_err(|e| crate::CsmError::Config(format!("rel_pos_enc layer {i}: {e}")))?;
            convs.push(conv);
        }
        Ok(Self { convs })
    }
}

impl Module for RelativePositionalEncoder {
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        // xs: (B, T, C). Conv1d operates on (B, C, T), so transpose.
        let mut h = xs.transpose(1, 2)?.contiguous()?;
        for conv in self.convs.iter() {
            h = conv.forward(&h)?;
            h = h.gelu()?;
        }
        // Transpose back; the result is the positional bias added to xs.
        let bias = h.transpose(1, 2)?.contiguous()?;
        xs + bias
    }
}

/// Emotion2Vec stores its inference device alongside the modules so
/// [`crate::ser::EmotionDetector::classify`] (which only takes `&[f32]`)
/// can build input tensors without an out-of-band device handle.
fn _device_struct_doc() {}

/// Top-level emotion2vec_plus_base model. Wires:
///   1. [`LocalEncoder`] (raw 16 kHz audio → 50 Hz, 512-d features)
///   2. [`ProjectFeatures`] (LN + 512→768 Linear)
///   3. [`RelativePositionalEncoder`] (additive conv positional bias)
///   4. [`ContextEncoder`] (4-block prenet + LN)
///   5. [`MainEncoder`] (8 blocks)
///   6. mean-pool over the time axis
///   7. [`Classifier`] (Linear 768→9)
///
/// Input: 16 kHz mono PCM as a `(B, 1, T)` tensor.
/// Output: `(B, 9)` raw logits — caller applies softmax + argmax.
#[derive(Debug)]
pub struct Emotion2Vec {
    cfg: Emotion2VecConfig,
    local_encoder: LocalEncoder,
    project_features: ProjectFeatures,
    rel_pos_enc: RelativePositionalEncoder,
    context_encoder: ContextEncoder,
    main_encoder: MainEncoder,
    classifier: Classifier,
    device: candle_core::Device,
}

impl Emotion2Vec {
    /// Build with random init from a VarBuilder rooted at the state-dict
    /// root. Useful for tests; production use should call
    /// [`Self::load_from_pickle`].
    ///
    /// `kernel` and `groups` for the relative positional encoder default
    /// to the `_base` config (kernel 19, groups 16).
    pub fn new(
        cfg: Emotion2VecConfig,
        vb: VarBuilder,
        device: candle_core::Device,
    ) -> CsmResult<Self> {
        // d2v_model.* sub-builder
        let vb_d = vb.pp("d2v_model");
        let vb_mod = vb_d.pp("modality_encoders").pp("AUDIO");
        let local_encoder = LocalEncoder::new(&cfg, vb_mod.pp("local_encoder").pp("conv_layers"))?;
        let project_features = ProjectFeatures::new(
            cfg.feature_dim,
            cfg.embed_dim,
            cfg.norm_eps,
            vb_mod.pp("project_features"),
        )?;
        let rel_pos_enc = RelativePositionalEncoder::new(
            cfg.embed_dim,
            5,  // _base: 5 conv layers
            19, // kernel
            16, // groups
            vb_mod.pp("relative_positional_encoder"),
        )?;
        let context_encoder = ContextEncoder::new(&cfg, vb_mod.pp("context_encoder"))?;
        let main_encoder = MainEncoder::new(&cfg, vb_d.pp("blocks"))?;
        // proj.* lives at the ROOT of the state dict, so we use `vb`,
        // not `vb_d`.
        let classifier = Classifier::new(cfg.embed_dim, cfg.num_classes, vb)?;
        Ok(Self {
            cfg,
            local_encoder,
            project_features,
            rel_pos_enc,
            context_encoder,
            main_encoder,
            classifier,
            device,
        })
    }

    /// Load the upstream `emotion2vec/emotion2vec_plus_base` checkpoint
    /// from a local `.pt` path. The pickle is a fairseq-style nested
    /// dict; the actual state lives under the `"model"` key.
    pub fn load_from_pickle<P: AsRef<std::path::Path>>(
        path: P,
        device: &candle_core::Device,
    ) -> CsmResult<Self> {
        let vb = VarBuilder::from_pth_with_state(
            path.as_ref(),
            candle_core::DType::F32,
            "model",
            device,
        )
        .map_err(|e| crate::CsmError::Config(format!("emotion2vec pth load: {e}")))?;
        Self::new(Emotion2VecConfig::plus_base(), vb, device.clone())
    }

    /// Forward pass: raw audio → 9-class logits.
    pub fn forward(&self, audio_16k: &Tensor) -> candle_core::Result<Tensor> {
        // (B, 1, T) → (B, 512, T/320)
        let feats = self.local_encoder.forward(audio_16k)?;
        // Transpose to channel-last for the transformer path.
        let feats = feats.transpose(1, 2)?.contiguous()?;
        // → (B, T', 768)
        let h = self.project_features.forward(&feats)?;
        // Add conv-based positional bias.
        let h = self.rel_pos_enc.forward(&h)?;
        // 4-block prenet.
        let h = self.context_encoder.forward(&h)?;
        // 8-block main encoder.
        let h = self.main_encoder.forward(&h)?;
        // Mean-pool over time → (B, 768).
        let pooled = h.mean(1)?;
        // (B, 9) logits.
        self.classifier.forward(&pooled)
    }

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

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

/// Implementation of the [`crate::ser::EmotionDetector`] trait so
/// `Emotion2Vec` is a drop-in replacement for the Phase 13.3
/// [`crate::ser::ProsodyDetector`] placeholder. Used by
/// `audio_to_manifest --auto-emotion-tag` and
/// `converse_server --reactive-emotion` once the user opts in via
/// `--use-emotion2vec`.
impl crate::ser::EmotionDetector for Emotion2Vec {
    fn classify(&self, samples_16k: &[f32]) -> CsmResult<crate::ser::EmotionLabel> {
        if samples_16k.is_empty() {
            return Ok(crate::ser::EmotionLabel::Neutral);
        }
        // Per-utterance zero-mean unit-variance normalization
        // (config.yaml: `normalize: true`). Without this the model
        // produces near-constant logits regardless of input — every
        // clip in the corpus would get the same label. Verified
        // empirically via examples/emotion2vec_probe before/after.
        let n = samples_16k.len();
        let mean = samples_16k.iter().sum::<f32>() / n as f32;
        let var = samples_16k.iter().map(|x| (x - mean).powi(2)).sum::<f32>() / n as f32;
        let std = var.sqrt().max(1e-7);
        let normed: Vec<f32> = samples_16k.iter().map(|x| (x - mean) / std).collect();
        let audio = Tensor::from_vec(normed, (1, 1, n), &self.device)
            .map_err(|e| crate::CsmError::Config(format!("emotion2vec audio tensor: {e}")))?;
        let logits = self
            .forward(&audio)
            .map_err(|e| crate::CsmError::Config(format!("emotion2vec forward: {e}")))?;
        let logits = logits
            .flatten_all()
            .and_then(|t| t.to_vec1::<f32>())
            .map_err(|e| crate::CsmError::Config(format!("emotion2vec logits: {e}")))?;
        let argmax = logits
            .iter()
            .enumerate()
            .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal))
            .map(|(i, _)| i)
            .unwrap_or(0);
        Ok(Classifier::tag_for_class(argmax as u32))
    }
}

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

    #[test]
    fn config_stride_product_matches_320() {
        let cfg = Emotion2VecConfig::plus_base();
        assert_eq!(cfg.conv_stride_total(), 320, "16 kHz → 50 Hz");
        assert_eq!(cfg.conv_layers.len(), 7);
    }

    #[test]
    fn project_features_shape_check() {
        let dev = Device::Cpu;
        let cfg = Emotion2VecConfig::plus_base();
        let vm = VarMap::new();
        let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
        let pf =
            ProjectFeatures::new(cfg.feature_dim, cfg.embed_dim, cfg.norm_eps, vb).expect("build");
        // (B=1, T=50, in=512) → expect (1, 50, 768)
        let xs = Tensor::zeros((1, 50, cfg.feature_dim), DType::F32, &dev).unwrap();
        let ys = <ProjectFeatures as Module>::forward(&pf, &xs).expect("forward");
        assert_eq!(ys.dims(), &[1, 50, cfg.embed_dim]);
    }

    #[test]
    fn block_residual_shape_check() {
        let dev = Device::Cpu;
        let cfg = Emotion2VecConfig::plus_base();
        let mlp_dim = (cfg.embed_dim as f32 * cfg.mlp_ratio) as usize;
        let vm = VarMap::new();
        let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
        let block =
            Block::new(cfg.embed_dim, cfg.num_heads, mlp_dim, cfg.norm_eps, vb).expect("build");
        // 8-frame input of zeros — residual should pass through, output
        // shape unchanged, and the LayerNorm guarantees no NaNs even when
        // input is all zeros.
        let xs = Tensor::randn(0f32, 1.0, (2, 8, cfg.embed_dim), &dev).unwrap();
        let ys = <Block as Module>::forward(&block, &xs).expect("forward");
        assert_eq!(ys.dims(), &[2, 8, cfg.embed_dim]);
        // Sanity: output is finite (not NaN/Inf).
        let v = ys.flatten_all().unwrap().to_vec1::<f32>().unwrap();
        assert!(
            v.iter().all(|x| x.is_finite()),
            "block produced non-finite output"
        );
    }

    #[test]
    fn context_encoder_chains_4_blocks_with_final_norm() {
        let dev = Device::Cpu;
        let cfg = Emotion2VecConfig::plus_base();
        let vm = VarMap::new();
        let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
        let enc = ContextEncoder::new(&cfg, vb).expect("build");
        let xs = Tensor::randn(0f32, 1.0, (1, 16, cfg.embed_dim), &dev).unwrap();
        let ys = <ContextEncoder as Module>::forward(&enc, &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 main_encoder_chains_8_blocks_no_final_norm() {
        let dev = Device::Cpu;
        let cfg = Emotion2VecConfig::plus_base();
        let vm = VarMap::new();
        let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
        let enc = MainEncoder::new(&cfg, vb).expect("build");
        let xs = Tensor::randn(0f32, 1.0, (1, 16, cfg.embed_dim), &dev).unwrap();
        let ys = <MainEncoder as Module>::forward(&enc, &xs).expect("forward");
        assert_eq!(ys.dims(), &[1, 16, cfg.embed_dim]);
    }

    #[test]
    fn classifier_emits_9_logits() {
        let dev = Device::Cpu;
        let cfg = Emotion2VecConfig::plus_base();
        let vm = VarMap::new();
        let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
        let cls = Classifier::new(cfg.embed_dim, cfg.num_classes, vb).expect("build");
        let pooled = Tensor::zeros((2, cfg.embed_dim), DType::F32, &dev).unwrap();
        let logits = <Classifier as Module>::forward(&cls, &pooled).expect("forward");
        assert_eq!(logits.dims(), &[2, cfg.num_classes]);
    }

    #[test]
    fn relative_positional_encoder_preserves_shape() {
        let dev = Device::Cpu;
        let cfg = Emotion2VecConfig::plus_base();
        let vm = VarMap::new();
        let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
        let rpe = RelativePositionalEncoder::new(cfg.embed_dim, 5, 19, 16, vb).expect("build");
        let xs = Tensor::randn(0f32, 1.0, (1, 32, cfg.embed_dim), &dev).unwrap();
        let ys = <RelativePositionalEncoder as Module>::forward(&rpe, &xs).expect("forward");
        // Same-padding (kernel 19 with pad 9) preserves time dim.
        assert_eq!(ys.dims(), &[1, 32, cfg.embed_dim]);
        let v = ys.flatten_all().unwrap().to_vec1::<f32>().unwrap();
        assert!(v.iter().all(|x| x.is_finite()));
    }

    #[test]
    fn emotion2vec_random_init_end_to_end_shape() {
        let dev = Device::Cpu;
        let cfg = Emotion2VecConfig::plus_base();
        let vm = VarMap::new();
        let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
        let model = Emotion2Vec::new(cfg.clone(), vb, dev.clone()).expect("build");
        // 1 second of synthetic audio. Random init won't give meaningful
        // emotion predictions but shape + finiteness should hold.
        let audio = Tensor::randn(0f32, 0.1, (1, 1, 16_000), &dev).unwrap();
        let logits = model.forward(&audio).expect("forward");
        assert_eq!(logits.dims(), &[1, cfg.num_classes]);
        let v = logits.flatten_all().unwrap().to_vec1::<f32>().unwrap();
        assert!(v.iter().all(|x| x.is_finite()), "logits had NaN/Inf");
    }

    #[test]
    fn classifier_class_to_emotion_label_mapping() {
        use crate::ser::EmotionLabel;
        // Phase 13.10: 9 raw classes map directly (no fold). The earlier
        // 9→5 fold was found to collapse all real audio to [excited]
        // empirically; raw labels give the LoRA trainer real prosodic
        // variety in its prompts.
        assert_eq!(Classifier::tag_for_class(0), EmotionLabel::Angry);
        assert_eq!(Classifier::tag_for_class(1), EmotionLabel::Disgusted);
        assert_eq!(Classifier::tag_for_class(2), EmotionLabel::Fearful);
        assert_eq!(Classifier::tag_for_class(3), EmotionLabel::Happy);
        assert_eq!(Classifier::tag_for_class(4), EmotionLabel::Neutral);
        assert_eq!(Classifier::tag_for_class(5), EmotionLabel::Excited); // emotion2vec "other"
        assert_eq!(Classifier::tag_for_class(6), EmotionLabel::Sad);
        assert_eq!(Classifier::tag_for_class(7), EmotionLabel::Surprised);
        assert_eq!(Classifier::tag_for_class(8), EmotionLabel::Unk);
    }

    #[test]
    fn local_encoder_random_init_shape_check() {
        // Build with a fresh VarMap → random Kaiming init weights, then
        // run a forward pass on synthetic audio and verify the output
        // shape matches the expected (B, 512, T/320).
        let dev = Device::Cpu;
        let cfg = Emotion2VecConfig::plus_base();
        let vm = VarMap::new();
        let vb = VarBuilder::from_varmap(&vm, DType::F32, &dev);
        let enc = LocalEncoder::new(&cfg, vb).expect("build");

        // 1 second of 16 kHz = 16_000 samples; 16_000 / 320 = 50 frames.
        let batch = 1;
        let t_in = 16_000;
        let xs = Tensor::zeros((batch, 1, t_in), DType::F32, &dev).unwrap();
        let ys = <LocalEncoder as Module>::forward(&enc, &xs).expect("forward");
        let dims = ys.dims();
        assert_eq!(dims.len(), 3, "expected 3D output, got {dims:?}");
        assert_eq!(dims[0], batch, "batch dim");
        assert_eq!(dims[1], cfg.feature_dim, "feature dim should be 512");
        // The first conv has kernel 10 stride 5; subsequent strides
        // multiply, but kernel-padding effects shrink the time dim a few
        // samples at each layer. The exact output length can be computed
        // from the conv arithmetic; for 16 000 samples the upstream
        // implementation produces 49 frames (one short of T/320 due to
        // valid-padding kernel effects).
        // Tolerance: within 2 frames of the ideal T/320=50.
        let ideal = t_in / cfg.conv_stride_total();
        assert!(
            dims[2] <= ideal && dims[2] >= ideal.saturating_sub(2),
            "time dim {} out of expected range [{}, {}]",
            dims[2],
            ideal.saturating_sub(2),
            ideal,
        );
    }
}