# MIT License
# 
# Copyright (c) 2026 audio-embeddings contributors
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

"""Trainable audio embeddings with the same extraction policy as HEAR."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

import torch
from torch.nn.utils.rnn import pad_sequence
from transformers import PreTrainedModel
from transformers.utils import ModelOutput

from .adapters import SpectrogramPatchAdapter, WaveformConvAdapter
from .adapters import resolve_adapter_spec
from .configuration_audio import AudioEmbeddingConfig


@dataclass
class AudioEmbeddingOutput(ModelOutput):
    last_hidden_state: torch.Tensor | None = None
    pooler_output: torch.Tensor | None = None
    attention_mask: torch.Tensor | None = None
    timestamps_ms: torch.Tensor | None = None


class AudioEmbeddingModel(PreTrainedModel):
    config_class = AudioEmbeddingConfig
    base_model_prefix = "adapter"
    main_input_name = "input_values"
    # RoPE modules are shared by every attention block. Save all buffer keys,
    # so loading does not need special tied-buffer handling.
    _supports_assign_param_buffer = False

    def __init__(self, config: AudioEmbeddingConfig) -> None:
        super().__init__(config)
        adapter_config = config.to_adapter_config()
        # Transformers 5 loads under a default meta device, but torchaudio's
        # filter-bank constructors need real values. Build on CPU; HF subsequently
        # loads the checkpoint tensors onto the requested device/dtype.
        with torch.device("cpu"):
            spec = resolve_adapter_spec(adapter_config)
            adapter_type = (
                SpectrogramPatchAdapter
                if spec.adapter_key == "spectrogram_patch"
                else WaveformConvAdapter
            )
            self.adapter = adapter_type(adapter_config, spec)
        self.post_init()

    def _init_weights(self, module: torch.nn.Module) -> None:
        """Preserve initialization performed by the research components themselves."""

    def save_pretrained(
        self, save_directory: str | Path, *args: Any, **kwargs: Any
    ) -> None:
        if kwargs.get("state_dict") is None:
            kwargs["state_dict"] = {
                key: value.detach().clone().contiguous()
                for key, value in self.state_dict().items()
            }
        return super().save_pretrained(save_directory, *args, **kwargs)

    def forward(
        self,
        input_values: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        return_dict: bool | None = None,
    ) -> AudioEmbeddingOutput | tuple[torch.Tensor, ...]:
        if input_values.ndim != 2 or min(input_values.shape) <= 0:
            raise ValueError(
                "input_values must have shape [batch, samples] with nonempty axes"
            )
        if (
            not input_values.is_floating_point()
            or not torch.isfinite(input_values).all()
        ):
            raise ValueError(
                "input_values must contain finite floating-point waveforms"
            )
        if attention_mask is None:
            lengths = [input_values.shape[1]] * input_values.shape[0]
        else:
            if attention_mask.shape != input_values.shape:
                raise ValueError(
                    "attention_mask must have the same shape as input_values"
                )
            if not torch.all((attention_mask == 0) | (attention_mask == 1)):
                raise ValueError("attention_mask must contain only zeros and ones")
            lengths_tensor = attention_mask.long().sum(dim=1)
            expected = (
                torch.arange(input_values.shape[1], device=attention_mask.device)[None]
                < lengths_tensor[:, None]
            )
            if not torch.equal(attention_mask.bool(), expected) or torch.any(
                lengths_tensor == 0
            ):
                raise ValueError(
                    "attention_mask must describe nonempty, right-padded waveforms"
                )
            lengths = lengths_tensor.tolist()
        # Clear non-buffer RoPE caches between calls: inference-mode caches cannot
        # be reused for autograd, and .to(device/dtype) does not move these caches.
        rope = self.adapter.encoder.rope
        if rope is not None:
            for name in ("cached_cos_sin", "cached_cos_sin_h", "cached_cos_sin_w"):
                if hasattr(rope, name):
                    setattr(rope, name, None)
        outputs = []
        for waveform, length in zip(input_values, lengths):
            if isinstance(self.adapter, SpectrogramPatchAdapter):
                minimum = self.adapter.spectrogram.mel_spec.n_fft // 2 + 1
                if length < minimum:
                    raise ValueError(
                        f"Audio requires at least {minimum} samples for this spectrogram; got {length}"
                    )
            outputs.append(
                self.adapter.extract(
                    waveform[:length], preset_name=self.config.extraction_preset
                )
            )
        hidden = pad_sequence(
            [item.timestamp_embeddings for item in outputs], batch_first=True
        )
        frame_lengths = torch.tensor(
            [item.timestamp_embeddings.shape[0] for item in outputs],
            device=hidden.device,
        )
        frame_mask = (
            torch.arange(hidden.shape[1], device=hidden.device)[None]
            < frame_lengths[:, None]
        )
        result = AudioEmbeddingOutput(
            last_hidden_state=hidden,
            pooler_output=torch.stack([item.scene_embedding for item in outputs]),
            attention_mask=frame_mask.long(),
            timestamps_ms=pad_sequence(
                [item.timestamps_ms for item in outputs],
                batch_first=True,
                padding_value=-1.0,
            ),
        )
        return (
            result
            if (self.config.return_dict if return_dict is None else return_dict)
            else result.to_tuple()
        )


AudioEmbeddingModel.register_for_auto_class("AutoModel")

from .adapters import __name__ as _bundled_adapters  # noqa: F401
from .extraction import __name__ as _bundled_extraction  # noqa: F401
from .patch_embed import __name__ as _bundled_patch_embed  # noqa: F401
from .spectrogram import __name__ as _bundled_spectrogram  # noqa: F401
from .vit import __name__ as _bundled_vit  # noqa: F401
from .rope import __name__ as _bundled_rope  # noqa: F401
from .transformer import __name__ as _bundled_transformer  # noqa: F401
from .normalization import __name__ as _bundled_normalization  # noqa: F401
from .waveform_feature_encoder import __name__ as _bundled_waveform_feature_encoder  # noqa: F401