#from transformers import BertConfig, BertPreTrainedModel, BertTokenizer, BertModel

from transformers import AutoConfig

from transformers import AutoModelForPreTraining

from transformers import AutoTokenizer

from transformers import AutoModel

from transformers.modeling_utils import SequenceSummary

from torch import nn

import torch.nn.functional as F

import torch

from loss import AMSoftmax

from pytorch_metric_learning import losses, miners

class UMLSFinetuneModel(nn.Module):

    def __init__(self, device, model_name_or_path, cui_label_count, cui_loss_type="ms_loss"):

        super(UMLSFinetuneModel, self).__init__()

        self.device = device

        self.model_name_or_path = model_name_or_path

        self.config = AutoConfig.from_pretrained(model_name_or_path)

        self.bert = AutoModel.from_pretrained(self.model_name_or_path, config=self.config)

        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path)

        self.dropout = nn.Dropout(0.1)

        self.feature_dim = 768

        self.cui_loss_type = cui_loss_type

        self.cui_label_count = cui_label_count

        if self.cui_loss_type == "softmax":

            self.cui_loss_fn = nn.CrossEntropyLoss()

            self.linear = nn.Linear(self.feature_dim, self.cui_label_count)

        if self.cui_loss_type == "am_softmax":

            self.cui_loss_fn = AMSoftmax(self.feature_dim, self.cui_label_count)

        if self.cui_loss_type == "ms_loss":

            self.cui_loss_fn = losses.MultiSimilarityLoss(alpha=2, beta=50)

            self.miner = miners.MultiSimilarityMiner(epsilon=0.1)

    def softmax(self, logits, label):

        loss = self.cui_loss_fn(logits, label)

        return loss

    def am_softmax(self, pooled_output, label):

        loss, _ = self.cui_loss_fn(pooled_output, label)

        return loss

    def ms_loss(self, pooled_output, label):

        pairs = self.miner(pooled_output, label)

        loss = self.cui_loss_fn(pooled_output, label, pairs)

        return loss

    def calculate_loss(self, pooled_output=None, logits=None, label=None):

        if self.cui_loss_type == "softmax":

            return self.softmax(logits, label)

        if self.cui_loss_type == "am_softmax":

            return self.am_softmax(pooled_output, label)

        if self.cui_loss_type == "ms_loss":

            return self.ms_loss(pooled_output, label)    

    def get_sentence_feature(self, input_ids, attention_mask):

        outputs = self.bert(input_ids, attention_mask)

        pooled_output = outputs[1]

        return pooled_output

    def forward(self, input_ids, cui_label, attention_mask):

        pooled_output = self.get_sentence_feature(input_ids, attention_mask)

        if self.cui_loss_type == "softmax":

            logits = self.linear(pooled_output)

        else:

            logits = None

        cui_loss = self.calculate_loss(pooled_output, logits, cui_label)            

        loss = cui_loss

        return loss