# Standard
import math

# PIP
import torch.nn as nn

# Custom


class ScaledDotProductAttention(nn.Module):
    def __init__(self, d_k):
        super().__init__()
        self.sqrt_d_k = math.sqrt(d_k)

    def forward(self, q, k, v):
        # q, k, v: (batch_size, seq_len, num_heads, d_k)

        q = q.transpose(1, 2)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        # q, k, v: (batch_size, num_heads, seq_len, d_k)

        k = k.transpose(2, 3)
        # k: (batch_size, num_heads, d_k, seq_len)

        out = torch.matmul(q, k) / self.sqrt_d_k
        # out: (batch_size, num_heads, seq_len, seq_len)

        out = F.softmax(out, dim=-1)
        # out: (batch_size, num_heads, seq_len, seq_len)

        out = torch.matmul(out, v)
        # out: (batch_size, num_heads, seq_len, d_k)

        out = out.transpose(1, 2)
        # out: (batch_size, seq_len, num_heads, d_k)

        return out


class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads):
        super().__init__()
        self.d_model = d_model
        self.num_heads = num_heads

        assert d_model % num_heads == 0
        self.d_k = d_model // num_heads

        self.value = nn.Linear(d_model, d_model)
        self.key = nn.Linear(d_model, d_model)
        self.query = nn.Linear(d_model, d_model)
        self.register_buffer('attention', ScaledDotProductAttention(self.d_k))
        self.fc = nn.Linear(d_model, d_model)

    def forward(self, v, k, q):
        # v, k, q: (batch_size, seq_len, d_model)

        # Linear Projection
        v = self.value(v)
        k = self.key(k)
        q = self.query(q)

        num_batches = v.shape[0]
        v = v.view(num_batches, -1, self.num_heads, self.d_k)
        k = k.view(num_batches, -1, self.num_heads, self.d_k)
        q = q.view(num_batches, -1, self.num_heads, self.d_k)
        # out: (batch_size, seq_len, num_heads, d_k)

        out = self.attention(q, k, v)
        # out: (batch_size, seq_len, num_heads, d_k)

        # TODO view + contiguous -> reshape
        out = out.contiguous().view(num_batches, -1, self.d_model)
        # out: (batch_size, seq_len, d_model)

        out = self.fc(out)
        # out: (batch_size, seq_len, d_model)

        return out
