""" Multi-Head Attention module """
import math
import torch
import torch.nn as nn
from torch.autograd import Variable


class MultiHeadedAttention(nn.Module):
    """
    Multi-Head Attention module from
    "Attention is All You Need"
    :cite:`DBLP:journals/corr/VaswaniSPUJGKP17`.
    Similar to standard `dot` attention but uses
    multiple attention distributions simulataneously
    to select relevant items.
    .. mermaid::
       graph BT
          A[key]
          B[value]
          C[query]
          O[output]
          subgraph Attn
            D[Attn 1]
            E[Attn 2]
            F[Attn N]
          end
          A --> D
          C --> D
          A --> E
          C --> E
          A --> F
          C --> F
          D --> O
          E --> O
          F --> O
          B --> O
    Also includes several additional tricks.
    Args:
       head_count (int): number of parallel heads
       model_dim (int): the dimension of keys/values/queries,
           must be divisible by head_count
       dropout (float): dropout parameter
    """

    def __init__(self, head_count, model_dim, d_k, d_v, dropout=0.1, clip_dist=0, use_neg_dist=False):
        super(MultiHeadedAttention, self).__init__()

        self.head_count = head_count
        self.model_dim = model_dim
        self.d_k = d_k
        self.d_v = d_v
        self.clip_dist = clip_dist
        self.use_neg_dist = use_neg_dist

        self.linear_keys = nn.Linear(model_dim, head_count * self.d_k)
        self.linear_query = nn.Linear(model_dim, head_count * self.d_k)
        self.linear_values = nn.Linear(model_dim, head_count * self.d_v)

        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)
        self.final_linear = nn.Linear(self.head_count * d_v, model_dim)

        # for relative-position aware self-attention
        if self.clip_dist > 0:
            self.edge_keys = nn.Parameter(torch.FloatTensor(2 * self.clip_dist + 1, self.d_k))
            self.edge_values = nn.Parameter(torch.FloatTensor(2 * self.clip_dist + 1, self.d_v))
            torch.nn.init.xavier_normal_(self.edge_keys)
            torch.nn.init.xavier_normal_(self.edge_values)

    def forward(self, key, value, query, mask=None):
        """
        Compute the context vector and the attention vectors.
        Args:
           key (`FloatTensor`): set of `key_len`
                key vectors `[batch, key_len, dim]`
           value (`FloatTensor`): set of `key_len`
                value vectors `[batch, key_len, dim]`
           query (`FloatTensor`): set of `query_len`
                 query vectors  `[batch, query_len, dim]`
           mask: binary mask indicating which keys have
                 non-zero attention `[batch, query_len, key_len]`
        Returns:
           (`FloatTensor`, `FloatTensor`) :
           * output context vectors `[batch, query_len, dim]`
           * one of the attention vectors `[batch, query_len, key_len]`
        """

        # CHECKS
        # batch, k_len, d = key.size()
        # batch_, k_len_, d_ = value.size()
        # aeq(batch, batch_)
        # aeq(k_len, k_len_)
        # aeq(d, d_)
        # batch_, q_len, d_ = query.size()
        # aeq(batch, batch_)
        # aeq(d, d_)
        # aeq(self.model_dim % 8, 0)
        # if mask is not None:
        #    batch_, q_len_, k_len_ = mask.size()
        #    aeq(batch_, batch)
        #    aeq(k_len_, k_len)
        #    aeq(q_len_ == q_len)
        # END CHECKS

        batch_size = key.size(0)
        head_count = self.head_count
        key_len = key.size(1)
        query_len = query.size(1)

        def shape(x, dim):
            """  projection """
            return x.view(batch_size, -1, head_count, dim).transpose(1, 2)

        def unshape(x, dim):
            """  compute context """
            return x.transpose(1, 2).contiguous().view(batch_size, -1, head_count * dim)

        if self.clip_dist > 0:
            dist_x = torch.arange(0, key_len).unsqueeze(0)
            dist_y = torch.arange(0, key_len).unsqueeze(1)
            distance = dist_x - dist_y
            if not self.use_neg_dist:
                distance = torch.abs(distance)
            distance = torch.clamp(distance, min=-self.clip_dist, max=self.clip_dist)
            distance = Variable((distance + self.clip_dist).long())     # [1,1,len,len]
            # distance = distance.repeat(head_count, 1, 1)  # nhead x seq_len x seq_len
            # distance = distance.repeat(batch_size, 1, 1, 1)  # bsz x nhead x seq_len x seq_len
            distance = distance.cuda() if key.is_cuda else distance

        # 1) Project key, value, and query.
        # key_up: bsz x nhead x key_len x d_k
        key_up = shape(self.linear_keys(key), self.d_k)  # x_j W^K
        # value_up: bsz x nhead x key_len x d_v
        value_up = shape(self.linear_values(value), self.d_v)  # x_j W^V
        # query_up: bsz x nhead x query_len x d_k
        query_up = shape(self.linear_query(query), self.d_k)  # x_j W^Q

        # 2) Calculate and scale scores.
        query_up = query_up / math.sqrt(self.d_k)  # (x_j W^Q) / sqrt(d_k)
        # bsz x nhead x query_len x key_len
        scores = torch.matmul(query_up, key_up.transpose(2, 3))
        if self.clip_dist > 0:
            out = self.edge_keys.index_select(0, distance.view(-1))
            # 1 x 1 x key_len x key_len x d_k
            # out = out.view(1, 1, query_len, query_len, self.d_k).contiguous()
            out = out.view(1, 1, query_len, query_len, self.d_k)
            # bsz x nhead x query_len x key_len
            add_term = torch.matmul(query_up.unsqueeze(3), out.transpose(3, 4)).squeeze(3)
            scores = scores + add_term

        if mask is not None:
            mask = mask.unsqueeze(1).expand_as(scores)
            scores = scores.masked_fill(mask, -1e20)

        # 3) Apply attention dropout and compute context vectors.
        attn = self.softmax(scores)
        drop_attn = self.dropout(attn)  # bsz x nhead x seq_len x seq_len
        context = torch.matmul(drop_attn, value_up)     # bsz x nhead x seq_len x d_v
        if self.clip_dist > 0:
            # TODO: modify `value_up` tensor to add a_{ij}^V [Eq.3 in https://arxiv.org/pdf/1803.02155.pdf]
            # also split Eq.3 to save space
            out = self.edge_values.index_select(0, distance.view(-1))
            # 1 x 1 x key_len x key_len x dim
            # out = out.view(1, 1, key_len, key_len, self.d_v).contiguous()
            out = out.view(1, 1, key_len, key_len, self.d_v)
            add_term = torch.matmul(drop_attn.unsqueeze(3), out).squeeze(3)     # bsz x nhead x seq_len x d_v
            context = context + add_term

        # bsz x nhead x seq_len x d_v --> bsz x seq_len x (nhead * d_v)
        context = unshape(context, self.d_v)

        # bsz x seq_len x (nhead * d_v) --> bsz x seq_len x model_dim
        output = self.final_linear(context)
        # CHECK
        # batch_, q_len_, d_ = output.size()
        # aeq(q_len, q_len_)
        # aeq(batch, batch_)
        # aeq(d, d_)

        # Return one attn
        top_attn = attn.view(batch_size, head_count, query_len, key_len)[:, 0, :, :].contiguous()

        return output, top_attn
