//! AudioSeal watermark — SEANet generator + detector.
//!
//! Architectural port of Meta's AudioSeal (Roman et al., ICML 2024,
//! [arXiv:2401.17264]) against candle 0.9.
//!
//! ## Module layout (matches `facebook/audioseal` reference exactly)
//!
//! Generator state_dict keys are organized as a `nn.Sequential` indexed by
//! integer position. With `n_filters=32, ratios=[8,5,4,2], n_residual_layers=1,
//! lstm=2, dimension=128` the encoder runs:
//!
//! ```text
//!  encoder.model.0  Conv1d(1, 32, k=7)               (SConv1d wrapper)
//!  encoder.model.1  ResidualBlock(32, dilation=1)     (n_residual=1, j=0)
//!  encoder.model.2  ELU(1.0)
//!  encoder.model.3  Conv1d(32, 64, k=4, stride=2)
//!  encoder.model.4  ResidualBlock(64, 1)
//!  encoder.model.5  ELU
//!  encoder.model.6  Conv1d(64, 128, k=8, stride=4)
//!  encoder.model.7  ResidualBlock(128, 1)
//!  encoder.model.8  ELU
//!  encoder.model.9  Conv1d(128, 256, k=10, stride=5)
//!  encoder.model.10 ResidualBlock(256, 1)
//!  encoder.model.11 ELU
//!  encoder.model.12 Conv1d(256, 512, k=16, stride=8)
//!  encoder.model.13 LSTM(512, 512, num_layers=2)      (skip-connected)
//!  encoder.model.14 ELU
//!  encoder.model.15 Conv1d(512, 128, k=7)             (dimension projection)
//! ```
//!
//! Decoder mirrors this: index 0 is the 128→512 init conv, index 1 is the
//! LSTM, indices 3,6,9,12 are ConvTranspose1d upsamples (with ratios
//! [8,5,4,2] in order), residuals at 4,7,10,13, final 32→1 conv at 15.
//!
//! The 16-bit message embedding lives at `msg_processor.msg_processor.weight`
//! shape `(32, 128)` and is broadcast-added to the encoder bottleneck
//! activations BEFORE the decoder runs.
//!
//! Detector reuses the same encoder + a mirrored upsample stack (the
//! reference instantiates a SEANetDecoder with output_channels=18 instead of
//! 1 and no `final_activation`); we expose this directly via `Detector`.
//!
//! ## Reference uses `weight_norm` parameterization
//!
//! The PyTorch checkpoint stores each Conv1d/ConvTranspose1d weight split
//! into `weight_g` (per-output-channel scale) + `weight_v` (unnormalized
//! direction). At forward time: `weight = weight_g * weight_v / ‖weight_v‖`.
//! Our offline converter merges these splits at conversion time so this
//! module's VarBuilder reads a single `weight` per layer.
//!
//! See `audioseal_convert.rs` example for the conversion path.
//!
//! [arXiv:2401.17264]: https://arxiv.org/abs/2401.17264

use crate::error::{CsmError, Result};
use crate::watermark::Watermarker;
use candle_core::{D, DType, Device, IndexOp, Module, Tensor};
use candle_nn::{
    Activation, Conv1d, Conv1dConfig, ConvTranspose1d, ConvTranspose1dConfig, Embedding, LSTM,
    LSTMConfig, RNN, VarBuilder, conv_transpose1d, conv1d, embedding, lstm, ops,
};
use std::path::Path;

pub const SAMPLE_RATE: u32 = 16_000;
/// AudioSeal embeds a 16-bit message per audio frame.
pub const MESSAGE_BITS: usize = 16;
/// Encoder downsample ratios (decoder iterates these in order; encoder reversed).
pub const RATIOS: [usize; 4] = [8, 5, 4, 2];
/// Total downsampling factor: prod(RATIOS) = 320.
pub const HOP_LENGTH: usize = 320;
/// Bottleneck channel dimension (the final 1×1 projection target).
pub const DIMENSION: usize = 128;
/// Initial filter count.
pub const N_FILTERS: usize = 32;
/// Channels at the LSTM bottleneck = N_FILTERS * 2^|ratios| = 512.
pub const LSTM_HIDDEN: usize = N_FILTERS * (1 << RATIOS.len());

// -- Conv1d/ConvTranspose1d wrappers with SEANet symmetric padding ----------

/// Apply `Conv1d` with SEANet's symmetric extra padding (non-causal mode).
/// Mirrors `audiocraft.modules.conv._get_extra_padding_for_conv1d` exactly:
///
/// ```text
/// padding_total = (kernel - 1) * dilation - (stride - 1)
/// n_frames      = (length - kernel + padding_total) / stride + 1
/// ideal_length  = (ceil(n_frames) - 1) * stride + (kernel - padding_total)
/// extra_padding = ideal_length - length
/// pad_right = padding_total // 2
/// pad_left  = padding_total - pad_right
/// padded    = pad_with_zeros(xs, pad_left, pad_right + extra_padding)
/// ```
fn padded_conv1d(
    xs: &Tensor,
    conv: &Conv1d,
    kernel: usize,
    stride: usize,
    dilation: usize,
) -> candle_core::Result<Tensor> {
    let length = xs.dim(D::Minus1)?;
    // Reference: (k - 1) * dilation - (stride - 1). For stride=1 this is (k-1)*d.
    let padding_total = ((kernel - 1) * dilation).saturating_sub(stride - 1);
    let n_frames_num = length as i64 + padding_total as i64 - kernel as i64;
    let n_frames = (n_frames_num as f64 / stride as f64) + 1.0;
    let n_frames_ceil = n_frames.ceil() as i64;
    let ideal_length =
        ((n_frames_ceil - 1) * stride as i64 + kernel as i64 - padding_total as i64) as usize;
    let extra = ideal_length.saturating_sub(length);
    let pad_right = padding_total / 2;
    let pad_left = padding_total - pad_right;
    let xs = xs.pad_with_zeros(D::Minus1, pad_left, pad_right + extra)?;
    xs.apply(conv)
}

/// Apply `ConvTranspose1d` then trim the SEANet asymmetric padding from the
/// output. Trim split: `trim_right = (k - stride) // 2`,
/// `trim_left = (k - stride) - trim_right` (non-causal, `trim_right_ratio=1.0`).
fn trimmed_conv_transpose1d(
    xs: &Tensor,
    conv: &ConvTranspose1d,
    kernel: usize,
    stride: usize,
) -> candle_core::Result<Tensor> {
    let xs = xs.apply(conv)?;
    let trim_total = kernel.saturating_sub(stride);
    let trim_right = trim_total / 2;
    let trim_left = trim_total - trim_right;
    let len = xs.dim(D::Minus1)?;
    let new_len = len.saturating_sub(trim_left + trim_right);
    if new_len == 0 {
        return Ok(xs);
    }
    xs.narrow(D::Minus1, trim_left, new_len)
}

// -- Residual block (matches `block.{1,3}.conv.conv.weight` layout) ---------

#[derive(Debug, Clone)]
pub struct SeanetResidualBlock {
    conv1: Conv1d,
    conv2: Conv1d,
    activation: Activation,
    k1: usize,
    d1: usize,
    k2: usize,
    d2: usize,
}

impl SeanetResidualBlock {
    /// `vb` is rooted at the residual block (e.g. `encoder.model.1`).
    pub fn new(dim: usize, dilation: usize, vb: VarBuilder) -> candle_core::Result<Self> {
        let hidden = dim / 2; // compress=2
        // Reference path: `block.1.conv.conv.weight` (the inner SConv1d→NormConv1d→Conv1d).
        let conv1 = conv1d(
            dim,
            hidden,
            3,
            Conv1dConfig {
                dilation,
                ..Default::default()
            },
            vb.pp("block.1.conv.conv"),
        )?;
        let conv2 = conv1d(
            hidden,
            dim,
            1,
            Conv1dConfig::default(),
            vb.pp("block.3.conv.conv"),
        )?;
        Ok(Self {
            conv1,
            conv2,
            activation: Activation::Elu(1.0),
            k1: 3,
            d1: dilation,
            k2: 1,
            d2: 1,
        })
    }
}

impl Module for SeanetResidualBlock {
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        let h = xs.apply(&self.activation)?;
        let h = padded_conv1d(&h, &self.conv1, self.k1, 1, self.d1)?;
        let h = h.apply(&self.activation)?;
        let h = padded_conv1d(&h, &self.conv2, self.k2, 1, self.d2)?;
        h + xs
    }
}

// -- LSTM bottleneck (skip-connected, matches `model.13.lstm` layout) -------

#[derive(Debug, Clone)]
pub struct LstmBottleneck {
    layers: Vec<LSTM>,
    hidden: usize,
}

impl LstmBottleneck {
    /// `vb` is rooted at the LSTM module (e.g. `encoder.model.13.lstm`). We
    /// then read `weight_ih_l0`, `weight_hh_l0`, `bias_ih_l0`, `bias_hh_l0`,
    /// `weight_ih_l1`, … directly via candle's lstm() helper.
    pub fn new(dim: usize, num_layers: usize, vb: VarBuilder) -> candle_core::Result<Self> {
        let mut layers = Vec::with_capacity(num_layers);
        for layer_idx in 0..num_layers {
            let cfg = LSTMConfig {
                layer_idx,
                ..Default::default()
            };
            layers.push(lstm(dim, dim, cfg, vb.clone())?);
        }
        Ok(Self {
            layers,
            hidden: dim,
        })
    }

    /// Input/output shape: `(B, C, T)`. SEANet uses skip = `lstm(x) + x`.
    pub fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        let (b, c, t) = xs.dims3()?;
        debug_assert_eq!(c, self.hidden);
        let mut h = xs.transpose(1, 2)?.contiguous()?;
        for layer in &self.layers {
            let init = layer.zero_state(b)?;
            let states = layer.seq_init(&h, &init)?;
            h = layer.states_to_tensor(&states)?;
        }
        let lstm_out = h.transpose(1, 2)?.contiguous()?;
        debug_assert_eq!(lstm_out.dims(), &[b, c, t]);
        lstm_out + xs
    }
}

// -- Encoder ----------------------------------------------------------------

/// Encoder stage = 1 residual block + ELU + downsample conv.
/// Module-list indices the stage occupies are `[res_idx, _, ds_idx]` since
/// the ELU at `res_idx + 1` carries no params.
#[derive(Debug, Clone)]
struct EncoderStage {
    residual: SeanetResidualBlock,
    downsample: Conv1d,
    ratio: usize,
}

#[derive(Debug, Clone)]
pub struct SeanetEncoder {
    init_conv: Conv1d,
    stages: Vec<EncoderStage>,
    lstm: LstmBottleneck,
    final_conv: Conv1d,
    activation: Activation,
}

impl SeanetEncoder {
    pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
        let m = vb.pp("model");
        let init_conv = conv1d(
            1,
            N_FILTERS,
            7,
            Conv1dConfig::default(),
            m.pp("0.conv.conv"),
        )?;

        let mut stages = Vec::with_capacity(RATIOS.len());
        let mut mult = 1usize;
        let mut idx = 1usize; // start at module index 1 (after init_conv at 0)
        for &ratio in RATIOS.iter().rev() {
            let residual = SeanetResidualBlock::new(
                mult * N_FILTERS,
                /* dilation */ 1,
                m.pp(idx.to_string()),
            )?;
            // ELU at idx+1, downsample at idx+2.
            let downsample = conv1d(
                mult * N_FILTERS,
                mult * N_FILTERS * 2,
                ratio * 2,
                Conv1dConfig {
                    stride: ratio,
                    ..Default::default()
                },
                m.pp((idx + 2).to_string()).pp("conv.conv"),
            )?;
            stages.push(EncoderStage {
                residual,
                downsample,
                ratio,
            });
            mult *= 2;
            idx += 3;
        }
        // After the final downsample (at idx-1=12 for our config), idx is now 13.
        // model.13 = LSTM, model.14 = ELU, model.15 = final conv.
        let lstm = LstmBottleneck::new(LSTM_HIDDEN, 2, m.pp(idx.to_string()).pp("lstm"))?;
        let final_conv = conv1d(
            LSTM_HIDDEN,
            DIMENSION,
            7,
            Conv1dConfig::default(),
            m.pp((idx + 2).to_string()).pp("conv.conv"),
        )?;
        Ok(Self {
            init_conv,
            stages,
            lstm,
            final_conv,
            activation: Activation::Elu(1.0),
        })
    }
}

impl Module for SeanetEncoder {
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        let mut h = padded_conv1d(xs, &self.init_conv, 7, 1, 1)?;
        for stage in &self.stages {
            h = stage.residual.forward(&h)?;
            h = h.apply(&self.activation)?;
            h = padded_conv1d(&h, &stage.downsample, stage.ratio * 2, stage.ratio, 1)?;
        }
        h = self.lstm.forward(&h)?;
        h = h.apply(&self.activation)?;
        padded_conv1d(&h, &self.final_conv, 7, 1, 1)
    }
}

// -- Decoder ----------------------------------------------------------------

#[derive(Debug, Clone)]
struct DecoderStage {
    upsample: ConvTranspose1d,
    residual: SeanetResidualBlock,
    ratio: usize,
}

/// Decoder produces `output_channels` (1 for generator, 2+nbits for detector).
#[derive(Debug, Clone)]
pub struct SeanetDecoder {
    init_conv: Conv1d,
    lstm: LstmBottleneck,
    stages: Vec<DecoderStage>,
    final_conv: Conv1d,
    activation: Activation,
}

impl SeanetDecoder {
    pub fn new(output_channels: usize, vb: VarBuilder) -> candle_core::Result<Self> {
        let m = vb.pp("model");
        // model.0 = init conv (DIMENSION → LSTM_HIDDEN), model.1 = LSTM,
        // model.2 = ELU, model.3 = upsample stage 0, model.4 = residual stage 0, ...
        let init_conv = conv1d(
            DIMENSION,
            LSTM_HIDDEN,
            7,
            Conv1dConfig::default(),
            m.pp("0.conv.conv"),
        )?;
        let lstm = LstmBottleneck::new(LSTM_HIDDEN, 2, m.pp("1.lstm"))?;
        let mut stages = Vec::with_capacity(RATIOS.len());
        let mut mult = 1usize << RATIOS.len(); // 16
        let mut idx = 3usize;
        for &ratio in RATIOS.iter() {
            let upsample = conv_transpose1d(
                mult * N_FILTERS,
                mult * N_FILTERS / 2,
                ratio * 2,
                ConvTranspose1dConfig {
                    stride: ratio,
                    ..Default::default()
                },
                m.pp(idx.to_string()).pp("convtr.convtr"),
            )?;
            let residual = SeanetResidualBlock::new(
                mult * N_FILTERS / 2,
                /* dilation */ 1,
                m.pp((idx + 1).to_string()),
            )?;
            stages.push(DecoderStage {
                upsample,
                residual,
                ratio,
            });
            mult /= 2;
            idx += 3; // upsample, residual, then ELU on next iter
        }
        // After 4 stages, idx is now 15. model.14 = ELU, model.15 = final conv.
        let final_conv = conv1d(
            N_FILTERS,
            output_channels,
            7,
            Conv1dConfig::default(),
            m.pp(idx.to_string()).pp("conv.conv"),
        )?;
        Ok(Self {
            init_conv,
            lstm,
            stages,
            final_conv,
            activation: Activation::Elu(1.0),
        })
    }
}

impl Module for SeanetDecoder {
    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        let mut h = padded_conv1d(xs, &self.init_conv, 7, 1, 1)?;
        h = self.lstm.forward(&h)?;
        for stage in &self.stages {
            h = h.apply(&self.activation)?;
            h = trimmed_conv_transpose1d(&h, &stage.upsample, stage.ratio * 2, stage.ratio)?;
            h = stage.residual.forward(&h)?;
        }
        h = h.apply(&self.activation)?;
        padded_conv1d(&h, &self.final_conv, 7, 1, 1)
    }
}

// -- 16-bit message embedding ----------------------------------------------

#[derive(Debug, Clone)]
pub struct MsgProcessor {
    table: Embedding,
    nbits: usize,
    hidden: usize,
    alpha: f32,
}

impl MsgProcessor {
    /// `vb` is rooted at the model root (NOT inside `msg_processor`); we
    /// extend by `msg_processor.msg_processor` to match the reference key
    /// `msg_processor.msg_processor.weight` of shape `(2*nbits, hidden)`.
    pub fn new(nbits: usize, hidden: usize, vb: VarBuilder) -> candle_core::Result<Self> {
        let table = embedding(2 * nbits, hidden, vb.pp("msg_processor.msg_processor"))?;
        Ok(Self {
            table,
            nbits,
            hidden,
            alpha: 1.0,
        })
    }

    /// Take encoder activations `xs: (B, C, T)` and a `u32` message of
    /// `nbits` bits; return `(B, C, T)` with the broadcast-added watermark.
    pub fn forward(
        &self,
        xs: &Tensor,
        message: u32,
        device: &Device,
        dtype: DType,
    ) -> candle_core::Result<Tensor> {
        let (b, c, t) = xs.dims3()?;
        debug_assert_eq!(c, self.hidden);
        let mut indices = Vec::with_capacity(self.nbits);
        for k in 0..self.nbits {
            let bit = (message >> k) & 1;
            indices.push(2 * k as u32 + bit);
        }
        let idx = Tensor::from_vec(indices, (self.nbits,), device)?;
        let looked = self.table.forward(&idx)?;
        let summed = looked.sum(0)?.to_dtype(dtype)?;
        let wm = summed
            .reshape((1, self.hidden, 1))?
            .broadcast_as((b, c, t))?;
        let scaled = (wm * self.alpha as f64)?;
        xs + scaled
    }
}

// -- Generator (encoder + msg + decoder) -----------------------------------

#[derive(Debug, Clone)]
pub struct Generator {
    pub encoder: SeanetEncoder,
    pub msg_processor: MsgProcessor,
    pub decoder: SeanetDecoder,
}

impl Generator {
    pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
        let encoder = SeanetEncoder::new(vb.pp("encoder"))?;
        let msg_processor = MsgProcessor::new(MESSAGE_BITS, DIMENSION, vb.clone())?;
        let decoder = SeanetDecoder::new(1, vb.pp("decoder"))?;
        Ok(Self {
            encoder,
            msg_processor,
            decoder,
        })
    }

    /// Forward `(B, 1, T) → (B, 1, T)`. Returns the **watermark residual**;
    /// `embed` adds it to the input audio with `alpha`. Trim/pad to input
    /// length so callers can sum directly without shape headaches.
    pub fn forward(&self, xs: &Tensor, message: u32) -> candle_core::Result<Tensor> {
        let device = xs.device().clone();
        let dtype = xs.dtype();
        let want = xs.dim(D::Minus1)?;
        let h = self.encoder.forward(xs)?;
        let h = self.msg_processor.forward(&h, message, &device, dtype)?;
        let h = self.decoder.forward(&h)?;
        let got = h.dim(D::Minus1)?;
        if got == want {
            Ok(h)
        } else if got > want {
            h.narrow(D::Minus1, 0, want)
        } else {
            h.pad_with_zeros(D::Minus1, 0, want - got)
        }
    }
}

// -- Detector --------------------------------------------------------------

/// Detector: `(B, 1, T) → (B, 2 + nbits, T)`.
///
/// Architecture (verbatim from `facebook/audioseal` detector_base.pth):
///   - `detector.0.model.*` — SeanetEncoder (full, ending with 128-channel
///     bottleneck at frame-rate ≈ T/320)
///   - `detector.0.reverse_convolution` — single ConvTranspose1d(128, 32,
///     kernel=320, stride=320, bias=True). Lifts frame-rate features back
///     to sample-rate (320× upsample) without weight_norm, no overlap.
///   - `detector.1` — Conv1d(32, 2+nbits, kernel=1, bias=True). Pointwise
///     head producing per-sample presence + message-bit logits.
#[derive(Debug, Clone)]
pub struct Detector {
    encoder: SeanetEncoder,
    reverse_convolution: ConvTranspose1d,
    head: Conv1d,
    nbits: usize,
}

impl Detector {
    pub fn new(vb: VarBuilder) -> candle_core::Result<Self> {
        let inner = vb.pp("detector.0");
        let encoder = SeanetEncoder::new(inner.clone())?;
        let reverse_convolution = conv_transpose1d(
            DIMENSION,
            N_FILTERS,
            HOP_LENGTH,
            ConvTranspose1dConfig {
                stride: HOP_LENGTH,
                ..Default::default()
            },
            inner.pp("reverse_convolution"),
        )?;
        let head = conv1d(
            N_FILTERS,
            2 + MESSAGE_BITS,
            1,
            Conv1dConfig::default(),
            vb.pp("detector.1"),
        )?;
        Ok(Self {
            encoder,
            reverse_convolution,
            head,
            nbits: MESSAGE_BITS,
        })
    }

    /// `(B, 1, T) → (B, 2+nbits, T)` per-sample logits.
    pub fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
        let h = self.encoder.forward(xs)?;
        // Single-shot 320× upsample. ConvTranspose1d with k=stride=320 → no overlap.
        let h = h.apply(&self.reverse_convolution)?;
        // Re-narrow to original T (the upsample may produce slightly more samples
        // than the input, depending on encoder rounding).
        let want = xs.dim(D::Minus1)?;
        let got = h.dim(D::Minus1)?;
        let h = if got == want {
            h
        } else if got > want {
            h.narrow(D::Minus1, 0, want)?
        } else {
            h.pad_with_zeros(D::Minus1, 0, want - got)?
        };
        h.apply(&self.head)
    }

    /// Decode per-sample presence + message bits from `(B, 2+nbits, T)` logits.
    pub fn decode(&self, logits: &Tensor) -> candle_core::Result<(Tensor, u16, f32)> {
        let presence_logits = logits.i((.., ..2, ..))?;
        let message_logits = logits.i((.., 2.., ..))?;
        let presence_probs = ops::softmax(&presence_logits, 1)?;
        let presence = presence_probs.i((.., 1, ..))?;
        let bit_probs = ops::sigmoid(&message_logits)?.mean(D::Minus1)?;
        let bits: Vec<f32> = bit_probs.i(0)?.to_dtype(DType::F32)?.to_vec1()?;
        let mut decoded: u16 = 0;
        for (k, p) in bits.iter().enumerate().take(self.nbits) {
            if *p > 0.5 {
                decoded |= 1 << k;
            }
        }
        let mean_presence = presence
            .mean_all()?
            .to_dtype(DType::F32)?
            .to_scalar::<f32>()?;
        Ok((presence, decoded, mean_presence))
    }
}

// -- Public watermarker (Watermarker trait surface) ------------------------

#[derive(Debug, Clone)]
pub struct DetectionResult {
    pub presence_per_sample: Vec<f32>,
    pub message: Option<u16>,
    pub mean_presence: f32,
}

pub struct AudioSealWatermarker {
    pub device: Device,
    pub message: u16,
    pub generator: Option<Generator>,
    pub detector: Option<Detector>,
    pub alpha: f32,
}

impl AudioSealWatermarker {
    pub fn new(device: Device, message: u16) -> Self {
        Self {
            device,
            message,
            generator: None,
            detector: None,
            alpha: 1.0,
        }
    }

    /// Construct from in-memory VarBuilders rooted at the generator and
    /// detector subtrees (after `weight_norm` merge by the converter).
    pub fn from_var_builders(
        generator_vb: VarBuilder,
        detector_vb: VarBuilder,
        device: Device,
        message: u16,
    ) -> Result<Self> {
        let generator = Generator::new(generator_vb)
            .map_err(|e| CsmError::Config(format!("AudioSeal generator load: {e}")))?;
        let detector = Detector::new(detector_vb)
            .map_err(|e| CsmError::Config(format!("AudioSeal detector load: {e}")))?;
        Ok(Self {
            device,
            message,
            generator: Some(generator),
            detector: Some(detector),
            alpha: 1.0,
        })
    }

    pub fn load<P: AsRef<Path>>(_weights_dir: P, device: Device, message: u16) -> Result<Self> {
        Ok(Self::new(device, message))
    }

    pub fn detect(&self, samples: &[f32]) -> Result<DetectionResult> {
        let detector = self.detector.as_ref().ok_or_else(|| {
            CsmError::Config(
                "AudioSeal::detect: detector weights not loaded — see audioseal_convert example"
                    .into(),
            )
        })?;
        let xs = Tensor::from_slice(samples, (1, 1, samples.len()), &self.device)
            .map_err(|e| CsmError::Config(format!("detect: input tensor: {e}")))?;
        let logits = detector
            .forward(&xs)
            .map_err(|e| CsmError::Config(format!("detect: forward: {e}")))?;
        let (presence, message, mean_presence) = detector
            .decode(&logits)
            .map_err(|e| CsmError::Config(format!("detect: decode: {e}")))?;
        let presence_per_sample = presence
            .i(0)
            .and_then(|t| t.to_dtype(DType::F32))
            .and_then(|t| t.to_vec1::<f32>())
            .map_err(|e| CsmError::Config(format!("detect: presence to_vec: {e}")))?;
        Ok(DetectionResult {
            presence_per_sample,
            message: Some(message),
            mean_presence,
        })
    }
}

impl AudioSealWatermarker {
    /// Internal implementation used by both `embed` (default message) and
    /// `embed_with_message` (per-call override).
    fn embed_inner(&self, audio: &[f32], message: u16) -> Result<Vec<f32>> {
        let generator = self.generator.as_ref().ok_or_else(|| {
            CsmError::Config(
                "AudioSeal::embed: generator weights not loaded — see audioseal_convert example"
                    .into(),
            )
        })?;
        let xs = Tensor::from_slice(audio, (1, 1, audio.len()), &self.device)
            .map_err(|e| CsmError::Config(format!("embed: input tensor: {e}")))?;
        let residual = generator
            .forward(&xs, message as u32)
            .map_err(|e| CsmError::Config(format!("embed: forward: {e}")))?;
        let scaled = (residual * self.alpha as f64)
            .map_err(|e| CsmError::Config(format!("embed: scale: {e}")))?;
        let out = (xs + scaled).map_err(|e| CsmError::Config(format!("embed: sum: {e}")))?;
        let len = out
            .dim(D::Minus1)
            .map_err(|e| CsmError::Config(e.to_string()))?;
        let flat = out
            .reshape((len,))
            .and_then(|t| t.to_dtype(DType::F32))
            .and_then(|t| t.to_vec1::<f32>())
            .map_err(|e| CsmError::Config(format!("embed: to_vec: {e}")))?;
        Ok(flat)
    }
}

impl Watermarker for AudioSealWatermarker {
    fn embed(&self, audio: &[f32]) -> Result<Vec<f32>> {
        self.embed_inner(audio, self.message)
    }

    fn embed_with_message(&self, audio: &[f32], message: u16) -> Result<Vec<f32>> {
        self.embed_inner(audio, message)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use candle_nn::{VarBuilder, 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 scaffold_constructs() {
        let w = AudioSealWatermarker::new(Device::Cpu, 0xABCD);
        assert_eq!(w.message, 0xABCD);
    }

    #[test]
    fn embed_returns_typed_error_without_weights() {
        let w = AudioSealWatermarker::new(Device::Cpu, 0);
        assert!(w.embed(&[0.0, 0.1, 0.2]).is_err());
    }

    #[test]
    fn detect_returns_typed_error_without_weights() {
        let w = AudioSealWatermarker::new(Device::Cpu, 0);
        assert!(w.detect(&[0.0, 0.1, 0.2]).is_err());
    }

    #[test]
    fn residual_block_preserves_shape() {
        let device = Device::Cpu;
        let (_vm, vb) = random_vb(&device);
        let block = SeanetResidualBlock::new(64, 1, vb).unwrap();
        let xs = Tensor::randn(0f32, 1f32, (2, 64, 100), &device).unwrap();
        let out = block.forward(&xs).unwrap();
        assert_eq!(out.dims(), &[2, 64, 100]);
    }

    #[test]
    fn lstm_bottleneck_preserves_shape() {
        let device = Device::Cpu;
        let (_vm, vb) = random_vb(&device);
        let lstm = LstmBottleneck::new(LSTM_HIDDEN, 2, vb).unwrap();
        let xs = Tensor::randn(0f32, 1f32, (2, LSTM_HIDDEN, 50), &device).unwrap();
        let out = lstm.forward(&xs).unwrap();
        assert_eq!(out.dims(), &[2, LSTM_HIDDEN, 50]);
    }

    #[test]
    fn msg_processor_preserves_shape() {
        let device = Device::Cpu;
        let (_vm, vb) = random_vb(&device);
        let msg = MsgProcessor::new(16, 128, vb).unwrap();
        let xs = Tensor::randn(0f32, 1f32, (1, 128, 25), &device).unwrap();
        let out = msg.forward(&xs, 0xABCDu32, &device, DType::F32).unwrap();
        assert_eq!(out.dims(), &[1, 128, 25]);
    }

    #[test]
    fn encoder_downsamples_by_320() {
        let device = Device::Cpu;
        let (_vm, vb) = random_vb(&device);
        let enc = SeanetEncoder::new(vb).unwrap();
        let xs = Tensor::randn(0f32, 1f32, (1, 1, 16000), &device).unwrap();
        let out = enc.forward(&xs).unwrap();
        let frames = out.dim(D::Minus1).unwrap();
        assert_eq!(out.dim(0).unwrap(), 1);
        assert_eq!(out.dim(1).unwrap(), DIMENSION);
        assert!(
            (frames as i64 - 50).abs() <= 2,
            "encoder frames expected ~50, got {frames}",
        );
    }

    #[test]
    fn generator_round_trips_shape() {
        let device = Device::Cpu;
        let (_vm, vb) = random_vb(&device);
        let g = Generator::new(vb).unwrap();
        let t = 16000;
        let xs = Tensor::randn(0f32, 1f32, (1, 1, t), &device).unwrap();
        let out = g.forward(&xs, 0xBEEFu32).unwrap();
        assert_eq!(out.dim(0).unwrap(), 1);
        assert_eq!(out.dim(1).unwrap(), 1);
        let got = out.dim(D::Minus1).unwrap();
        let drift = (got as i64 - t as i64).abs() as usize;
        assert!(
            drift * 100 < t,
            "generator length drift too large: {drift} samples"
        );
    }

    #[test]
    fn detector_emits_18_channel_logits() {
        let device = Device::Cpu;
        let (_vm, vb) = random_vb(&device);
        let det = Detector::new(vb).unwrap();
        let t = 16000;
        let xs = Tensor::randn(0f32, 1f32, (1, 1, t), &device).unwrap();
        let logits = det.forward(&xs).unwrap();
        assert_eq!(logits.dim(0).unwrap(), 1);
        assert_eq!(logits.dim(1).unwrap(), 2 + MESSAGE_BITS);
        assert_eq!(logits.dim(D::Minus1).unwrap(), t);
    }

    #[test]
    fn detector_decode_yields_u16_message() {
        let device = Device::Cpu;
        let (_vm, vb) = random_vb(&device);
        let det = Detector::new(vb).unwrap();
        let xs = Tensor::randn(0f32, 1f32, (1, 1, 16000), &device).unwrap();
        let logits = det.forward(&xs).unwrap();
        let (presence, _message, mean_presence) = det.decode(&logits).unwrap();
        assert_eq!(presence.dim(D::Minus1).unwrap(), 16000);
        assert!(mean_presence.is_finite());
    }
}