"""Implements positional embedding modules for 2D, 3D and time-oriented inputs."""

import abc
import math
import typing

import torch
import torch.fx
from torch import Tensor, nn

from unipercept.ops.embed import fourier_embed_2d
from unipercept.vision.geometry import (
    GridMode,
    GridType,
    generate_coord_grid_like,
    generate_directions_grid,
)

__all__ = ["CoordPos2d", "RadialPos2d", "TrigPos2d", "TrigPos2dTime", "TrigPosTime"]


class _TrigPosBase(nn.Module, metaclass=abc.ABCMeta):
    """Trigonometric positional embedding base class."""

    scale: typing.Final[float | None]
    temperature: typing.Final[float]
    eps: typing.Final[float]
    cache: Tensor | None

    def __init__(
        self,
        dim_input: int,
        *,
        temperature: int = 10000,
        scale: float | None = 2 * math.pi,
        eps: float = 1e-6,
    ) -> None:
        super().__init__()

        self.scale = scale
        self.dim_input = dim_input
        self.temperature = temperature
        self.eps = eps

        self.register_buffer("cache", None, persistent=False)
        self._register_bands(dim_input // 2, dtype=torch.float32, device="cpu")  # type: ignore

    @abc.abstractmethod
    def _generate(self, input: Tensor, mask: Tensor | None = None) -> Tensor: ...

    @abc.abstractmethod
    def _register_bands(self, num_bands: int, **kwargs: typing.Any) -> None: ...  # noqa: ANN401

    @typing.override
    @torch.no_grad()
    def forward(self, input: Tensor, mask: Tensor | None = None) -> Tensor:
        if mask is not None or torch.compiler.is_compiling():
            return self._generate(input, mask)
        if self.cache is None or self.cache.size() != input.size():
            self.cache = self._generate(input, mask)
        return self.cache


class TrigPos2d(_TrigPosBase):
    """Trigonometric positional embedding for 2D inputs."""

    t: Tensor

    @typing.override
    def _register_bands(self, num_bands: int, **kwargs: typing.Any) -> None:
        with torch.no_grad():
            t = torch.arange(num_bands, **kwargs)
            t = self.temperature ** (2 * (t // 2) / num_bands)
        self.register_buffer("t", t, persistent=False)

    @typing.override
    def _generate(self, input: Tensor, mask: Tensor | None = None) -> Tensor:
        """Generate a 2D trigonometric positional embedding."""
        if mask is None:
            H = input.size(-2)
            W = input.size(-1)
            mask = input.new_ones((1, H, W))
        else:
            mask = ~mask
        with torch.autocast(device_type=input.device.type, enabled=False):
            pe = self._embed(mask.float())
        return pe.type_as(input).expand(input.size(0), -1, -1, -1)

    def _embed(self, mask: Tensor) -> Tensor:
        y_embed = mask.cumsum(-2, dtype=torch.float32)
        x_embed = mask.cumsum(-1, dtype=torch.float32)
        t = self.t.to(torch.float32)
        if self.scale is not None:
            y_embed = y_embed / (y_embed[..., -1:, :] + self.eps) * self.scale
            x_embed = x_embed / (x_embed[..., :, -1:] + self.eps) * self.scale
        pos_x = x_embed.unsqueeze(-1) / t
        pos_y = y_embed.unsqueeze(-1) / t
        pos_x = torch.stack(
            (pos_x[..., 0::2].sin(), pos_x[..., 1::2].cos()), dim=4
        ).flatten(3)
        pos_y = torch.stack(
            (pos_y[..., 0::2].sin(), pos_y[..., 1::2].cos()), dim=4
        ).flatten(3)
        return torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)


class TrigPos2dTime(_TrigPosBase):
    """Trigonometric positional embedding for 3D/2+1D inputs."""

    t: Tensor
    t_z: Tensor

    @typing.override
    def _register_bands(self, num_bands: int, **kwargs: typing.Any) -> None:
        with torch.no_grad():
            t = torch.arange(num_bands, **kwargs)
            t = self.temperature ** (2 * (t // 2) / num_bands)

            t_z = torch.arange((num_bands * 2), **kwargs)
            t_z = self.temperature ** (2 * (t_z // 2) / (num_bands * 2))

        self.register_buffer("t", t, persistent=False)
        self.register_buffer("t_z", t_z, persistent=False)

    @typing.override
    def _generate(self, input: Tensor, mask: Tensor | None = None) -> Tensor:
        if mask is None:
            T = input.size(-4)
            H = input.size(-2)
            W = input.size(-1)
            mask = input.new_ones((1, T, H, W))
        else:
            mask = ~mask
        with torch.autocast(device_type=input.device.type, enabled=False):
            pe = self._embed(mask.float())
        return pe.type_as(input).expand_as(input)

    def _embed(self, mask: Tensor) -> Tensor:
        """Generate a (2+1)D trigonometric positional embedding."""
        t_embed = mask.cumsum(1)
        y_embed = mask.cumsum(2)
        x_embed = mask.cumsum(3)

        if self.scale is not None:
            t_embed = t_embed / (t_embed[:, -1:, :, :] + self.eps) * self.scale
            y_embed = y_embed / (y_embed[:, :, -1:, :] + self.eps) * self.scale
            x_embed = x_embed / (x_embed[:, :, :, -1:] + self.eps) * self.scale

        pos_x = x_embed[:, :, :, :, None] / self.t
        pos_x = torch.stack(
            (pos_x[:, :, :, :, 0::2].sin(), pos_x[:, :, :, :, 1::2].cos()), dim=5
        ).flatten(4)

        pos_y = y_embed[:, :, :, :, None] / self.t
        pos_y = torch.stack(
            (pos_y[:, :, :, :, 0::2].sin(), pos_y[:, :, :, :, 1::2].cos()), dim=5
        ).flatten(4)

        pos_t = t_embed[:, :, :, :, None] / self.t_z
        pos_t = torch.stack(
            (pos_t[:, :, :, :, 0::2].sin(), pos_t[:, :, :, :, 1::2].cos()), dim=5
        ).flatten(4)

        return (torch.cat((pos_y, pos_x), dim=4) + pos_t).permute(0, 1, 4, 2, 3)


class TrigPosTime(_TrigPosBase):
    """Trigonometric positional embedding for time-oriented inputs."""

    t: Tensor

    @typing.override
    def _register_bands(self, num_bands: int, **kwargs: typing.Any) -> None:
        t = torch.arange(num_bands * 2, requires_grad=False, **kwargs)
        t = self.temperature ** (2 * (t // 2) / num_bands)
        self.register_buffer("t", t, persistent=False)

    @typing.override
    def _generate(self, input: Tensor, mask: Tensor | None = None) -> Tensor:
        if mask is None:
            T = input.size(-2)
            mask = input.new_ones((1, T))
        else:
            mask = ~mask
        with torch.autocast(device_type=input.device.type, enabled=False):
            p = self._embed(mask.float())
        return p.type_as(input).expand_as(input)

    def _embed(
        self,
        mask: Tensor,
    ) -> Tensor:
        """Generate trigonometric positional embedding for time-oriented inputs."""
        t_embed = mask.cumsum(0)
        if self.scale is not None:
            t_embed = t_embed / (t_embed[-1:, :] + self.eps) * self.scale
        pos_t = t_embed[:, :, None] / self.t
        return torch.stack(
            (pos_t[:, :, 0::2].sin(), pos_t[:, :, 1::2].cos()), dim=3
        ).flatten(2)


class CoordPos2d(nn.Module):
    """Embed 2D coordinates."""

    stride: typing.Final[int]
    f_max: typing.Final[float | None]
    cache: Tensor | None

    def __init__(
        self,
        dim_embed: int,
        *,
        stride: int = 1,
        f_max: float | None = None,
        grid_mode: GridType = GridMode.PIXEL_CENTER,
    ) -> None:
        super().__init__()
        self.dim_embed = dim_embed
        self.f_max = f_max
        self.stride = stride
        self.grid_mode = str(grid_mode)

        self.register_buffer("cache", None, persistent=False)

    @typing.override
    @torch.no_grad()
    def forward(self, input: Tensor) -> Tensor:
        if torch.compiler.is_compiling():
            return self._generate(input)
        if self.cache is None or self.cache.size() != input.size():
            self.cache = self._generate(input)
        return self.cache

    def _generate(self, input: Tensor) -> Tensor:
        with torch.autocast(input.device.type, enabled=False):
            pos_matrix = self._generate_grid(input)
            pos_embed = fourier_embed_2d(
                pos_matrix.float(),
                dim_embed=self.dim_embed,
                f_max=self.f_max,
            )
        return pos_embed.type_as(input)

    def _generate_grid(self, input: Tensor) -> Tensor:
        return (
            generate_coord_grid_like(input, mode=self.grid_mode, stride=self.stride)
            .permute(2, 0, 1)
            .unsqueeze(0)
            .repeat(input.size(0), 1, 1, 1)
        )


class RadialPos2d(CoordPos2d):
    """Embed 2D radiance."""

    @typing.override
    def _generate_grid(self, input: Tensor) -> Tensor:
        B, _C, H, W = input.shape

        # Allow scaling using a common stride factor
        H = H * self.stride
        W = W * self.stride

        # Generate homogeneous radial transform as a function of the input size
        r = ((H**2) + (W**2)) ** 0.5
        u = W / 2
        v = H / 2
        transform = input.new_tensor(
            [[[r, 0, u, 0], [0, r, v, 0], [0, 0, 0, 1], [0, 0, 1, 0]]],
        )

        # Generate the directions grid
        dirs = generate_directions_grid(
            transform,
            (H, W),
            dim=1,
            grid_mode=self.grid_mode,
        )
        if self.stride != 1:
            dirs = torch.nn.functional.interpolate(
                dirs,
                scale_factor=1 / self.stride,
                mode="bilinear",
                align_corners=False,
                antialias=False,
            )

        # Repeat along the batch dimension
        return dirs.repeat(B, 1, 1, 1)
