from torch import nn
from transformers import AutoModel
from transformers.models.modernvbert.modeling_modernvbert import (
    ModernVBertConnector,
    ModernVBertModel,
    ModernVBertPreTrainedModel,
)
from transformers.models.ultrabert.modeling_ultrabert import 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()


__all__ = ["UltraVBertModel"]