import math

import torch
import torch.nn as nn


class PositionalEncoding(nn.Module):
    """
    Standard sinusoidal positional encoding from
    "Attention Is All You Need".
    """

    def __init__(self, embed_dim, dropout=0.1, max_len=10000):
        super().__init__()

        self.dropout = nn.Dropout(dropout)

        pe = torch.zeros(max_len, embed_dim)

        position = torch.arange(
            0,
            max_len,
            dtype=torch.float
        ).unsqueeze(1)

        div_term = torch.exp(
            torch.arange(
                0,
                embed_dim,
                2
            ).float() *
            (-math.log(10000.0) / embed_dim)
        )

        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)

        pe = pe.unsqueeze(0)

        # Stored as a buffer so it moves with the model
        # but isn't trained.
        self.register_buffer("pe", pe)

    def forward(self, x):
        """
        Args:
            x: Tensor of shape
               (batch_size, sequence_length, embed_dim)
        """

        seq_len = x.size(1)

        x = x + self.pe[:, :seq_len]

        return self.dropout(x)