import torch
import torch.nn as nn

from transformers import (
    AutoModel,
    PreTrainedModel
)
from transformers.modeling_outputs import SequenceClassifierOutput

from .configuration_veritas import BERTClassifierConfig

from transformers import AutoModel, AutoConfig

class BERTClassifier(PreTrainedModel):
    config_class = BERTClassifierConfig

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

        bert_config = AutoConfig.from_pretrained(
            "neuralmind/bert-base-portuguese-cased"
        )

        self.bert = AutoModel.from_config(bert_config)

        for param in self.bert.encoder.layer[:6].parameters():
            param.requires_grad = False

        hidden = self.bert.config.hidden_size

        self.classifier = nn.Sequential(
            nn.Linear(hidden, 32),
            nn.GELU(),
            nn.Dropout(config.drop_rate),
            nn.Linear(32, 16),
            nn.GELU(),
            nn.Linear(16, config.num_labels)
        )

        self.post_init()

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        labels=None
    ):
        outputs = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask
        )

        pooler = outputs.last_hidden_state[:, 0, :]
        logits = self.classifier(pooler)

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

        return SequenceClassifierOutput(
            loss=loss,
            logits=logits
        )