import torch
from torch import nn


class PositionalEncodingsFixed(nn.Module):
    def __init__(self, emb_dim, temperature=10000):
        super(PositionalEncodingsFixed, self).__init__()

        self.emb_dim = emb_dim
        self.temperature = temperature

    def _1d_pos_enc(self, mask, dim):
        temp = torch.arange(self.emb_dim // 2).float().to(mask.device)
        temp = self.temperature ** (
            2 * (temp.div(2, rounding_mode="floor")) / self.emb_dim
        )

        enc = (~mask).cumsum(dim).float().unsqueeze(-1) / temp
        enc = torch.stack(
            [enc[..., 0::2].sin(), enc[..., 1::2].cos()], dim=-1
        ).flatten(-2)

        return enc

    def forward(self, bs, h, w, device):
        mask = torch.zeros(
            bs, h, w, dtype=torch.bool, requires_grad=False, device=device
        )
        x = self._1d_pos_enc(mask, dim=2)
        y = self._1d_pos_enc(mask, dim=1)

        return torch.cat([y, x], dim=3).permute(0, 3, 1, 2)
