from typing import Optional, Tuple

import torch
import torch.nn.functional as F
from torch import nn
from transformers.modeling_outputs import BaseModelOutputWithPooling
from transformers import PreTrainedModel

try:
    from .ganda_config import UnModelConfig
except ImportError:
    from ganda_config import UnModelConfig


class UnModel(PreTrainedModel):
    """Transformer sentence encoder with a normalized 3072-dimensional output."""

    config_class = UnModelConfig
    main_input_name = "input_ids"

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

        self.d_model = config.hidden_size
        self.embed = nn.Embedding(
            config.vocab_size,
            config.hidden_size,
            padding_idx=config.pad_token_id,
        )
        self.pos_proj = nn.Embedding(
            config.max_position_embeddings,
            config.hidden_size,
        )

        # Keep the encoder_layers attribute as well as encoder.layers. The
        # original training class created both, so this preserves its state_dict
        # key layout (including encoder_layers.*, which is not used in forward).
        self.encoder_layers = nn.TransformerEncoderLayer(
            d_model=config.hidden_size,
            nhead=config.num_attention_heads,
            batch_first=True,
            norm_first=True,
            dim_feedforward=config.intermediate_size,
            dropout=config.dropout,
            activation=config.hidden_act,
        )
        self.encoder = nn.TransformerEncoder(
            encoder_layer=self.encoder_layers,
            num_layers=config.num_hidden_layers,
            norm=nn.RMSNorm(config.hidden_size),
        )
        self.prog_head = nn.Linear(config.hidden_size, config.embedding_dim)

        self.post_init()

    def _init_weights(self, module):
        # PyTorch layers initialize themselves on construction. Leaving those
        # values untouched also preserves the original model's initialization
        # behavior when a new model is instantiated before loading a checkpoint.
        return

    def get_input_embeddings(self):
        return self.embed

    def set_input_embeddings(self, value):
        self.embed = value

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        output_hidden_states: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ):
        if input_ids is None and inputs_embeds is None:
            raise ValueError("Pass either input_ids or inputs_embeds.")
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("Pass input_ids or inputs_embeds, not both.")

        if inputs_embeds is None:
            if input_ids.ndim != 2:
                raise ValueError(
                    f"input_ids must have shape [batch, sequence], got {tuple(input_ids.shape)}."
                )
            inputs_embeds = self.embed(input_ids)

        batch_size, sequence_length, _ = inputs_embeds.shape
        if sequence_length > self.config.max_position_embeddings:
            raise ValueError(
                f"Sequence length {sequence_length} exceeds the model limit of "
                f"{self.config.max_position_embeddings}."
            )

        if attention_mask is None:
            attention_mask = torch.ones(
                (batch_size, sequence_length),
                dtype=torch.long,
                device=inputs_embeds.device,
            )
        elif attention_mask.shape != (batch_size, sequence_length):
            raise ValueError(
                "attention_mask must have shape "
                f"{(batch_size, sequence_length)}, got {tuple(attention_mask.shape)}."
            )

        position_ids = torch.arange(sequence_length, device=inputs_embeds.device)
        hidden_states = inputs_embeds + self.pos_proj(position_ids).unsqueeze(0)
        src_key_padding_mask = attention_mask.eq(0)

        output_hidden_states = (
            self.config.output_hidden_states
            if output_hidden_states is None
            else output_hidden_states
        )
        output_attentions = (
            self.config.output_attentions
            if output_attentions is None
            else output_attentions
        )
        return_dict = self.config.use_return_dict if return_dict is None else return_dict

        if output_attentions:
            raise NotImplementedError(
                "PyTorch TransformerEncoderLayer does not expose attention weights "
                "through this model implementation."
            )

        all_hidden_states: Optional[Tuple[torch.Tensor, ...]] = None
        if output_hidden_states:
            history = [hidden_states]
            for layer in self.encoder.layers:
                hidden_states = layer(
                    hidden_states,
                    src_key_padding_mask=src_key_padding_mask,
                )
                history.append(hidden_states)
            if self.encoder.norm is not None:
                hidden_states = self.encoder.norm(hidden_states)
                history[-1] = hidden_states
            all_hidden_states = tuple(history)
        else:
            hidden_states = self.encoder(
                hidden_states,
                src_key_padding_mask=src_key_padding_mask,
            )

        # Masked mean pooling, matching the original training implementation.
        pooling_mask = attention_mask.unsqueeze(-1).to(dtype=hidden_states.dtype)
        summed = torch.sum(hidden_states * pooling_mask, dim=1)
        denominator = torch.clamp(pooling_mask.sum(dim=1), min=1e-9)
        pooled = summed / denominator

        embeddings = self.prog_head(pooled)
        embeddings = F.normalize(embeddings)

        if not return_dict:
            output = (hidden_states, embeddings)
            if all_hidden_states is not None:
                output += (all_hidden_states,)
            return output

        return BaseModelOutputWithPooling(
            last_hidden_state=hidden_states,
            pooler_output=embeddings,
            hidden_states=all_hidden_states,
        )

    @torch.no_grad()
    def generate_embeddings(
        self,
        input_ids: torch.LongTensor,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """Return L2-normalized sentence embeddings with shape [batch, embedding_dim]."""
        if input_ids.ndim == 1:
            input_ids = input_ids.unsqueeze(0)
        if attention_mask is not None and attention_mask.ndim == 1:
            attention_mask = attention_mask.unsqueeze(0)
        return self(
            input_ids=input_ids,
            attention_mask=attention_mask,
            return_dict=True,
        ).pooler_output


UnModel.register_for_auto_class("AutoModel")