# -*- coding: utf-8 -*-
# @Time    : 2021/8/19 10:54 上午
# @Author  : JianingWang
# @File    : classification.py
import torch
from torch import nn
from torch.nn import CrossEntropyLoss, MSELoss, BCEWithLogitsLoss
from transformers import RobertaModel
from transformers.activations import ACT2FN
from transformers.models.electra import ElectraModel
from transformers.models.roformer import RoFormerModel
from transformers.models.albert import AlbertModel
from transformers.models.bert import BertModel, BertPreTrainedModel
from transformers.models.deberta_v2 import DebertaV2Model, DebertaV2PreTrainedModel
from transformers.modeling_outputs import SequenceClassifierOutput
from transformers.models.roberta import RobertaPreTrainedModel
from transformers.models.bert.modeling_bert import BertForSequenceClassification
from transformers.models.megatron_bert import MegatronBertPreTrainedModel, MegatronBertModel

PRETRAINED_MODEL_MAP = {
    "bert": BertPreTrainedModel,
    "deberta-v2": DebertaV2PreTrainedModel,
    "roberta": RobertaPreTrainedModel,
    "erlangshen": MegatronBertPreTrainedModel
}


class BertPooler(nn.Module):
    def __init__(self, hidden_size, hidden_act, hidden_dropout_prob):
        super().__init__()
        self.dense = nn.Linear(hidden_size, hidden_size)
        # self.activation = nn.Tanh()
        self.activation = ACT2FN[hidden_act]
        # self.dropout = nn.Dropout(hidden_dropout_prob)

    def forward(self, features):
        x = features[:, 0, :]  # take <s> token (equiv. to [CLS])
        # x = self.dropout(x)
        x = self.dense(x)
        x = self.activation(x)
        return x


def build_cls_model(config):
    BaseClass = PRETRAINED_MODEL_MAP[config.model_type]

    class BertForClassification(BaseClass):

        def __init__(self, config):
            super().__init__(config)
            self.num_labels = config.num_labels
            self.config = config
            self.model_type = config.model_type
            self.problem_type = config.problem_type

            if self.model_type == "bert":
                self.bert = BertModel(config)
            elif self.model_type == "albert":
                self.albert = AlbertModel(config)
            # elif self.model_type == "chinesebert":
            #     self.bert = ChineseBertModel(config)
            elif self.model_type == "roformer":
                self.roformer = RoFormerModel(config)
            elif self.model_type == "electra":
                self.electra = ElectraModel(config)
            elif self.model_type == "deberta-v2":
                self.deberta = DebertaV2Model(config)
            elif self.model_type == "roberta":
                self.roberta = RobertaModel(config)
            elif self.model_type == "erlangshen":
                self.bert = MegatronBertModel(config)
            self.pooler = BertPooler(config.hidden_size, config.hidden_act, config.hidden_dropout_prob)
            if hasattr(config, "cls_dropout_rate"):
                cls_dropout_rate = config.cls_dropout_rate
            else:
                cls_dropout_rate = config.hidden_dropout_prob
            self.dropout = nn.Dropout(cls_dropout_rate)
            add_feature_dims = config.additional_feature_dims if hasattr(config, "additional_feature_dims") else 0
            # self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
            cls_hidden = config.hidden_size + add_feature_dims
            if hasattr(config, "is_relation_task"):
                cls_hidden = config.hidden_size * 2
            self.classifier = nn.Linear(cls_hidden, config.num_labels)

            self.init_weights()

        def forward(
                self,
                input_ids=None,
                attention_mask=None,
                token_type_ids=None,
                position_ids=None,
                head_mask=None,
                inputs_embeds=None,
                labels=None,
                output_attentions=None,
                output_hidden_states=None,
                return_dict=None,
                pseudo_label=None,
                pinyin_ids=None,
                additional_features=None
        ):
            return_dict = return_dict if return_dict is not None else self.config.use_return_dict
            logits, outputs = None, None
            inputs = {"input_ids": input_ids, "attention_mask": attention_mask, "token_type_ids": token_type_ids, "position_ids": position_ids,
                      "head_mask": head_mask, "inputs_embeds": inputs_embeds, "output_attentions": output_attentions,
                      "output_hidden_states": output_hidden_states, "return_dict": return_dict, "pinyin_ids": pinyin_ids}
            inputs = {k: v for k, v in inputs.items() if v is not None}
            if self.model_type == "chinesebert":
                outputs = self.bert(**inputs)
            elif self.model_type == "bert":
                outputs = self.bert(**inputs)
            elif self.model_type == "albert":
                outputs = self.albert(**inputs)
            elif self.model_type == "electra":
                outputs = self.electra(**inputs)
            elif self.model_type == "roformer":
                outputs = self.roformer(**inputs)
            elif self.model_type == "deberta-v2":
                outputs = self.deberta(**inputs)
            elif self.model_type == "roberta":
                outputs = self.roberta(**inputs)
            elif self.model_type == "erlangshen":
                outputs = self.bert(**inputs)

            if hasattr(self.config, "is_relation_task"):
                w = torch.logical_and(input_ids >= min(self.config.start_token_ids), input_ids <= max(self.config.start_token_ids))
                start_index = w.nonzero()[:, 1].view(-1, 2)
                # <start_entity> + <end_entity> 进分类
                pooler_output = torch.cat([torch.cat([x[y[0], :], x[y[1], :]]).unsqueeze(0) for x, y in zip(outputs.last_hidden_state, start_index)])
                # [CLS] + <start_entity> + <end_entity> 进分类
                # pooler_output = torch.cat([torch.cat([z, x[y[0], :], x[y[1], :]]).unsqueeze(0) for x, y, z in zip(outputs.last_hidden_state, start_index, outputs.last_hidden_state[:, 0])])

            elif "pooler_output" in outputs:
                pooler_output = outputs.pooler_output
            else:
                pooler_output = self.pooler(outputs[0])
            pooler_output = self.dropout(pooler_output)
            # pooler_output = self.LayerNorm(pooler_output)
            if additional_features is not None:
                pooler_output = torch.cat((pooler_output, additional_features), dim=1)
            logits = self.classifier(pooler_output)

            loss = None
            if labels is not None:
                if self.problem_type == "regression":
                    loss_fct = MSELoss()
                    if self.num_labels == 1:
                        loss = loss_fct(logits.squeeze(), labels.squeeze())
                    else:
                        loss = loss_fct(logits, labels)
                elif self.problem_type == "multi_label_classification":
                    loss_fct = BCEWithLogitsLoss()
                    loss = loss_fct(logits.view(-1, self.num_labels), labels.float().view(-1, self.num_labels))
                # elif self.problem_type in ["single_label_classification"] or hasattr(self.config, "is_relation_task"):
                else:
                    # loss_fct = FocalLoss()
                    loss_fct = CrossEntropyLoss()
                    # 伪标签
                    if pseudo_label is not None:
                        train_logits, pseudo_logits = logits[pseudo_label > 0.9], logits[pseudo_label < 0.1]
                        train_labels, pseudo_labels = labels[pseudo_label > 0.9], labels[pseudo_label < 0.1]
                        train_loss = loss_fct(train_logits.view(-1, self.num_labels), train_labels.view(-1)) if train_labels.nelement() else 0
                        pseudo_loss = loss_fct(pseudo_logits.view(-1, self.num_labels), pseudo_labels.view(-1)) if pseudo_labels.nelement() else 0
                        loss = 0.9 * train_loss + 0.1 * pseudo_loss
                    else:
                        loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            return SequenceClassifierOutput(
                loss=loss,
                logits=logits,
                hidden_states=outputs.hidden_states,
                attentions=outputs.attentions,
            )

    return BertForClassification