import copy

from transformers import PreTrainedModel, PretrainedConfig, AutoModel, AutoTokenizer
import torch.nn as nn
import torch
from typing import List


class AMC_OPT_conf(PretrainedConfig):
    model_type = "bert"
    def __init__(self,
                 out_labels=3,
                 emb_size=768,
                 drop_out=0.1,
                 pretrained_model="m-polignano-uniba/bert_uncased_L-12_H-768_A-12_italian_alb3rt0",
                 **kwargs):
        self.out_labels = out_labels
        self.drop_out = drop_out
        self.emb_size = emb_size
        self.pretrained_model = pretrained_model
        self.fine_tuned_model = "data/models/amc_opt_msmd_nocoadapt.pt"
        super().__init__(**kwargs)


class AMC_OPT_sub(nn.Module):
    def __init__(self, config):
        super(AMC_OPT_sub, self).__init__()

        self.model = copy.deepcopy(AutoModel.from_pretrained(config.pretrained_model))
        self.dropout1 = nn.Dropout(config.drop_out)
        self.linear1 = nn.Linear(config.emb_size, config.out_labels)

        self.loss_fct = nn.CrossEntropyLoss()
        self.hyper_params = {'learning_rate': 6.599917952321265e-05,
                             'weight_decay': 0.02157165894420757,
                             'warmup_steps': 0.8999999999999999,
                             'num_epochs': 11}
        self.params_name = "alberto_multiclass_opt_msmd.pt"

    def forward(self, labels, input_ids, attention_mask, **args):
        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask, **args)
        x = self.dropout1(outputs[1])
        logits = self.linear1(x)
        if labels != None:
            loss = self.loss_fct(logits, labels)
            return {'logits':logits, 'loss':loss}
        else:
            return {'logits':logits}

class AMC_OPT(PreTrainedModel):
    config_class = AMC_OPT_conf

    def __init__(self, config):
        super().__init__(config)
        self.model = AMC_OPT_sub(config)
        self.model.load_state_dict(torch.load(config.fine_tuned_model))

    def forward(self, labels, input_ids, attention_mask, **args):
        return self.model(labels, input_ids, attention_mask, **args)

tokenizer = AutoTokenizer.from_pretrained("m-polignano-uniba/bert_uncased_L-12_H-768_A-12_italian_alb3rt0")
tokenizer.model_max_length = 128