"""
Custom HuggingFace-compatible wrapper for the DeBERTa-v3-small MCQ scorer
trained in the Kaggle notebook (DeBERTaMCQModel).

This makes the model loadable on the Hub via:
    AutoConfig.from_pretrained(repo_id, trust_remote_code=True)
    AutoModel.from_pretrained(repo_id, trust_remote_code=True)
"""

import torch
import torch.nn as nn
from transformers import PretrainedConfig, PreTrainedModel, AutoModel


class DebertaMCQConfig(PretrainedConfig):
    model_type = "deberta_mcq"

    def __init__(
        self,
        base_model_name="microsoft/deberta-v3-small",
        hidden_dropout=0.3,
        num_options=5,
        max_len=64,
        **kwargs,
    ):
        self.base_model_name = base_model_name
        self.hidden_dropout = hidden_dropout
        self.num_options = num_options
        self.max_len = max_len
        super().__init__(**kwargs)


class DebertaMCQForMultipleChoice(PreTrainedModel):
    config_class = DebertaMCQConfig

    def __init__(self, config: DebertaMCQConfig):
        super().__init__(config)
        self.deberta = AutoModel.from_pretrained(
            config.base_model_name,
            torch_dtype=torch.float32,
            low_cpu_mem_usage=False,
            device_map=None,
        )
        hidden_size = self.deberta.config.hidden_size
        self.classifier = nn.Sequential(
            nn.Dropout(config.hidden_dropout),
            nn.Linear(hidden_size, 1),
        )
        self.post_init()

    def forward(self, input_ids, attention_mask=None, labels=None):
        # input_ids / attention_mask shape: (batch, num_options, seq_len)
        batch_size, num_options, seq_len = input_ids.shape
        input_ids = input_ids.view(-1, seq_len)
        attention_mask = attention_mask.view(-1, seq_len)

        outputs = self.deberta(input_ids=input_ids, attention_mask=attention_mask)
        cls_output = outputs.last_hidden_state[:, 0, :].float()
        scores = self.classifier(cls_output).view(batch_size, num_options)

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

        return {"loss": loss, "logits": scores} if loss is not None else {"logits": scores}