from __future__ import annotations

import math

import torch
from torch import nn


class RotaryEmbedding(nn.Module):
    def __init__(self, d_head: int, max_seq_len: int, theta: float = 10_000.0) -> None:
        super().__init__()
        if d_head < 2 or d_head % 2 != 0:
            raise ValueError("RoPE requires a positive, even attention head dimension.")
        if max_seq_len < 1:
            raise ValueError("max_seq_len must be positive.")
        theta = float(theta)
        if not math.isfinite(theta) or theta <= 0:
            raise ValueError("RoPE theta must be positive.")

        self.d_head = d_head
        self.max_seq_len = max_seq_len
        self.theta = theta
        self.rope_dim = d_head // 2
        self.register_buffer("_cos", torch.empty(0), persistent=False)
        self.register_buffer("_sin", torch.empty(0), persistent=False)

    def forward(
        self,
        query: torch.Tensor,
        key: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        if query.shape != key.shape:
            raise ValueError("RoPE query and key tensors must have the same shape.")
        return self.rotate(query), self.rotate(key)

    def rotate(self, x: torch.Tensor) -> torch.Tensor:
        if x.ndim != 4 or x.shape[-1] != self.d_head:
            raise ValueError(f"RoPE expects shape [batch, heads, sequence, {self.d_head}].")
        seq_len = x.shape[-2]
        if seq_len > self.max_seq_len:
            raise ValueError(f"Sequence length {seq_len} exceeds RoPE limit {self.max_seq_len}.")

        cos, sin = self._cos_sin(x.device)
        cos = cos[:, :, :seq_len].to(dtype=x.dtype)
        sin = sin[:, :, :seq_len].to(dtype=x.dtype)
        first_half = x[..., : self.rope_dim]
        second_half = x[..., self.rope_dim :]
        return torch.cat(
            (
                cos * second_half + sin * first_half,
                -sin * second_half + cos * first_half,
            ),
            dim=-1,
        )

    def _cos_sin(self, device: torch.device) -> tuple[torch.Tensor, torch.Tensor]:
        if self._cos.numel() == 0 or self._cos.device != device:
            inverse_frequencies = self.theta ** (
                -torch.arange(self.rope_dim, dtype=torch.float32, device=device) / self.rope_dim
            )
            frequencies = torch.outer(
                torch.arange(self.max_seq_len, dtype=torch.float32, device=device),
                inverse_frequencies,
            )
            self._cos = frequencies.cos()[None, None, :, :]
            self._sin = frequencies.sin()[None, None, :, :]
        return self._cos, self._sin