"""Self-contained inference model for the VocBulwark HiFi-GAN / BigVGAN vocoder.

Loadable with `AutoModel.from_pretrained(path, trust_remote_code=True)` without
the training repository. This is the lean vocoder: it takes a mel-spectrogram and
a precomputed speaker embedding, and does NOT include the speaker encoder (that
is a separate model), the training losses, discriminators, or watermark heads.
"""
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 every weight-norm reparametrization (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_hifigan import HiFiGANConfig
from .bigvgan_model import BigVGAN

# transformers' trust_remote_code loader copies only modules *directly* imported
# by this entry file, and recognizes the `from .X import Y` form. Import a real
# symbol from every bundled leaf module so all files get copied into the
# dynamic-module cache (they are used transitively by the imports above).
from .bigvgan_activations import Snake as _Snake  # noqa: F401
from .alias_free_act import Activation1d as _Activation1d  # noqa: F401
from .alias_free_resample import UpSample1d as _UpSample1d  # noqa: F401
from .alias_free_filter import LowPassFilter1d as _LowPassFilter1d  # noqa: F401
from .temporal_adapter import TemporalAdapterBlock as _TemporalAdapterBlock  # noqa: F401


@dataclass
class HiFiGANOutput(ModelOutput):
    audio: torch.Tensor = None


class HiFiGANArchitecture(PreTrainedModel):
    config_class = HiFiGANConfig

    def __init__(self, config):
        super().__init__(config)

        # Lean vocoder: the speaker embedding is supplied at inference time, so the
        # generator only needs its dimension to build the conditioning conv.
        global_channels = int(getattr(config, "speaker_embedding_size", 0)) or -1
        self.hifi_gan = BigVGAN(
            num_mels=config.hifigan_in_channels,
            upsample_initial_channel=config.hifigan_channels,
            resblock_kernel_sizes=config.hifigan_resblock_kernel_sizes,
            resblock_dilation_sizes=config.hifigan_resblock_dilations,
            upsample_kernel_sizes=config.hifigan_upsample_kernel_sizes,
            upsample_rates=config.hifigan_upsample_scales,
            snake_logscale=config.snake_logscale,
            activation=config.hifigan_nonlinear_activation,
            use_bias_at_final=config.hifigan_bias,
            global_channels=global_channels,
            watermark_bits=config.watermark_bits if config.add_film_watermark else 0,
            watermark_film_hidden=config.watermark_film_hidden,
            vocbulwark_bits=config.vocbulwark_bits if config.add_watermark_vocbulwark else 0,
            vocbulwark_ta_hidden=config.vocbulwark_ta_hidden,
            vocbulwark_zero_conv_std=getattr(config, "vocbulwark_zero_conv_std", 0.02),
        )

        # Export folds weight-norm into plain .weight; strip the reparametrization
        # here so module keys match (and inference is independent of torch's
        # weight_norm convention).
        _strip_weight_norm(self)

    @torch.no_grad()
    def forward(
        self,
        mel_spectrogram=None,
        speaker_embedding=None,
        input_features=None,
        return_loss=False,
        **kwargs,
    ):
        """Vocode a (log-)mel spectrogram into a waveform.

        Args:
            mel_spectrogram:   [B, mel_channels, T] input features.
            speaker_embedding: [B, speaker_embedding_size] speaker conditioning,
                               produced by the companion speaker-encoder model.
        Returns:
            HiFiGANOutput with `.audio` of shape [B, 1, samples] @ target_sample_rate.

        The model's fixed provenance watermark is embedded in every generated clip
        and cannot be changed or disabled through this interface.
        """
        hifi_gan_device = self.hifi_gan.conv_pre.bias.device
        if torch.is_autocast_enabled():
            target_dtype = torch.get_autocast_gpu_dtype()
        else:
            target_dtype = self.hifi_gan.conv_pre.bias.dtype

        if input_features is not None:
            features = input_features.to(device=hifi_gan_device, dtype=target_dtype)
        else:
            features = mel_spectrogram.to(device=hifi_gan_device, dtype=target_dtype)

        if speaker_embedding is None:
            raise ValueError(
                "speaker_embedding is required: pass a "
                "[B, config.speaker_embedding_size] tensor from the companion "
                "speaker-encoder model.")
        g = speaker_embedding.to(device=hifi_gan_device, dtype=target_dtype)
        if g.dim() == 2:                       # [B, E] -> [B, E, 1]
            g = g.unsqueeze(2)

        # Provenance watermark: the model's fixed signature is embedded in every
        # generated clip. It is intentionally not overridable through this
        # interface — there is no way to disable or change it here.
        watermark = None
        fw = getattr(self.config, "fixed_watermark", None)
        if fw is not None:
            watermark = torch.tensor(fw, dtype=target_dtype, device=hifi_gan_device)
            watermark = watermark.unsqueeze(0).expand(features.shape[0], -1)

        audio = self.hifi_gan(features, g=g, watermark=watermark)
        return HiFiGANOutput(audio=audio)