from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn
from transformers import AutoConfig, AutoModel, PreTrainedModel
from transformers.utils import ModelOutput

from .configuration_genaid import GenAIDConfig


@dataclass
class GenAIDOutput(ModelOutput):
    embedding: torch.FloatTensor = None
    accent_logits: Optional[torch.FloatTensor] = None
    speaker_logits: Optional[torch.FloatTensor] = None


class GenAIDModel(PreTrainedModel):
    config_class = GenAIDConfig
    base_model_prefix = "genaid"
    main_input_name = "input_values"
    # GenAID has no tied parameters. Transformers 5.x expects custom models to
    # expose this mapping while finalizing low-memory checkpoint loading.
    all_tied_weights_keys = {}

    def __init__(self, config):
        super().__init__(config)
        encoder_dict = dict(config.encoder_config)
        model_type = encoder_dict.pop("model_type")
        encoder_config = AutoConfig.for_model(model_type, **encoder_dict)
        self.encoder = AutoModel.from_config(encoder_config)
        hidden = encoder_config.hidden_size
        dim = config.bottleneck_dim
        self.bottleneck = nn.Sequential(
            nn.Linear(hidden, dim), nn.GELU(), nn.Linear(dim, dim), nn.GELU()
        )
        self.accent_classifier = nn.Linear(dim, config.num_accents, bias=False)
        self.speaker_classifier = nn.Linear(dim, config.num_speakers, bias=False)

    @staticmethod
    def masked_mean(hidden_states, attention_mask):
        if attention_mask is None:
            return hidden_states.mean(1)
        lengths = attention_mask.sum(-1)
        frame_lengths = (lengths * hidden_states.shape[1] / attention_mask.shape[1]).ceil().long()
        frame_lengths = frame_lengths.clamp(1, hidden_states.shape[1])
        frame_mask = torch.arange(hidden_states.shape[1], device=hidden_states.device)[None]
        frame_mask = frame_mask < frame_lengths[:, None]
        return (hidden_states * frame_mask.unsqueeze(-1)).sum(1) / frame_lengths.unsqueeze(-1)

    def forward(self, input_values, attention_mask=None, return_dict=True, **kwargs):
        encoded = self.encoder(
            input_values=input_values,
            attention_mask=attention_mask,
            return_dict=True,
            **kwargs,
        ).last_hidden_state
        embedding = self.bottleneck(self.masked_mean(encoded, attention_mask))
        accent_logits = self.accent_classifier(embedding)
        speaker_logits = self.speaker_classifier(embedding)
        if not return_dict:
            return embedding, accent_logits, speaker_logits
        return GenAIDOutput(
            embedding=embedding,
            accent_logits=accent_logits,
            speaker_logits=speaker_logits,
        )