import torch
import torch.nn as nn


def build_rope_cache(head_dim: int, max_seq_len: int, theta: float = 1_000_000.0):
    """Precompute the cos/sin rotation tables used by RoPE.

    Returns two tensors of shape (max_seq_len, head_dim), ready to index by position.
    """
    # each of the head_dim/2 pairs gets its own rotation speed (fast -> slow)
    idx = torch.arange(0, head_dim, 2, dtype=torch.float32)
    inv_freq = 1.0 / (theta ** (idx / head_dim))

    # outer product: rotation angle for every (position, pair) combination
    pos = torch.arange(max_seq_len, dtype=torch.float32)
    angles = torch.outer(pos, inv_freq)                      # (max_seq_len, head_dim/2)

    # duplicate each angle so it lines up with both halves of the vector
    angles = torch.cat([angles, angles], dim=-1)             # (max_seq_len, head_dim)
    return angles.cos(), angles.sin()


def _rotate_half(x: torch.Tensor) -> torch.Tensor:
    # turns [x1, x2] into [-x2, x1] — the 90-degree partner needed for rotation
    half = x.shape[-1] // 2
    x1, x2 = x[..., :half], x[..., half:]
    return torch.cat([-x2, x1], dim=-1)


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """Rotate query or key vectors by their positions.

    x: (batch, heads, seq_len, head_dim); cos/sin: (seq_len, head_dim).
    """
    # broadcast cos/sin over batch and head dims
    cos = cos[None, None, :, :]
    sin = sin[None, None, :, :]
    # the rotation identity: x*cos + rotate_half(x)*sin
    return x * cos + _rotate_half(x) * sin