#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

from trans import TransE

class UMLSPretrainedModel(nn.Module):

    def __init__(self, device, model_name_or_path,

                 cui_label_count, rel_label_count, sty_label_count,

                 re_weight=1.0, sty_weight=0.1,

                 cui_loss_type="ms_loss",

                 trans_loss_type="TransE", trans_margin=1.0):

        super(UMLSPretrainedModel, self).__init__()

        self.device = device

        self.model_name_or_path = model_name_or_path

        if self.model_name_or_path.find("large") >= 0:

            self.feature_dim = 1024

        else:

            self.feature_dim = 768

        self.bert = AutoModel.from_pretrained(model_name_or_path)

        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)

        self.dropout = nn.Dropout(0.1)

        self.rel_label_count = rel_label_count

        self.re_weight = re_weight

        self.sty_label_count = sty_label_count

        self.linear_sty = nn.Linear(self.feature_dim, self.sty_label_count)

        self.sty_loss_fn = nn.CrossEntropyLoss()

        self.sty_weight = sty_weight

        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)

        self.trans_loss_type = trans_loss_type

        if self.trans_loss_type == "TransE":

            self.re_loss_fn = TransE(trans_margin)

        self.re_embedding = nn.Embedding(

            self.rel_label_count, self.feature_dim)

        self.standard_dataloader = None

        self.sequence_summary = SequenceSummary(AutoConfig.from_pretrained(model_name_or_path)) # Now only used for XLNet

    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):

        # bert, albert, roberta

        if self.model_name_or_path.find("xlnet") < 0:

            outputs = self.bert(input_ids)

            pooled_output = outputs[1]

            return pooled_output

        # xlnet

        outputs = self.bert(input_ids)

        pooled_output = self.sequence_summary(outputs[0])

        return pooled_output

    # @profile

    def forward(self,

                input_ids_0, input_ids_1, input_ids_2,

                cui_label_0, cui_label_1, cui_label_2,

                sty_label_0, sty_label_1, sty_label_2,

                re_label):

        input_ids = torch.cat((input_ids_0, input_ids_1, input_ids_2), 0)

        cui_label = torch.cat((cui_label_0, cui_label_1, cui_label_2))

        sty_label = torch.cat((sty_label_0, sty_label_1, sty_label_2))

        #print(input_ids.shape, cui_label.shape, sty_label.shape)

        use_len = input_ids_0.shape[0]

        pooled_output = self.get_sentence_feature(

            input_ids)  # (3 * pair) * re_label

        logits_sty = self.linear_sty(pooled_output)

        sty_loss = self.sty_loss_fn(logits_sty, sty_label)

        if self.cui_loss_type == "softmax":

            logits = self.linear(pooled_output)

        else:

            logits = None

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

        cui_0_output = pooled_output[0:use_len]

        cui_1_output = pooled_output[use_len:2 * use_len]

        cui_2_output = pooled_output[2 * use_len:]

        re_output = self.re_embedding(re_label)

        re_loss = self.re_loss_fn(

            cui_0_output, cui_1_output, cui_2_output, re_output)

        loss = self.sty_weight * sty_loss + cui_loss + self.re_weight * re_loss

        #print(sty_loss.device, cui_loss.device, re_loss.device)

        return loss, (sty_loss, cui_loss, re_loss)

    """

    def predict(self, input_ids):

        if self.loss_type == "softmax":

            return self.predict_by_softmax(input_ids)

        if self.loss_type == "am_softmax":

            return self.predict_by_amsoftmax(input_ids)        

    def predict_by_softmax(self, input_ids):

        pooled_output = self.get_sentence_feature(input_ids)

        logits = self.linear(pooled_output)

        return torch.max(logits, dim=1)[1], logits

    def predict_by_amsoftmax(self, input_ids):

        pooled_output = self.get_sentence_feature(input_ids)

        logits = self.loss_fn.predict(pooled_output)

        return torch.max(logits, dim=1)[1], logits

    """

    def init_standard_feature(self):

        if self.standard_dataloader is not None:

            for index, batch in enumerate(self.standard_dataloader):

                input_ids = batch[0].to(self.device)

                outputs = self.get_sentence_feature(input_ids)

                normalized_standard_feature = torch.norm(

                    outputs, p=2, dim=1, keepdim=True).clamp(min=1e-12)

                normalized_standard_feature = torch.div(

                    outputs, normalized_standard_feature)

                if index == 0:

                    self.standard_feature = normalized_standard_feature

                else:

                    self.standard_feature = torch.cat(

                        (self.standard_feature, normalized_standard_feature), 0)

            assert self.standard_feature.shape == (

                self.num_label, self.feature_dim), self.standard_feature.shape

        return None

    def predict_by_cosine(self, input_ids):

        pooled_output = self.get_sentence_feature(input_ids)

        normalized_feature = torch.norm(

            pooled_output, p=2, dim=1, keepdim=True).clamp(min=1e-12)

        normalized_feature = torch.div(pooled_output, normalized_feature)

        sim_mat = torch.matmul(normalized_feature, torch.t(

            self.standard_feature))  # batch_size * num_label

        return torch.max(sim_mat, dim=1)[1], sim_mat