"""HF-style StrataBERT model classes."""

from __future__ import annotations

import torch
from torch import nn

from .attention import StrataBertAttentionLayer
from .bidirectional_ssm import BidirectionalSSMLayer
from .configuration_stratabert import StrataBertConfig
from .heads import StrataBertClassificationHead, StrataBertLMHead, StrataBertRTDHead, StrataBertTokenClassificationHead
from .losses import masked_lm_loss, replaced_token_detection_loss, sequence_classification_loss, token_classification_loss
from .modeling_outputs import (
    StrataBertMaskedLMOutput,
    StrataBertModelOutput,
    StrataBertPreTrainingOutput,
    StrataBertSequenceClassifierOutput,
    StrataBertTokenClassifierOutput,
)
from .padding import make_attention_mask, masked_hidden
from .pooling import StrataBertPooler

try:
    from transformers import PreTrainedModel
except ModuleNotFoundError:

    class PreTrainedModel(nn.Module):  # type: ignore[no-redef]
        config_class = None

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

        def post_init(self):
            return None


class StrataBertPreTrainedModel(PreTrainedModel):
    config_class = StrataBertConfig
    base_model_prefix = "stratabert"
    supports_gradient_checkpointing = True

    def _init_weights(self, module):
        if isinstance(module, nn.Linear):
            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                nn.init.zeros_(module.bias)
        elif isinstance(module, nn.Embedding):
            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                nn.init.zeros_(module.weight[module.padding_idx])


class StrataBertEmbeddings(nn.Module):
    def __init__(self, config: StrataBertConfig):
        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.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
        self.norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout)

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor,
        token_type_ids: torch.Tensor | None = None,
        position_ids: torch.Tensor | None = None,
    ) -> torch.Tensor:
        batch, length = input_ids.shape
        if position_ids is None:
            position_ids = torch.arange(length, device=input_ids.device).unsqueeze(0).expand(batch, -1)
        if token_type_ids is None:
            token_type_ids = torch.zeros_like(input_ids)
        x = (
            self.word_embeddings(input_ids)
            + self.position_embeddings(position_ids)
            + self.token_type_embeddings(token_type_ids)
        )
        return masked_hidden(self.dropout(self.norm(x)), attention_mask)


class StrataBertEncoder(nn.Module):
    def __init__(self, config: StrataBertConfig):
        super().__init__()
        layers = []
        for layer_type in config.layer_types:
            if layer_type == "ssm":
                layers.append(BidirectionalSSMLayer(config))
            else:
                layers.append(StrataBertAttentionLayer(config, layer_type))
        self.layers = nn.ModuleList(layers)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: torch.Tensor,
        segment_ids: torch.Tensor | None = None,
        output_hidden_states: bool = False,
    ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...] | None]:
        all_hidden = [] if output_hidden_states else None
        for layer in self.layers:
            if all_hidden is not None:
                all_hidden.append(hidden_states)
            if isinstance(layer, BidirectionalSSMLayer):
                hidden_states = layer(hidden_states, attention_mask, segment_ids)
            else:
                hidden_states = layer(hidden_states, attention_mask, segment_ids)
        if all_hidden is not None:
            all_hidden.append(hidden_states)
        return hidden_states, tuple(all_hidden) if all_hidden is not None else None


class StrataBertModel(StrataBertPreTrainedModel):
    def __init__(self, config: StrataBertConfig):
        super().__init__(config)
        self.embeddings = StrataBertEmbeddings(config)
        self.encoder = StrataBertEncoder(config)
        self.pooler = StrataBertPooler(config.hidden_size, config.pooling_type)
        self.post_init()

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        token_type_ids: torch.Tensor | None = None,
        position_ids: torch.Tensor | None = None,
        segment_ids: torch.Tensor | None = None,
        output_hidden_states: bool | None = None,
        **kwargs,
    ) -> StrataBertModelOutput:
        del kwargs
        attention_mask = make_attention_mask(input_ids, attention_mask, self.config.pad_token_id)
        output_hidden_states = self.config.output_hidden_states if output_hidden_states is None else output_hidden_states
        hidden_states = self.embeddings(input_ids, attention_mask, token_type_ids, position_ids)
        hidden_states, all_hidden = self.encoder(hidden_states, attention_mask, segment_ids, output_hidden_states)
        pooled = self.pooler(hidden_states, attention_mask)
        return StrataBertModelOutput(
            last_hidden_state=hidden_states,
            pooler_output=pooled,
            hidden_states=all_hidden,
            attentions=None,
            ssm_states=None,
        )


class StrataBertForSequenceClassification(StrataBertPreTrainedModel):
    def __init__(self, config: StrataBertConfig):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.stratabert = StrataBertModel(config)
        self.classifier = StrataBertClassificationHead(config)
        self.post_init()

    def forward(self, input_ids, attention_mask=None, token_type_ids=None, position_ids=None, segment_ids=None, labels=None, **kwargs):
        outputs = self.stratabert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            segment_ids=segment_ids,
            **kwargs,
        )
        logits = self.classifier(outputs.pooler_output)
        loss = None
        if labels is not None:
            loss = sequence_classification_loss(logits, labels, self.num_labels, self.config.problem_type)
        return StrataBertSequenceClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states)


class StrataBertForTokenClassification(StrataBertPreTrainedModel):
    def __init__(self, config: StrataBertConfig):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.stratabert = StrataBertModel(config)
        self.classifier = StrataBertTokenClassificationHead(config)
        self.post_init()

    def forward(self, input_ids, attention_mask=None, token_type_ids=None, position_ids=None, segment_ids=None, labels=None, **kwargs):
        outputs = self.stratabert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            segment_ids=segment_ids,
            **kwargs,
        )
        logits = self.classifier(outputs.last_hidden_state)
        loss = None
        if labels is not None:
            mask = make_attention_mask(input_ids, attention_mask, self.config.pad_token_id)
            loss = token_classification_loss(logits, labels, mask, self.num_labels)
        return StrataBertTokenClassifierOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states)


class StrataBertForMaskedLM(StrataBertPreTrainedModel):
    _tied_weights_keys = {
        "lm_head.decoder.weight": "stratabert.embeddings.word_embeddings.weight",
        "lm_head.decoder.bias": "lm_head.bias",
    }

    def __init__(self, config: StrataBertConfig):
        super().__init__(config)
        self.stratabert = StrataBertModel(config)
        self.lm_head = StrataBertLMHead(config, self.stratabert.embeddings.word_embeddings.weight)
        self.post_init()

    def forward(self, input_ids, attention_mask=None, token_type_ids=None, position_ids=None, segment_ids=None, labels=None, **kwargs):
        outputs = self.stratabert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            segment_ids=segment_ids,
            **kwargs,
        )
        logits = self.lm_head(outputs.last_hidden_state)
        loss = masked_lm_loss(logits, labels) if labels is not None else None
        return StrataBertMaskedLMOutput(loss=loss, logits=logits, hidden_states=outputs.hidden_states)


class StrataBertForPreTraining(StrataBertPreTrainedModel):
    _tied_weights_keys = {
        "lm_head.decoder.weight": "stratabert.embeddings.word_embeddings.weight",
        "lm_head.decoder.bias": "lm_head.bias",
    }

    def __init__(self, config: StrataBertConfig):
        super().__init__(config)
        self.stratabert = StrataBertModel(config)
        self.lm_head = StrataBertLMHead(config, self.stratabert.embeddings.word_embeddings.weight)
        self.rtd_head = StrataBertRTDHead(config)
        self.post_init()

    def forward(
        self,
        input_ids,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        segment_ids=None,
        labels=None,
        rtd_labels=None,
        mlm_weight: float = 1.0,
        rtd_weight: float = 0.25,
        **kwargs,
    ):
        outputs = self.stratabert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            segment_ids=segment_ids,
            **kwargs,
        )
        prediction_logits = self.lm_head(outputs.last_hidden_state)
        rtd_logits = self.rtd_head(outputs.last_hidden_state)
        mlm = masked_lm_loss(prediction_logits, labels) if labels is not None else None
        rtd = None
        if rtd_labels is not None:
            mask = make_attention_mask(input_ids, attention_mask, self.config.pad_token_id)
            rtd = replaced_token_detection_loss(rtd_logits, rtd_labels, mask)
        loss = None
        if mlm is not None and rtd is not None:
            loss = mlm_weight * mlm + rtd_weight * rtd
        elif mlm is not None:
            loss = mlm
        elif rtd is not None:
            loss = rtd
        return StrataBertPreTrainingOutput(
            loss=loss,
            prediction_logits=prediction_logits,
            rtd_logits=rtd_logits,
            mlm_loss=mlm,
            rtd_loss=rtd,
            hidden_states=outputs.hidden_states,
        )