"""Rotary Position Embedding (RoPE) — implémenté à la main.

RoPE encode la position en **faisant tourner** par paires les dimensions de Q et K,
au lieu d'**ajouter** un vecteur de position aux embeddings. Propriétés :

* la position devient **relative** dans le produit scalaire Q·K ;
* aucun paramètre appris, aucune contrainte sur la longueur ;
* meilleure extrapolation en longueur qu'un embedding positionnel absolu appris.

Justification complète : docs/architecture.md §2.1.

Convention : les dimensions ``(2i, 2i+1)`` d'un vecteur forment la paire ``i``.
"""

from __future__ import annotations

import torch


def build_rope_cache(
    seq_len: int,
    head_dim: int,
    theta: float = 10000.0,
    device: torch.device | None = None,
    dtype: torch.dtype = torch.float32,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Précalcule ``cos`` et ``sin`` pour les positions ``0..seq_len-1``.

    Retourne deux tenseurs de forme ``(seq_len, head_dim)``, où chaque paire
    ``(2i, 2i+1)`` partage le même angle.
    """
    if head_dim % 2 != 0:
        raise ValueError(f"head_dim doit être pair pour RoPE, reçu {head_dim}")

    half = head_dim // 2
    # Fréquences décroissantes : inv_freq[i] = theta^(-2i/head_dim)
    inv_freq = 1.0 / (theta ** (torch.arange(0, half, device=device, dtype=torch.float32) / half))
    positions = torch.arange(seq_len, device=device, dtype=torch.float32)

    freqs = torch.outer(positions, inv_freq)  # (seq_len, half)
    emb = torch.cat([freqs, freqs], dim=-1)   # (seq_len, head_dim)
    return emb.cos().to(dtype), emb.sin().to(dtype)


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """Échange et inverse les deux moitiés : ``[a, b] -> [-b, a]`` sur la dernière dim."""
    half = x.shape[-1] // 2
    x1, x2 = x[..., :half], x[..., half:]
    return torch.cat([-x2, x1], dim=-1)


def apply_rope(
    q: torch.Tensor,
    k: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Applique RoPE à Q et K.

    ``q``, ``k`` : ``(batch, n_head, seq_len, head_dim)``.
    ``cos``, ``sin`` : ``(seq_len, head_dim)`` (broadcast sur batch et têtes).
    """
    cos = cos[None, None, :, :]
    sin = sin[None, None, :, :]
    q_rot = q * cos + rotate_half(q) * sin
    k_rot = k * cos + rotate_half(k) * sin
    return q_rot, k_rot