# coding=utf-8
"""Spiral Rotary Positional Encoding (SRPE).

Implements Section VI of the Wiola paper. For position ``p`` and dimension-pair
index ``j in [d_h/2]``:

    omega_j   = theta_0 ** (-2j / d_h)
    Theta_j(p)= p * omega_j * (1 + 1/k_s)
    r_j(p)    = 1 + a_s * sin(p * f_s * omega_j)
    c_j(p)    = r_j(p) * cos(Theta_j(p))
    s_j(p)    = r_j(p) * sin(Theta_j(p))

and the rotation is applied to q/k with the standard rotate-half trick, which
is algebraically identical to Eqs. (12)-(13):

    out_lo = q_lo * c - q_hi * s
    out_hi = q_hi * c + q_lo * s

Unlike RoPE, ``Theta`` carries a second winding angle (via ``1 + 1/k_s``) and
the magnitude is modulated by a radial term ``r_j(p)``, embedding positions on a
3D helical manifold with no learned parameters.
"""

import torch
import torch.nn as nn


def rotate_half(x: torch.Tensor) -> torch.Tensor:
    """Rotate the two halves of the last dimension: [x1, x2] -> [-x2, x1]."""
    half = x.shape[-1] // 2
    x1 = x[..., :half]
    x2 = x[..., half:]
    return torch.cat((-x2, x1), dim=-1)


class SpiralRotaryEmbedding(nn.Module):
    def __init__(
        self,
        head_dim: int,
        max_position_embeddings: int = 2048,
        theta: float = 10000.0,
        spiral_divisor: int = 8,
        radial_amplitude: float = 0.1,
        radial_frequency: float = 0.01,
    ):
        super().__init__()
        if head_dim % 2 != 0:
            raise ValueError(f"head_dim must be even, got {head_dim}.")
        self.head_dim = head_dim
        self.max_position_embeddings = max_position_embeddings
        self.theta = theta
        self.k_s = spiral_divisor
        self.a_s = radial_amplitude
        self.f_s = radial_frequency

        # omega_j = theta ** (-2j / d_h), j in [0, d_h/2).
        j = torch.arange(0, head_dim, 2, dtype=torch.float32)
        inv_freq = theta ** (-(j / head_dim))  # shape [d_h/2]
        self.register_buffer("inv_freq", inv_freq, persistent=False)

        self._cached_len = 0
        self._build_cache(max_position_embeddings)

    def _build_cache(self, seq_len: int, device=None, dtype=torch.float32):
        positions = torch.arange(
            seq_len,
            dtype=torch.float32,
            device=device if device is not None else self.inv_freq.device,
        )
        inv_freq = self.inv_freq.to(positions.device)
        # Combined winding angle Theta_j(p) = p * omega_j * (1 + 1/k_s).
        combined = 1.0 + 1.0 / self.k_s
        angle = torch.outer(positions, inv_freq) * combined  # [T, d_h/2]
        # Radial modulation r_j(p) = 1 + a_s * sin(p * f_s * omega_j).
        radial = 1.0 + self.a_s * torch.sin(torch.outer(positions, inv_freq) * self.f_s)
        cos = radial * torch.cos(angle)
        sin = radial * torch.sin(angle)
        # Duplicate to full head_dim so it lines up with rotate_half.
        cos = torch.cat((cos, cos), dim=-1)  # [T, d_h]
        sin = torch.cat((sin, sin), dim=-1)
        self.register_buffer("cos_cached", cos.to(dtype), persistent=False)
        self.register_buffer("sin_cached", sin.to(dtype), persistent=False)
        self._cached_len = seq_len

    @torch.no_grad()
    def _maybe_extend(self, seq_len: int, device, dtype):
        if seq_len > self._cached_len or self.cos_cached.device != device:
            self._build_cache(max(seq_len, self._cached_len), device=device, dtype=dtype)

    def forward(self, x: torch.Tensor, position_ids: torch.Tensor):
        """Return (cos, sin) gathered for ``position_ids``.

        Args:
            x: any tensor only used to read dtype/device.
            position_ids: LongTensor of shape [batch, seq_len].
        Returns:
            cos, sin each of shape [batch, seq_len, head_dim].
        """
        max_pos = int(position_ids.max().item()) + 1
        self._maybe_extend(max_pos, x.device, torch.float32)
        cos = self.cos_cached[position_ids]  # [B, S, d_h]
        sin = self.sin_cached[position_ids]
        return cos.to(x.dtype), sin.to(x.dtype)


def apply_srpe(q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor):
    """Apply SRPE rotation to query and key tensors.

    Args:
        q, k: shape [batch, num_heads, seq_len, head_dim].
        cos, sin: shape [batch, seq_len, head_dim].
    """
    cos = cos.unsqueeze(1)  # [B, 1, S, d_h] -> broadcast over heads
    sin = sin.unsqueeze(1)
    q_rot = (q * cos) + (rotate_half(q) * sin)
    k_rot = (k * cos) + (rotate_half(k) * sin)
    return q_rot, k_rot