"""`tile_kernels.transform` -- rotary position embedding.

This is the widest gap between the two signatures, and the one worth reading carefully:

* the compiled kernel takes `cos_sin_cache`, a single tensor of ``cat([cos, sin], -1)``, and
  **rotates in place**, returning ``None``;
* the reference takes `cos` and `sin` as separate tensors and **returns a new tensor**.

So the adapter has to split the cache, call the reference, and write the result back into the
caller's tensor -- because the call site relies on the in-place contract
(`rotary_embedding.py` documents it and discards the return value).

Measured against the compiled kernel: identical except one bf16 last-bit tie per 65536 elements,
which is rounding, not a logic difference.
"""

from __future__ import annotations

import torch

from deepinfer.tk_compat import use_references

if not use_references():
    from tile_kernels import transform as _compiled

from tile_kernels import torch as _ref_raw

from deepinfer.tk_compat._sync import reference_module

# Calls through `_ref` may sync; see `_sync.py` for why that is the reference's design.
_ref = reference_module(_ref_raw)


def apply_rotary(
    query: torch.Tensor,
    cos_sin_cache: torch.Tensor,
    key: torch.Tensor | None = None,
    positions: torch.Tensor | None = None,
    interleaved: bool = False,
    conjugate: bool = False,
    seqlen_offset: int = 0,
) -> None:
    """Rotate `query` (and `key`) in place. Returns None, like the compiled kernel."""
    if not use_references():
        return _compiled.apply_rotary(
            query,
            cos_sin_cache,
            key,
            positions,
            interleaved,
            conjugate,
            seqlen_offset,
        )

    # The cache stores cos and sin halves side by side; rot_dim is their total.
    rot_dim = cos_sin_cache.shape[-1]
    if rot_dim % 2:
        raise ValueError(f"cos_sin_cache last dim must be even, got {rot_dim}")
    half = rot_dim // 2
    cos = cos_sin_cache[..., :half]
    sin = cos_sin_cache[..., half:]

    def rotate(x: torch.Tensor) -> None:
        # The reference wants (seqlen, heads, rot_dim) or 4-D with a batch axis; the caller passes
        # (num_tokens, num_heads, rot_dim). Cast for the reshape and back after.
        original_shape = x.shape
        if x.dim() == 3:
            seqlen, heads, dim = original_shape
            view = x.view(seqlen, heads, dim)
        elif x.dim() == 4:
            batch, seqlen, heads, dim = original_shape
            view = x.view(batch * seqlen, heads, dim)
        else:
            raise ValueError(f"apply_rotary expects 3-D or 4-D input, got {x.dim()}-D")

        # `positions` is per-token; a flattened batch axis would need them tiled, which the
        # compiled kernel also does not accept.
        if x.dim() == 4 and positions is not None:
            raise NotImplementedError(
                "4-D query with explicit positions has no reference equivalent: the reference "
                "indexes positions along its batch axis. Pass None to use the default ordering."
            )

        rotated = _ref.rotary_embedding_ref(
            view,
            cos,
            sin,
            positions,
            seqlen_offsets=seqlen_offset,
            interleaved=interleaved,
            conjugate=conjugate,
        )
        # In-place contract: the caller keeps using its own tensor.
        x.copy_(rotated.view(original_shape))

    rotate(query)
    if key is not None:
        rotate(key)


def backend() -> str:
    return "ref" if use_references() else "cuda"