from __future__ import annotations

import torch
import torch.nn as nn


def build_1d_sincos_pos_emb(
    length: int, dim: int, temperature: float = 10000.0
) -> torch.Tensor:
    if dim % 2 != 0:
        raise ValueError(
            f"dim must be even for sin/cos positional embedding, got {dim}"
        )
    pos = torch.arange(length, dtype=torch.float32)
    freqs = torch.arange(dim // 2, dtype=torch.float32)
    omega = 1.0 / (temperature ** (freqs / (dim // 2)))
    out = torch.outer(pos, omega)
    emb = torch.cat([torch.sin(out), torch.cos(out)], dim=-1)
    return emb.unsqueeze(0)


class PositionalEmbeddingAdder(nn.Module):
    def __init__(self, dim: int, max_length: int):
        super().__init__()
        pos = build_1d_sincos_pos_emb(max_length, dim)
        self.register_buffer("pos_emb", pos, persistent=False)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        length = x.shape[1]
        return x + self.pos_emb[:, :length].to(device=x.device, dtype=x.dtype)


class PositionalQuery(nn.Module):
    """Returns a fixed sinusoidal query sequence for decoder targets."""

    def __init__(self, dim: int, max_length: int):
        super().__init__()
        pos = build_1d_sincos_pos_emb(max_length, dim)
        self.register_buffer("query", pos, persistent=False)

    def forward(
        self, batch_size: int, length: int, device: torch.device, dtype: torch.dtype
    ) -> torch.Tensor:
        q = self.query[:, :length].to(device=device, dtype=dtype)
        return q.expand(batch_size, -1, -1)
