"""Exact Hugging Face-compatible implementation of BashkirRoBERTa Pre-LN."""

import torch
from torch import nn
from transformers import PreTrainedModel
from transformers.modeling_outputs import MaskedLMOutput

try:
    from .configuration_bashkir_roberta import BashkirRobertaConfig
except (ImportError, ValueError):
    from configuration_bashkir_roberta import BashkirRobertaConfig


class BashkirRobertaEmbeddings(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.word_embeddings = nn.Embedding(
            config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id
        )
        self.position_embeddings = nn.Embedding(
            config.max_position_embeddings, config.hidden_size
        )
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=1e-12)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, input_ids):
        sequence_length = input_ids.shape[1]
        positions = torch.arange(
            sequence_length, dtype=torch.long, device=input_ids.device
        ).unsqueeze(0)
        hidden_states = self.word_embeddings(input_ids) + self.position_embeddings(positions)
        return self.dropout(self.LayerNorm(hidden_states))


class BashkirRobertaLayer(nn.Module):
    """Pre-LayerNorm transformer block, matching ``train_flagship.py`` exactly."""

    def __init__(self, config):
        super().__init__()
        self.self_attn = nn.MultiheadAttention(
            config.hidden_size,
            config.num_attention_heads,
            dropout=config.hidden_dropout_prob,
            batch_first=True,
        )
        self.norm1 = nn.LayerNorm(config.hidden_size, eps=1e-12)
        self.ffn = nn.Sequential(
            nn.Linear(config.hidden_size, config.intermediate_size),
            nn.GELU(),
            nn.Dropout(config.hidden_dropout_prob),
            nn.Linear(config.intermediate_size, config.hidden_size),
            nn.Dropout(config.hidden_dropout_prob),
        )
        self.norm2 = nn.LayerNorm(config.hidden_size, eps=1e-12)

    def forward(self, hidden_states, padding_mask=None):
        normalized = self.norm1(hidden_states)
        attention, _ = self.self_attn(
            normalized, normalized, normalized, key_padding_mask=padding_mask
        )
        hidden_states = hidden_states + attention
        return hidden_states + self.ffn(self.norm2(hidden_states))


class BashkirRobertaForMaskedLM(PreTrainedModel):
    config_class = BashkirRobertaConfig
    base_model_prefix = "bashkir_roberta"
    main_input_name = "input_ids"
    _tied_weights_keys = {"lm_head.weight": "embeddings.word_embeddings.weight"}

    def __init__(self, config):
        super().__init__(config)
        self.embeddings = BashkirRobertaEmbeddings(config)
        self.layers = nn.ModuleList(
            [BashkirRobertaLayer(config) for _ in range(config.num_hidden_layers)]
        )
        self.norm_final = nn.LayerNorm(config.hidden_size, eps=1e-12)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        self.post_init()

    def get_input_embeddings(self):
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    def get_output_embeddings(self):
        return self.lm_head

    def set_output_embeddings(self, new_embeddings):
        self.lm_head = new_embeddings

    def tie_weights(self, **kwargs):
        if self.config.tie_word_embeddings:
            self.lm_head.weight = self.embeddings.word_embeddings.weight

    def forward(self, input_ids, attention_mask=None, labels=None, return_dict=True, **kwargs):
        if input_ids.shape[1] > self.config.max_position_embeddings:
            raise ValueError(
                f"Sequence length {input_ids.shape[1]} exceeds "
                f"max_position_embeddings={self.config.max_position_embeddings}."
            )
        padding_mask = input_ids.eq(self.config.pad_token_id) if attention_mask is None else attention_mask.eq(0)
        hidden_states = self.embeddings(input_ids)
        for layer in self.layers:
            hidden_states = layer(hidden_states, padding_mask=padding_mask)
        hidden_states = self.norm_final(hidden_states)
        logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            loss = nn.functional.cross_entropy(
                logits.view(-1, self.config.vocab_size), labels.view(-1), ignore_index=-100
            )
        if not return_dict:
            return (loss, logits) if loss is not None else (logits,)
        return MaskedLMOutput(loss=loss, logits=logits)