from torch import nn
from torch.nn import CrossEntropyLoss
from transformers import AutoModel
from transformers.models.modernvbert.modeling_modernvbert import (
    ModernVBertConnector,
    ModernVBertForSequenceClassification,
    ModernVBertForTokenClassification,
    ModernVBertMaskedLMOutput,
    ModernVBertModel,
    ModernVBertPreTrainedModel,
)
from transformers.models.ultrabert.modeling_ultrabert import UltrabertPredictionHead, UltrabertRMSNorm

from .configuration_ultravbert import UltraVBertConfig


class UltraVBertConnector(ModernVBertConnector):
    def __init__(self, config: UltraVBertConfig):
        super().__init__(config)
        # UltraBERT's first layer takes embeddings unnormalized; raw projections would swamp its token embeddings.
        self.norm = (
            UltrabertRMSNorm(config.text_config.hidden_size, eps=config.text_config.rms_norm_eps)
            if config.connector_norm
            else nn.Identity()
        )

    def forward(self, image_hidden_states):
        return self.norm(super().forward(image_hidden_states))


class UltraVBertModel(ModernVBertModel):
    config_class = UltraVBertConfig
    _keys_to_ignore_on_load_unexpected = [r"vision_model\.head\."]

    def __init__(self, config: UltraVBertConfig):
        # ModernVBertModel.__init__ builds ModernVBertConnector; the rest of its body follows unchanged.
        ModernVBertPreTrainedModel.__init__(self, config)
        self.padding_idx = config.text_config.pad_token_id
        self.vocab_size = config.text_config.vocab_size
        self.vision_model = AutoModel.from_config(config.vision_config)
        self.connector = UltraVBertConnector(config)
        self.text_model = AutoModel.from_config(config.text_config)
        self.image_seq_len = int(
            ((config.vision_config.image_size // config.vision_config.patch_size) ** 2)
            / (config.pixel_shuffle_factor**2)
        )
        self.image_token_id = config.image_token_id
        self.post_init()


class UltraVBertForMaskedLM(ModernVBertPreTrainedModel):
    config_class = UltraVBertConfig
    _keys_to_ignore_on_load_unexpected = [r"vision_model\.head\."]
    _tied_weights_keys = {"head.decoder.weight": "model.text_model.embed_tokens.weight"}

    def __init__(self, config: UltraVBertConfig):
        super().__init__(config)
        self.vocab_size = config.text_config.vocab_size
        self.model = UltraVBertModel(config)
        self.head = UltrabertPredictionHead(config.text_config)
        self.post_init()

    def get_output_embeddings(self):
        return self.head.decoder

    def set_output_embeddings(self, new_embeddings):
        self.head.decoder = new_embeddings

    def forward(self, input_ids=None, attention_mask=None, position_ids=None, inputs_embeds=None, pixel_values=None,
                pixel_attention_mask=None, image_hidden_states=None, labels=None, **kwargs):
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
            pixel_values=pixel_values,
            pixel_attention_mask=pixel_attention_mask,
            image_hidden_states=image_hidden_states,
            **kwargs,
        )
        logits = self.head(outputs.last_hidden_state)
        loss = None
        if labels is not None:
            loss = CrossEntropyLoss()(logits.view(-1, self.vocab_size), labels.view(-1))
        return ModernVBertMaskedLMOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
            image_hidden_states=outputs.image_hidden_states,
        )


def _build_classifier(model, config: UltraVBertConfig):
    # ModernVBert's own __init__ builds a ModernBERT prediction head that UltraBERT's text config cannot configure.
    ModernVBertPreTrainedModel.__init__(model, config)
    model.num_labels = config.num_labels
    model.model = UltraVBertModel(config)
    model.head = nn.Identity()
    model.drop = nn.Dropout(config.classifier_dropout)
    model.classifier = nn.Linear(config.text_config.hidden_size, config.num_labels)
    model.post_init()


class UltraVBertForSequenceClassification(ModernVBertForSequenceClassification):
    config_class = UltraVBertConfig
    _keys_to_ignore_on_load_unexpected = [r"vision_model\.head\."]

    def __init__(self, config: UltraVBertConfig):
        _build_classifier(self, config)


class UltraVBertForTokenClassification(ModernVBertForTokenClassification):
    config_class = UltraVBertConfig
    _keys_to_ignore_on_load_unexpected = [r"vision_model\.head\."]

    def __init__(self, config: UltraVBertConfig):
        _build_classifier(self, config)


__all__ = [
    "UltraVBertModel",
    "UltraVBertForMaskedLM",
    "UltraVBertForSequenceClassification",
    "UltraVBertForTokenClassification",
]