# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from __future__ import annotations

import torch
from torch import Tensor, nn


def rotate_half_dimension(inputs: Tensor) -> Tensor:
    """Rotate pairs of feature halves by 90 degrees."""
    if inputs.shape[-1] % 2:
        raise ValueError("The final input dimension must be even.")
    first_half, second_half = inputs.chunk(2, dim=-1)
    return torch.cat((-second_half, first_half), dim=-1)


def apply_rotary_embedding(
    inputs: Tensor,
    cosine: Tensor,
    sine: Tensor,
) -> Tensor:
    """Apply precomputed rotary tables to query or key tensors."""
    sequence_length = inputs.shape[-2]
    cosine = cosine[:, :sequence_length, :]
    sine = sine[:, :sequence_length, :]

    return inputs * cosine + rotate_half_dimension(inputs) * sine


class RotaryEmbedding(nn.Module):
    """Rotary position embedding used for attention queries and keys."""

    def __init__(self, embedding_dimension: int) -> None:
        super().__init__()
        if embedding_dimension <= 0 or embedding_dimension % 2:
            raise ValueError("embedding_dimension must be a positive even integer.")

        inverse_frequencies = 1.0 / (
            10000
            ** (
                torch.arange(0, embedding_dimension, 2).float()
                / embedding_dimension
            )
        )
        self.register_buffer("inv_freq", inverse_frequencies)

        self._cached_sequence_length = 0
        self._cached_cosine: Tensor | None = None
        self._cached_sine: Tensor | None = None

    def _get_cosine_sine_tables(
        self,
        inputs: Tensor,
        sequence_dimension: int = 1,
    ) -> tuple[Tensor, Tensor]:
        sequence_length = inputs.shape[sequence_dimension]

        if (
            sequence_length != self._cached_sequence_length
            or self._cached_cosine is None
            or self._cached_sine is None
            or self._cached_cosine.device != inputs.device
            or self._cached_cosine.dtype != inputs.dtype
        ):
            self._cached_sequence_length = sequence_length
            positions = torch.arange(
                sequence_length,
                device=inputs.device,
                dtype=self.inv_freq.dtype,
            )
            frequencies = torch.einsum(
                "i,j->ij", positions, self.inv_freq
            )
            embedding = torch.cat(
                (frequencies, frequencies), dim=-1
            ).to(dtype=inputs.dtype)

            self._cached_cosine = embedding.cos()[None, :, :]
            self._cached_sine = embedding.sin()[None, :, :]

        return self._cached_cosine, self._cached_sine

    def forward(
        self,
        query: Tensor,
        key: Tensor,
    ) -> tuple[Tensor, Tensor]:
        """Rotate query and key tensors with shape ``[B, H, L, D]``."""
        if query.shape != key.shape:
            raise ValueError("query and key must have identical shapes.")
        self._cached_cosine, self._cached_sine = (
            self._get_cosine_sine_tables(
                key,
                sequence_dimension=-2,
            )
        )

        return (
            apply_rotary_embedding(
                query, self._cached_cosine, self._cached_sine
            ),
            apply_rotary_embedding(
                key, self._cached_cosine, self._cached_sine
            ),
        )