# Copyright (c) Alibaba, Inc. and its affiliates.

import os
from typing import Dict, Union

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn.init import xavier_uniform_
from transformers import BertModel

from modelscope.metainfo import Models
from modelscope.models import TorchModel
from modelscope.models.base import Tensor
from modelscope.models.builder import MODELS
from modelscope.outputs import DialogueUserSatisfactionEstimationModelOutput
from modelscope.utils.constant import ModelFile, Tasks
from .transformer import TransformerEncoder

__all__ = ['UserSatisfactionEstimation']


@MODELS.register_module(Tasks.text_classification, module_name=Models.use)
class UserSatisfactionEstimation(TorchModel):

    def __init__(self,
                 model_dir: str,
                 bert_name: str = None,
                 device: str = None):
        """initialize the user satisfaction estimation model from the `model_dir` path. The default preprocessor
        for this task is DialogueClassificationUsePreprocessor.

        Args:
            model_dir: The model dir containing the model.
            bert_name: The pretrained model, default bert-base-chinese
            device: The device of running model, default cpu
        """
        super().__init__(model_dir)
        self.model_dir = model_dir
        self.bert_name = bert_name if bert_name is not None else 'bert-base-chinese'
        self.device = 'cpu'
        if device is not None and torch.cuda.is_available():
            self.device = device
        self.model = self.init_model()
        model_ckpt = os.path.join(model_dir, ModelFile.TORCH_MODEL_FILE)
        self.model.load_state_dict(
            torch.load(model_ckpt, map_location=torch.device('cpu')))

    def init_model(self):
        configs = {
            'bert_name': self.bert_name,
            'cache_dir': self.model_dir,
            'dropout': 0.1
        }
        model = USE(configs)
        return model

    def forward(
        self, input_ids: Tensor
    ) -> Union[DialogueUserSatisfactionEstimationModelOutput, Dict[str,
                                                                   Tensor]]:
        """Compute the logits of satisfaction polarities for a dialogue.

        Args:
           input_ids (Tensor): the preprocessed dialogue input
        Returns:
           output (Dict[str, Any] or DialogueUserSatisfactionEstimationModelOutput): The results of user satisfaction.

        Example:
            >>> {'logits': tensor([[-2.1795,  1.1323,  1.8605]])}
        """
        logits = self.model(input_ids)
        return DialogueUserSatisfactionEstimationModelOutput(logits=logits)


def init_params(model):
    for name, param in model.named_parameters():
        if param.data.dim() > 1:
            xavier_uniform_(param.data)
        else:
            pass


def universal_sentence_embedding(sentences, mask, sqrt=True):
    sentence_sums = torch.bmm(
        sentences.permute(0, 2, 1),
        mask.float().unsqueeze(-1)).squeeze(-1)
    divisor = (mask.sum(dim=1).view(-1, 1).float())
    if sqrt:
        divisor = divisor.sqrt()
    sentence_sums /= divisor
    return sentence_sums


class BERTBackbone(nn.Module):

    def __init__(self, **config):
        super().__init__()
        bert_name = config.get('bert_name', 'bert-base-chinese')
        cache_dir = config.get('cache_dir')
        self.bert = BertModel.from_pretrained(bert_name, cache_dir=cache_dir)
        self.d_model = 768 * 2

    def forward(self, input_ids):
        attention_mask = input_ids.ne(0).detach()
        outputs = self.bert(input_ids, attention_mask)
        h = universal_sentence_embedding(outputs[0], attention_mask)
        cls = outputs[1]
        out = torch.cat([cls, h], dim=-1)
        return out


class MLP(nn.Module):

    def __init__(self, input_size, output_size, hidden_size):
        super().__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.fc2 = nn.Linear(hidden_size, output_size)

    def forward(self, din):
        dout = F.relu(self.fc1(din))
        dout = F.relu(self.fc2(dout))
        return dout


class USE(nn.Module):

    def __init__(self, args):
        super().__init__()
        self.drop_out = nn.Dropout(args['dropout'])
        self.private = BERTBackbone(
            bert_name=args['bert_name'], cache_dir=args['cache_dir'])
        d_model = self.private.d_model
        self.encoder = TransformerEncoder(d_model, d_model * 2, 8, 2, 0.1)
        self.content_gru = nn.GRU(
            d_model,
            d_model,
            num_layers=1,
            bidirectional=False,
            batch_first=True)
        self.sat_classifier = nn.Linear(d_model, 3)

        self.U_c = nn.Linear(d_model, d_model)
        self.w_c = nn.Linear(d_model, 1, bias=False)

        init_params(self.encoder)
        init_params(self.sat_classifier)
        init_params(self.U_c)
        init_params(self.w_c)

    def forward(self, input_ids):
        self.content_gru.flatten_parameters()
        batch_size, dialog_len, utt_len = input_ids.size()
        attention_mask = input_ids[:, :, 0].squeeze(-1).ne(0).detach()
        input_ids = input_ids.view(-1, utt_len)

        private_out = self.private(input_ids=input_ids)
        private_out = private_out.view(batch_size, dialog_len, -1)
        H = self.encoder(private_out, attention_mask)
        H = self.drop_out(H)

        H, _ = self.content_gru(H)
        att_c = self.w_c(torch.tanh(self.U_c(H))).squeeze(-1)
        att_c = F.softmax(
            att_c.masked_fill(mask=~attention_mask, value=-np.inf), dim=1)
        hidden = torch.bmm(H.permute(0, 2, 1), att_c.unsqueeze(-1)).squeeze(-1)

        sat_res = self.sat_classifier(hidden)
        return sat_res
