import inspect

import torch.nn as nn
from transformers import AutoConfig, AutoModel, PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput

from custom_text_config import CustomTextConfig


class CustomTextForSequenceClassification(PreTrainedModel):
    config_class = CustomTextConfig

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

        if not getattr(config, "_name_or_path", None) and not getattr(
            config, "_is_meta", False
        ):
            self.backbone = AutoModel.from_pretrained(config.backbone_name_or_path)
        else:
            backbone_config = AutoConfig.from_pretrained(config.backbone_name_or_path)
            self.backbone = AutoModel.from_config(backbone_config)

        hidden_size = self.backbone.config.hidden_size
        self.classifier = nn.Linear(hidden_size, config.num_labels)
        self.post_init()

    def forward(
        self,
        input_ids,
        attention_mask=None,
        token_type_ids=None,
        labels=None,
        **kwargs,
    ):
        backbone_inputs = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "token_type_ids": token_type_ids,
            **kwargs,
        }
        accepted_args = set(inspect.signature(self.backbone.forward).parameters)
        backbone_inputs = {
            key: value
            for key, value in backbone_inputs.items()
            if value is not None and key in accepted_args
        }

        outputs = self.backbone(**backbone_inputs)
        cls_output = outputs.last_hidden_state[:, 0, :]
        logits = self.classifier(cls_output)

        loss = None
        if labels is not None:
            loss = nn.CrossEntropyLoss()(logits, labels)

        return SequenceClassifierOutput(loss=loss, logits=logits)


CustomTextForSequenceClassification.register_for_auto_class(
    "AutoModelForSequenceClassification"
)