"""Standalone VocBulwark speaker encoder.

Loadable with `AutoModel.from_pretrained(path, trust_remote_code=True)` without
the training repository. Maps a raw waveform (mono, at `config.raw_sample_rate`)
to a fixed-size speaker embedding used to condition the companion VocBulwark
vocoder.
"""
from dataclasses import dataclass

import torch
import torch.nn.utils.parametrize as parametrize
from torch.nn.utils import remove_weight_norm
from transformers import PreTrainedModel
from transformers.modeling_outputs import ModelOutput


def _strip_weight_norm(root):
    """Remove weight-norm reparametrizations (both conventions), leaving plain
    ``.weight`` parameters that match the folded export weights."""
    for m in root.modules():
        if parametrize.is_parametrized(m, "weight"):
            parametrize.remove_parametrizations(m, "weight", leave_parametrized=True)
        elif hasattr(m, "weight_g"):
            try:
                remove_weight_norm(m)
            except (ValueError, RuntimeError):
                pass

from .configuration_speaker import SpeakerEncoderConfig
from .speaker_embedding_config import SpeakerEmbeddingConfig
from .speaker_embedding_architecture import SpeakerEmbeddingArchitecture

# Force the loader to copy the (transitively used) leaf module into the cache.
from .ge2e_loss import GE2ELoss as _GE2ELoss  # noqa: F401


@dataclass
class SpeakerEmbeddingOutput(ModelOutput):
    embeddings: torch.Tensor = None


class SpeakerEncoder(PreTrainedModel):
    config_class = SpeakerEncoderConfig

    def __init__(self, config):
        super().__init__(config)
        spk_cfg = config.speaker_embed_config
        if isinstance(spk_cfg, dict):
            spk_cfg = SpeakerEmbeddingConfig.from_dict(spk_cfg)
        self.speaker_embedding_model = SpeakerEmbeddingArchitecture(spk_cfg)
        self.speaker_embedding_model.eval()
        # Export folds weight-norm into plain .weight; strip the reparametrization
        # so module keys match and inference is convention-independent.
        _strip_weight_norm(self)

    @torch.no_grad()
    def forward(self, raw_audio=None, attention_mask=None, input_values=None, **kwargs):
        """Encode a raw waveform into a speaker embedding.

        Args:
            raw_audio: [B, samples] (or [samples]) mono audio @ config.raw_sample_rate.
            attention_mask: optional [B, samples] mask.
        Returns:
            SpeakerEmbeddingOutput with `.embeddings` of shape [B, embedding_size].
        """
        audio = raw_audio if raw_audio is not None else input_values
        if audio is None:
            raise ValueError("Pass raw_audio [B, samples] @ config.raw_sample_rate.")
        p = next(self.parameters())
        audio = audio.to(device=p.device, dtype=p.dtype)
        if audio.dim() == 1:
            audio = audio.unsqueeze(0)
        if attention_mask is not None:
            attention_mask = attention_mask.to(device=p.device)
        emb = self.speaker_embedding_model(audio, attention_mask=attention_mask).embeddings
        return SpeakerEmbeddingOutput(embeddings=emb)

    @torch.no_grad()
    def embed(self, raw_audio, attention_mask=None):
        """Convenience: return just the [B, embedding_size] embedding tensor."""
        return self.forward(raw_audio=raw_audio, attention_mask=attention_mask).embeddings