import logging

import torch
import torch.nn as nn

from transformers import AutoModel

from utils.nn_utils import gelu
from modules.token_embedders.bert_encoder import BertLinear

logger = logging.getLogger(__name__)


class PretrainedEncoder(nn.Module):
    """This class using pre-trained model to encode token,
    then fine-tuning the pre-trained model
    """
    def __init__(self, pretrained_model_name, trainable=False, output_size=0, activation=gelu, dropout=0.0):
        """This function initialize pertrained model

        Arguments:
            pretrained_model_name {str} -- pre-trained model name

        Keyword Arguments:
            output_size {float} -- output size (default: {None})
            activation {nn.Module} -- activation function (default: {gelu})
            dropout {float} -- dropout rate (default: {0.0})
        """

        super().__init__()
        self.pretrained_model = AutoModel.from_pretrained(pretrained_model_name)
        logger.info("Load pre-trained model {} successfully.".format(pretrained_model_name))

        self.output_size = output_size

        if trainable:
            logger.info("Start fine-tuning pre-trained model {}.".format(pretrained_model_name))
        else:
            logger.info("Keep fixed pre-trained model {}.".format(pretrained_model_name))

        for param in self.pretrained_model.parameters():
            param.requires_grad = trainable

        if self.output_size > 0:
            self.mlp = BertLinear(input_size=self.pretrained_model.config.hidden_size,
                                  output_size=self.output_size,
                                  activation=activation)
        else:
            self.output_size = self.pretrained_model.config.hidden_size
            self.mlp = lambda x: x

        if dropout > 0:
            self.dropout = nn.Dropout(p=dropout)
        else:
            self.dropout = lambda x: x

    def get_output_dims(self):
        return self.output_size

    def forward(self, seq_inputs, token_type_inputs=None):
        """forward calculates forward propagation results, get token embedding

        Args:
            seq_inputs {tensor} -- sequence inputs (tokenized)
            token_type_inputs (tensor, optional): token type inputs. Defaults to None.

        Returns:
            tensor: bert output for tokens
        """

        if token_type_inputs is None:
            token_type_inputs = torch.zeros_like(seq_inputs)
        mask_inputs = (seq_inputs != 0).long()

        outputs = self.pretrained_model(input_ids=seq_inputs,
                                        token_type_ids=token_type_inputs,
                                        attention_mask=mask_inputs)
        last_hidden_state = outputs[0]
        pooled_output = outputs[1]

        return self.dropout(self.mlp(last_hidden_state)), self.dropout(self.mlp(pooled_output))