# -*- coding: utf-8 -*-
import copy
import math

import torch
import transformers


class QueryBertLikeSelfAttention(torch.nn.Module):
    def __init__(
        self,
        origin: transformers.models.bert.modeling_bert.BertSelfAttention,
        share_weights: bool = False,
    ):
        super().__init__()
        self.num_attention_heads = origin.num_attention_heads
        self.attention_head_size = origin.attention_head_size
        self.all_head_size = origin.all_head_size

        if share_weights:
            self.query = origin.query
            self.key = origin.key
            self.value = origin.value
            self.dropout = origin.dropout
        else:
            # https://discuss.pytorch.org/t/are-there-any-recommended-methods-to-clone-a-model/483
            self.query = copy.deepcopy(origin.query)
            self.key = copy.deepcopy(origin.key)
            self.value = copy.deepcopy(origin.value)
            self.dropout = copy.deepcopy(origin.dropout)

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (
            self.num_attention_heads,
            self.attention_head_size,
        )
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)

    def forward(
        self,
        query_states,
        hidden_states,
        attention_mask=None,
        head_mask=None,
        output_attentions=False,
    ):
        query_layer = self.transpose_for_scores(self.query(query_states))
        key_layer = self.transpose_for_scores(self.key(hidden_states))
        value_layer = self.transpose_for_scores(self.value(hidden_states))

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
            attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = torch.nn.Softmax(dim=-1)(attention_scores)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = torch.matmul(attention_probs, value_layer)

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)

        outputs = (
            (context_layer, attention_probs) if output_attentions else (context_layer,)
        )
        return outputs


class QueryBertLikeAttention(torch.nn.Module):
    def __init__(
        self,
        origin: transformers.models.bert.modeling_bert.BertAttention,
        share_weights: bool = False,
    ):
        super().__init__()
        self.self = QueryBertLikeSelfAttention(origin.self, share_weights=share_weights)
        if share_weights:
            self.output = origin.output
        else:
            self.output = copy.deepcopy(origin.output)

    def forward(self, query_states, hidden_states, **kwargs):
        self_outputs = self.self(query_states, hidden_states, **kwargs)
        attention_output = self.output(self_outputs[0], query_states)
        outputs = (attention_output,) + self_outputs[
            1:
        ]  # add attentions if we output them
        return outputs


class QueryBertLikeLayer(torch.nn.Module):
    def __init__(
        self,
        origin: transformers.models.bert.modeling_bert.BertLayer,
        share_weights: bool = False,
    ):
        super().__init__()
        self.attention = QueryBertLikeAttention(
            origin.attention, share_weights=share_weights
        )
        if share_weights:
            self.intermediate = origin.intermediate
            self.output = origin.output
        else:
            self.intermediate = copy.deepcopy(origin.intermediate)
            self.output = copy.deepcopy(origin.output)

    def forward(self, query_states, hidden_states, **kwargs):
        self_attention_outputs = self.attention(query_states, hidden_states, **kwargs)
        attention_output = self_attention_outputs[0]

        outputs = self_attention_outputs[
            1:
        ]  # add self attentions if we output attention weights

        intermediate_output = self.intermediate(attention_output)
        layer_output = self.output(intermediate_output, attention_output)
        outputs = (layer_output,) + outputs

        return outputs


class QueryAlbertSelfAttention(QueryBertLikeSelfAttention):
    def __init__(
        self,
        origin: transformers.models.albert.modeling_albert.AlbertAttention,
        share_weights: bool = False,
    ):
        torch.nn.Module.__init__(self)
        self.num_attention_heads = origin.num_attention_heads
        self.attention_head_size = origin.attention_head_size
        self.all_head_size = origin.all_head_size

        if share_weights:
            self.query = origin.query
            self.key = origin.key
            self.value = origin.value
            self.dropout = origin.attention_dropout
        else:
            self.query = copy.deepcopy(origin.query)
            self.key = copy.deepcopy(origin.key)
            self.value = copy.deepcopy(origin.value)
            self.dropout = copy.deepcopy(origin.attention_dropout)


class AlbertSelfOutput(transformers.models.bert.modeling_bert.BertOutput):
    def __init__(
        self,
        origin: transformers.models.albert.modeling_albert.AlbertAttention,
        share_weights: bool = False,
    ):
        torch.nn.Module.__init__(self)
        if share_weights:
            self.dense = origin.dense
            self.LayerNorm = origin.LayerNorm
            self.dropout = origin.output_dropout
        else:
            self.dense = copy.deepcopy(origin.dense)
            self.LayerNorm = copy.deepcopy(origin.LayerNorm)
            self.dropout = copy.deepcopy(origin.output_dropout)


class QueryAlbertAttention(QueryBertLikeAttention):
    def __init__(
        self,
        origin: transformers.models.albert.modeling_albert.AlbertAttention,
        share_weights: bool = False,
    ):
        torch.nn.Module.__init__(self)
        self.self = QueryAlbertSelfAttention(origin, share_weights=share_weights)
        self.output = AlbertSelfOutput(origin, share_weights=share_weights)


class AlbertIntermediate(transformers.models.bert.modeling_bert.BertIntermediate):
    def __init__(
        self,
        origin: transformers.models.albert.modeling_albert.AlbertLayer,
        share_weights: bool = False,
    ):
        torch.nn.Module.__init__(self)
        if share_weights:
            self.dense = origin.ffn
            self.intermediate_act_fn = origin.activation
        else:
            self.dense = copy.deepcopy(origin.ffn)
            self.intermediate_act_fn = copy.deepcopy(origin.activation)


class AlbertOutput(transformers.models.bert.modeling_bert.BertOutput):
    def __init__(
        self,
        origin: transformers.models.albert.modeling_albert.AlbertLayer,
        share_weights: bool = False,
    ):
        torch.nn.Module.__init__(self)
        if share_weights:
            self.dense = origin.ffn_output
            self.LayerNorm = origin.full_layer_layer_norm
            self.dropout = origin.dropout
        else:
            self.dense = copy.deepcopy(origin.ffn_output)
            self.LayerNorm = copy.deepcopy(origin.full_layer_layer_norm)
            self.dropout = copy.deepcopy(origin.dropout)


class QueryAlbertLayer(QueryBertLikeLayer):
    def __init__(
        self,
        origin: transformers.models.albert.modeling_albert.AlbertLayer,
        share_weights: bool = False,
    ):
        torch.nn.Module.__init__(self)
        self.attention = QueryAlbertAttention(
            origin.attention, share_weights=share_weights
        )
        self.intermediate = AlbertIntermediate(origin, share_weights=share_weights)
        self.output = AlbertOutput(origin, share_weights=share_weights)


class QueryBertLikeEncoder(torch.nn.Module):
    """A module of architecture corresponding to `BertEncoder`, where the `query_states` are different from `hidden_states`.
    Typically, the `hidden_states` are pre-calculated by a `BertModel`.
    To some extent, this module replaces the self-attention with cross-attention in `BertEncoder`.

    References
    ----------
    transformers/models/bert/modeling_bert.py
    """

    def __init__(
        self,
        origin: transformers.models.bert.modeling_bert.BertEncoder,
        num_layers: int = None,
        share_weights: bool = False,
    ):
        super().__init__()
        self.num_layers = (
            origin.config.num_hidden_layers if num_layers is None else num_layers
        )

        if isinstance(
            origin, transformers.models.albert.modeling_albert.AlbertTransformer
        ):
            assert len(origin.albert_layer_groups) == 1
            assert len(origin.albert_layer_groups[0].albert_layers) == 1
            self.layer = QueryAlbertLayer(
                origin.albert_layer_groups[0].albert_layers[0],
                share_weights=share_weights,
            )
        else:
            self.layer = torch.nn.ModuleList(
                [
                    QueryBertLikeLayer(layer, share_weights=share_weights)
                    for layer in origin.layer[-self.num_layers :]
                ]
            )

    def _layers(self):
        for i in range(self.num_layers):
            if isinstance(self.layer, QueryAlbertLayer):
                yield self.layer
            else:
                yield self.layer[i]

    def forward(
        self,
        query_states,
        all_hidden_states,
        attention_mask=None,
        head_mask=None,
        output_attentions=False,
        output_query_states=False,
        return_dict=True,
    ):
        assert len(all_hidden_states) == self.num_layers + 1

        all_query_states = () if output_query_states else None
        all_self_attentions = () if output_attentions else None
        # The top query states are computed:
        #     (1) by the top Transformer block
        #     (2) from the second-top query states and second-top hidden states (key/value states)
        # Hence, the top hidden states are *not* used, because there is no more upper Transformer blocks
        for i, (layer_module, hidden_states) in enumerate(
            zip(self._layers(), all_hidden_states[:-1])
        ):
            if output_query_states:
                all_query_states = all_query_states + (query_states,)

            layer_head_mask = head_mask[i] if head_mask is not None else None
            layer_outputs = layer_module(
                query_states,
                hidden_states,
                attention_mask=attention_mask,
                head_mask=layer_head_mask,
                output_attentions=output_attentions,
            )
            # Update query_states to layer+1
            query_states = layer_outputs[0]

            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)

        if output_query_states:
            all_query_states = all_query_states + (query_states,)

        if not return_dict:
            return tuple(
                v for v in [query_states, all_query_states, all_self_attentions]
            )
        else:
            return dict(
                last_query_state=query_states,
                query_states=all_query_states,
                attentions=all_self_attentions,
            )
