"""Continuous positional embedding helpers for voxel and query tokens."""

from __future__ import annotations

import math

import torch
import torch.nn as nn


class ContinuousFourierPositionEmbed(nn.Module):
    """Encode continuous coordinates into a learned token-space positional embedding."""

    def __init__(
        self,
        *,
        input_dim: int,
        embed_dim: int,
        num_bands: int = 6,
        hidden_dim: int | None = None,
        include_input: bool = True,
    ) -> None:
        super().__init__()
        self.input_dim = int(input_dim)
        self.num_bands = int(num_bands)
        self.include_input = bool(include_input)
        hidden_dim = int(hidden_dim) if hidden_dim is not None else max(int(embed_dim), 64)
        self.register_buffer(
            "freq_bands",
            2.0 ** torch.arange(self.num_bands, dtype=torch.float32) * math.pi,
            persistent=False,
        )
        encoded_dim = (self.input_dim if self.include_input else 0) + self.input_dim * self.num_bands * 2
        self.net = nn.Sequential(
            nn.Linear(encoded_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, embed_dim),
        )

    def _encode(self, x: torch.Tensor) -> torch.Tensor:
        pieces = [x] if self.include_input else []
        angles = x.unsqueeze(-1) * self.freq_bands
        pieces.append(torch.sin(angles).flatten(-2))
        pieces.append(torch.cos(angles).flatten(-2))
        return torch.cat(pieces, dim=-1)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.net(self._encode(x))


__all__ = ["ContinuousFourierPositionEmbed"]