import math
import torch
import torch.nn as nn


def positionalencoding1d(d_model, length):
    """
    :param d_model: dimension of the model
    :param length: length of positions
    :return: length*d_model position matrix
    """
    if d_model % 2 != 0:
        raise ValueError("Cannot use sin/cos positional encoding with "
                         "odd dim (got dim={:d})".format(d_model))
    pe = torch.zeros(length, d_model)
    position = torch.arange(0, length).unsqueeze(1)
    div_term = torch.exp((torch.arange(0, d_model, 2, dtype=torch.float) *
                         -(math.log(10000.0) / d_model)))
    pe[:, 0::2] = torch.sin(position.float() * div_term)
    pe[:, 1::2] = torch.cos(position.float() * div_term)

    return pe


def positionalencoding2d(d_model, height, width):
    """
    :param d_model: dimension of the model
    :param height: height of the positions
    :param width: width of the positions
    :return: d_model*height*width position matrix
    """
    if d_model % 4 != 0:
        raise ValueError("Cannot use sin/cos positional encoding with "
                         "odd dimension (got dim={:d})".format(d_model))
    pe = torch.zeros(d_model, height, width)
    # Each dimension use half of d_model
    d_model = int(d_model / 2)
    div_term = torch.exp(torch.arange(0., d_model, 2) *
                         -(math.log(10000.0) / d_model))
    pos_w = torch.arange(0., width).unsqueeze(1)
    pos_h = torch.arange(0., height).unsqueeze(1)
    pe[0:d_model:2, :, :] = torch.sin(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1)
    pe[1:d_model:2, :, :] = torch.cos(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1)
    pe[d_model::2, :, :] = torch.sin(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)
    pe[d_model + 1::2, :, :] = torch.cos(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)

    return pe


class PositionalEncoding2d(nn.Module):

    def __init__(self, d_model, dropout=0, max_len=100):
        super().__init__()

        assert d_model % 4 == 0, "d_model must be devisible by 4"

        self.dropout = nn.Dropout(p=dropout)

        with torch.no_grad():
            pe = torch.zeros(d_model, max_len, max_len)
            # Each dimension use half of d_model
            d_model = d_model // 2

            pos_w = torch.arange(0., max_len).unsqueeze(1)
            pos_h = torch.arange(0., max_len).unsqueeze(1)

            div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
            pe[0:d_model:2, :, :] = torch.sin(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, max_len, 1)
            pe[1:d_model:2, :, :] = torch.cos(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, max_len, 1)
            pe[d_model::2, :, :] = torch.sin(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, max_len)
            pe[d_model + 1::2, :, :] = torch.cos(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, max_len)

        self.register_buffer('pe', pe)

    def forward(self, x):
        batch_size, d_model, height, width = x.size()
        x = x + self.pe[:, :height, :width]
        return self.dropout(x)


# class PositionalEncoding2d(torch.nn.Module):
#     def __init__(self, d_model):
#         super().__init__()
#         if d_model % 4 != 0:
#             raise ValueError("Cannot use sin/cos positional encoding with "
#                              "odd dimension (got dim={:d})".format(d_model))
#         self.d_model = d_model
#
#     def forward(self, x):
#         batch_size, d_model, height, width = x.size()
#         assert d_model == self.d_model
#         pe = torch.zeros(d_model, height, width)
#         # Each dimension use half of d_model
#         d_model = d_model // 2
#         div_term = torch.exp(torch.arange(0., d_model, 2) *
#                              -(math.log(10000.0) / d_model))
#         pos_w = torch.arange(0., width).unsqueeze(1)
#         pos_h = torch.arange(0., height).unsqueeze(1)
#         pe[0:d_model:2, :, :] = torch.sin(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1)
#         pe[1:d_model:2, :, :] = torch.cos(pos_w * div_term).transpose(0, 1).unsqueeze(1).repeat(1, height, 1)
#         pe[d_model::2, :, :] = torch.sin(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)
#         pe[d_model + 1::2, :, :] = torch.cos(pos_h * div_term).transpose(0, 1).unsqueeze(2).repeat(1, 1, width)
#         pe = pe.unsqueeze(0).repeat(batch_size, 1, 1, 1)
#         return x + pe.detach  # is detach needed?
