"""Operations for embedding data into a (usually higher dimensional) space."""

import enum as E
import math
import typing

import torch


class FourierEmbedMode(E.StrEnum):
    """Various modes of generating Fourier embeddings."""

    SIN = E.auto()
    SIN_COS = E.auto()
    COS = E.auto()
    COS_SIN = E.auto()


type FourierEmbedType = (
    FourierEmbedMode | typing.Literal["sin", "cos", "sin_cos", "cos_sin"]
)


def fourier_embed_2d(
    input: torch.Tensor,
    dim_embed: int,
    *,
    f_max: float | None = None,
    mode: FourierEmbedType = FourierEmbedMode.SIN,
    linear: bool = False,
) -> torch.Tensor:
    """Embed dense 2D data using a Fourier basis."""
    B, C, H, W = input.shape
    factory_kwargs = {"device": input.device, "dtype": input.dtype}
    num_bands = (
        round(dim_embed / (2 * C))
        if mode in (FourierEmbedMode.SIN_COS, FourierEmbedMode.COS_SIN)
        else round(dim_embed / C)
    )
    f_max = max(H, W) / 2.0 if f_max is None else float(f_max)

    if not linear:
        s = 2 ** torch.linspace(
            0.0,
            math.log2(f_max),
            num_bands,
            **factory_kwargs,
        )
    else:
        s = torch.linspace(
            1.0,
            f_max / 2,
            num_bands,
            **factory_kwargs,
        )
    s = s * torch.pi

    input = torch.einsum("bchw,s->bcshw", input, s).reshape(B, -1, H, W)
    match mode:
        case FourierEmbedMode.SIN:
            input = input.sin()
        case FourierEmbedMode.COS:
            input = input.cos()
        case FourierEmbedMode.SIN_COS:
            input = torch.cat([input.sin(), input.cos()], dim=1)
        case FourierEmbedMode.COS_SIN:
            input = torch.cat([input.cos(), input.sin()], dim=1)
        case _:
            msg = f"Unknown Fourier embedding mode: {mode}"
            raise NotImplementedError(msg)

    npad = dim_embed - input.size(1)
    return torch.nn.functional.pad(
        input, (0, 0, 0, 0, 0, npad), mode="constant", value=0
    )
