"""Rotary embeddings for the Evo1 HF port.

Two modes:
  * **Fast path** - when ``flash_attn`` is installed, we delegate to
    ``flash_attn.layers.rotary.RotaryEmbedding``, whose Triton kernel does
    the rotary multiply in fp32 internally (and is bit-exact with our
    pure-PyTorch path below).
  * **Fallback** - pure-PyTorch implementation, mathematically identical to
    flash_attn's kernel (multiply done in fp32 then cast back to bf16). Used
    when ``flash_attn`` isn't available.

The ``LinearlyScaledRotaryEmbedding`` subclass (used for the 131k variant)
overrides ``_update_cos_sin_cache`` to scale position indices, which works
identically against either parent class.
"""

from __future__ import annotations

import torch
import torch.nn as nn


try:
    from flash_attn.layers.rotary import RotaryEmbedding as _FlashRotaryEmbedding
    _HAS_FLASH_ROTARY = True
except ImportError:  # pragma: no cover - optional dep
    _FlashRotaryEmbedding = None  # type: ignore[assignment]
    _HAS_FLASH_ROTARY = False


def _rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor:
    """[x1, x2] -> [-x2, x1] over halves, or over even/odd pairs if interleaved."""
    if not interleaved:
        x1, x2 = x.chunk(2, dim=-1)
        return torch.cat((-x2, x1), dim=-1)
    x1, x2 = x[..., ::2], x[..., 1::2]
    return torch.stack((-x2, x1), dim=-1).flatten(-2)


@torch.no_grad()
def restore_rotary_inv_freq(
    rotary_emb: nn.Module,
    *,
    device=None,
    invalidate_cache: bool = True,
) -> None:
    """Rebuild ``inv_freq`` from RoPE metadata in fp32.

    Checkpoint-loaded values may already contain bf16 rounding even when the
    destination buffer is fp32. Reconstructing from ``base`` and ``dim`` is
    the only way to recover the exact frequencies.
    """
    if not hasattr(rotary_emb, "inv_freq"):
        return
    device = rotary_emb.inv_freq.device if device is None else device
    if hasattr(rotary_emb, "_compute_inv_freq"):
        inv_freq = rotary_emb._compute_inv_freq(device=device)
    else:
        inv_freq = 1.0 / (
            float(rotary_emb.base)
            ** (
                torch.arange(
                    0,
                    int(rotary_emb.dim),
                    2,
                    device=device,
                    dtype=torch.float32,
                )
                / int(rotary_emb.dim)
            )
        )
    rotary_emb.inv_freq = inv_freq.to(dtype=torch.float32)
    if invalidate_cache:
        rotary_emb._seq_len_cached = 0
        for name in (
            "_cos_cached",
            "_sin_cached",
            "_cos_k_cached",
            "_sin_k_cached",
        ):
            if hasattr(rotary_emb, name):
                setattr(rotary_emb, name, None)


def _apply_rotary(
    x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, interleaved: bool = False
) -> torch.Tensor:
    """Apply non-interleaved RoPE to the last `2 * cos.shape[-1]` dims of x.

    cos / sin shape: (T, rot_dim/2). x shape: (..., T, ..., D), where the rot
    is applied along the last dim. We expand cos/sin to broadcast over the
    leading dims.

    The multiplication is performed in fp32 internally (then cast back to
    x.dtype) to match flash_attn's Triton rotary kernel bit-exactly. Doing
    the multiply in bf16 directly compounds rounding error of ~3e-2 per
    layer, which becomes a ~1% relative error after 32 transformer blocks.
    """
    rot_dim = cos.shape[-1] * 2
    x_rot = x[..., :rot_dim]
    x_pass = x[..., rot_dim:]
    orig_dtype = x.dtype
    if interleaved:
        cos_full = torch.repeat_interleave(cos, 2, dim=-1).float()
        sin_full = torch.repeat_interleave(sin, 2, dim=-1).float()
    else:
        cos_full = torch.cat((cos, cos), dim=-1).float()
        sin_full = torch.cat((sin, sin), dim=-1).float()
    x_rot_f = x_rot.float()
    rotated = (x_rot_f * cos_full) + (_rotate_half(x_rot_f, interleaved) * sin_full)
    rotated = rotated.to(orig_dtype)
    return torch.cat((rotated, x_pass), dim=-1)


def _set_cos_sin_tables(rope: nn.Module, freqs: torch.Tensor, seqlen: int, dtype) -> None:
    """Fill a RoPE module's cos/sin caches from ``freqs``.

    Same as flash_attn / vortex RotaryEmbedding: with xPos (``rope.scale``
    set), queries use cos*scale and keys cos/scale, multiplied in fp32.
    """
    if rope.scale is None:
        rope._cos_cached = torch.cos(freqs).to(dtype)
        rope._sin_cached = torch.sin(freqs).to(dtype)
        return
    power = (
        torch.arange(seqlen, dtype=rope.scale.dtype, device=rope.scale.device) - seqlen // 2
    ) / rope.scale_base
    scale = rope.scale.to(device=power.device) ** power[:, None]
    rope._cos_cached = (torch.cos(freqs) * scale).to(dtype)
    rope._sin_cached = (torch.sin(freqs) * scale).to(dtype)
    rope._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)
    rope._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)


class _PureRotaryEmbedding(nn.Module):
    """Pure-PyTorch fallback RoPE (used when flash_attn is unavailable).

    Mirrors the public surface of ``flash_attn.layers.rotary.RotaryEmbedding``
    for the subset used by the Evo1 attention block: exposes ``inv_freq`` as
    a non-persistent runtime buffer and a forward(qkv) -> qkv method that
    rotates Q and K.
    """

    def __init__(
        self,
        dim: int,
        base: float = 10000.0,
        interleaved: bool = False,
        scale_base: float | None = None,
        pos_idx_in_fp32: bool = True,
        device=None,
    ):
        super().__init__()
        self.dim = dim
        self.base = float(base)
        self.interleaved = interleaved
        self.scale_base = scale_base
        self.pos_idx_in_fp32 = pos_idx_in_fp32

        inv_freq = self._compute_inv_freq(device=device)
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        scale = (
            (torch.arange(0, dim, 2, device=device, dtype=torch.float32) + 0.4 * dim) / (1.4 * dim)
            if scale_base is not None
            else None
        )
        self.register_buffer("scale", scale, persistent=False)

        self._seq_len_cached = 0
        self._cos_cached: torch.Tensor | None = None
        self._sin_cached: torch.Tensor | None = None
        self._cos_k_cached: torch.Tensor | None = None
        self._sin_k_cached: torch.Tensor | None = None

    def _compute_inv_freq(self, device=None) -> torch.Tensor:
        return 1.0 / (
            self.base
            ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim)
        )

    def _update_cos_sin_cache(self, seqlen: int, device=None, dtype=None):
        if (
            seqlen > self._seq_len_cached
            or self._cos_cached is None
            or self._cos_cached.device != device
            or self._cos_cached.dtype != dtype
            or (self.training and self._cos_cached.is_inference())
        ):
            self._seq_len_cached = seqlen
            if self.pos_idx_in_fp32:
                t = torch.arange(seqlen, device=device, dtype=torch.float32)
                inv_freq = self._compute_inv_freq(device=device)
            else:
                t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
                inv_freq = self.inv_freq

            _set_cos_sin_tables(self, torch.outer(t, inv_freq), seqlen, dtype)

    def forward(
        self,
        qkv: torch.Tensor,
        seqlen_offset: int | torch.Tensor = 0,
        max_seqlen: int | None = None,
    ) -> torch.Tensor:
        """Rotate Q and K of a packed (B, T, 3, H, D) qkv tensor.

        seqlen_offset is supported as int only (no per-sample offsets); for
        the inference KV-cache fast path we fall back to int(seqlen_offset).
        """
        if isinstance(seqlen_offset, torch.Tensor):
            seqlen_offset = int(seqlen_offset.max().item())
        T = qkv.shape[1]
        seqlen = max_seqlen if max_seqlen is not None else (T + seqlen_offset)
        self._update_cos_sin_cache(seqlen, device=qkv.device, dtype=qkv.dtype)

        span = slice(seqlen_offset, seqlen_offset + T)
        cos = self._cos_cached[span][None, :, None, :]
        sin = self._sin_cached[span][None, :, None, :]
        if self.scale is None:
            cos_k, sin_k = cos, sin
        else:
            cos_k = self._cos_k_cached[span][None, :, None, :]
            sin_k = self._sin_k_cached[span][None, :, None, :]
        q, k, v = qkv.unbind(dim=2)
        q = _apply_rotary(q, cos, sin, self.interleaved)
        k = _apply_rotary(k, cos_k, sin_k, self.interleaved)
        return torch.stack((q, k, v), dim=2)


if _HAS_FLASH_ROTARY:
    import inspect as _inspect

    # flash-attn >= 2.8 dropped this argument and always uses fp32 positions.
    _FLASH_TAKES_POS_IDX = (
        "pos_idx_in_fp32" in _inspect.signature(_FlashRotaryEmbedding.__init__).parameters
    )

    class RotaryEmbedding(_FlashRotaryEmbedding):
        """Flash RoPE with exact fp32 frequency reconstruction."""

        def __init__(self, *args, pos_idx_in_fp32: bool = True, **kwargs):
            if _FLASH_TAKES_POS_IDX:
                kwargs["pos_idx_in_fp32"] = pos_idx_in_fp32
            elif not pos_idx_in_fp32:
                raise ValueError(
                    "pos_idx_in_fp32=False needs flash-attn < 2.8, which always uses fp32 "
                    "positions; uninstall flash-attn to use the PyTorch fallback instead."
                )
            super().__init__(*args, **kwargs)
            self.pos_idx_in_fp32 = pos_idx_in_fp32

        def _update_cos_sin_cache(self, seqlen, device=None, dtype=None):
            needs_update = (
                seqlen > self._seq_len_cached
                or self._cos_cached is None
                or self._cos_cached.device != device
                or self._cos_cached.dtype != dtype
                or (self.training and self._cos_cached.is_inference())
            )
            if needs_update and self.pos_idx_in_fp32:
                restore_rotary_inv_freq(
                    self,
                    device=device,
                    invalidate_cache=False,
                )
            return super()._update_cos_sin_cache(
                seqlen,
                device=device,
                dtype=dtype,
            )
else:
    RotaryEmbedding = _PureRotaryEmbedding


class LinearlyScaledRotaryEmbedding(RotaryEmbedding):
    """RoPE with linear interpolation of position indices.

    Used for evo-1-131k-base: positions are divided by ``scaling_factor``
    before the cos/sin tables are computed, effectively stretching the
    trained context. The override is the same shape regardless of whether
    the parent class is flash_attn's RotaryEmbedding or our pure-PyTorch
    fallback (both expose the same ``_update_cos_sin_cache`` hook).
    """

    def __init__(self, dim: int, scaling_factor: float = 1.0, **kwargs):
        super().__init__(dim=dim, **kwargs)
        self._linear_scaling_factor = float(scaling_factor)

    def _update_cos_sin_cache(self, seqlen, device=None, dtype=None):
        # Mirrors the parent body but divides position indices by the linear
        # scaling factor before computing the cos/sin tables.
        if (
            seqlen <= self._seq_len_cached
            and self._cos_cached is not None
            and self._cos_cached.device == device
            and self._cos_cached.dtype == dtype
            and not (self.training and self._cos_cached.is_inference())
        ):
            return

        self._seq_len_cached = seqlen
        if self.pos_idx_in_fp32:
            t = torch.arange(seqlen, device=device, dtype=torch.float32)
            t = t / self._linear_scaling_factor
            inv_freq = self._compute_inv_freq(device=device)
        else:
            t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
            t = t / self._linear_scaling_factor
            inv_freq = self.inv_freq

        _set_cos_sin_tables(self, torch.outer(t, inv_freq), seqlen, dtype)


def swap_mha_rope(mha, new_rope=LinearlyScaledRotaryEmbedding, kwargs_new_rope=None):
    """Replace ``mha.rotary_emb`` with a freshly-constructed scaled RoPE.

    Mirrors ``stripedhyena.positional_embeddings.swap_mha_rope``: inherits
    dim/base/interleaved/scale_base/pos_idx_in_fp32 from the existing rope,
    deletes the old module, and attaches a new one of ``new_rope`` type
    configured with ``kwargs_new_rope``.
    """
    weight_attr = "Wq" if getattr(mha, "cross_attn", False) else "Wqkv"
    weight = getattr(mha, weight_attr).weight
    dtype = weight.dtype
    kwargs_old_rope = dict(
        dim=mha.rotary_emb.dim,
        base=mha.rotary_emb.base,
        interleaved=mha.rotary_emb.interleaved,
        scale_base=mha.rotary_emb.scale_base,
        pos_idx_in_fp32=mha.rotary_emb.pos_idx_in_fp32,
        device=mha.rotary_emb.inv_freq.device,
    )
    del mha.rotary_emb
    kwargs_new_rope = kwargs_new_rope or {"scaling_factor": 1.0}
    scaled = new_rope(**kwargs_new_rope, **kwargs_old_rope).to(dtype)
    restore_rotary_inv_freq(scaled)
    mha.rotary_emb = scaled
    return mha