"""Custom Rotary Positional Encoding (RoPE) module for DEPToE.

Inspired from:
 - https://github.com/RobertCsordas/moeut
 - https://www.kaggle.com/code/aeryss/rotary-postional-encoding-rope-pytorch
 - https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/layers/rotary.py

"""

from typing import cast

import torch
from einops import rearrange, repeat
from flash_attn.layers.rotary import apply_rotary_emb, rotate_half


class RotaryPositionalEncoding(torch.nn.Module):
    """Rotary Positional Encoding (RoPE) module."""

    def __init__(
        self,
        attn_heads_projection_size: int,
        base: int = 10000,
        scale_base: int | None = None,
        *,
        interleaved: bool = False,
        use_torch_apply_rotary_emb: bool = False,
    ) -> None:
        """Initialize the RotaryPositionalEncoding module.

        Parameters
        ----------
        attn_heads_projection_size : int
            The dimensionality of the attention head.
        base : int, optional
            The base for the frequency calculation, default is 10000.
        scale_base : int, optional
            If provided, scales the base by this factor
            (default is None, which means no scaling).
        interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style)
            instead of 1st half and 2nd half (GPT-NeoX style).

        use_torch_apply_rotary_emb : bool, optional
            If True, use the torch implementation of apply_rotary_emb instead of the

        """
        super().__init__()
        inv_freq = 1.0 / (
            base
            ** (
                torch.arange(0, attn_heads_projection_size, 2).float()
                / attn_heads_projection_size
            )
        )
        self.register_buffer("inv_freq", inv_freq, persistent=False)
        self.base = float(base)
        self.interleaved = interleaved

        self.scale_base = scale_base
        scale: torch.Tensor | None = (
            (
                torch.arange(
                    0,
                    attn_heads_projection_size,
                    2,
                    dtype=torch.float32,
                )
                + 0.4 * attn_heads_projection_size
            )
            / (1.4 * attn_heads_projection_size)
            if scale_base is not None
            else None
        )
        self.register_buffer("scale", scale, persistent=False)
        self.scale: torch.Tensor | None = cast("torch.Tensor | None", self.scale)  # type: ignore[reportUnnecessaryCast]

        # Cache variables (similar to flash_attn)
        self._seq_len_cached = 0
        self._cos_cached = None
        self._sin_cached = None
        self._cos_k_cached = None
        self._sin_k_cached = None
        self.use_torch_apply_rotary_emb = use_torch_apply_rotary_emb

    def _update_cos_sin_cache(
        self,
        sequence_length: int,
        device: torch.device | None = None,
        dtype: torch.dtype | None = None,
    ) -> None:
        """Update the cached cosine and sine tables based on the sequence length.

        Parameters
        ----------
        sequence_length : int
            The length of the sequence for which to compute the cosine and sine tables.
        device : torch.device, optional
            The device on which to create the tensors (default is None, which uses the
            current device).
        dtype : torch.dtype, optional
            The data type of the tensors (default is None, which uses the current
            dtype).

        """
        # Reset the tables if the sequence length has changed,
        # if we're on a new device (possibly due to tracing for instance),
        # or if we're switching from inference mode to training
        recompute_tables = (
            sequence_length > 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 recompute_tables:
            self._seq_len_cached = sequence_length

            # Use fp32 for precision (similar to flash_attn)
            t = torch.arange(sequence_length, device=device, dtype=torch.float32)

            # Recompute inv_freq if it's not in fp32
            if self.inv_freq.dtype != torch.float32:
                inv_freq = 1.0 / (
                    self.base
                    ** (
                        torch.arange(
                            start=0,
                            end=int(cast("torch.Tensor", self.inv_freq).shape[0] * 2),
                            step=2,
                            device=device,
                            dtype=torch.float32,
                        )
                        / (cast("torch.Tensor", self.inv_freq).shape[0] * 2)
                    )
                )
            else:
                inv_freq = self.inv_freq

            # Use torch.outer instead of einsum for better precision
            freqs = torch.outer(t, cast("torch.Tensor", inv_freq))
            if self.scale is None or self.scale_base is None:
                self._cos_cached = torch.cos(freqs).to(dtype)
                self._sin_cached = torch.sin(freqs).to(dtype)
            else:
                power = (
                    torch.arange(
                        sequence_length,
                        dtype=self.scale.dtype,
                        device=device,
                    )
                    - sequence_length // 2
                ) / self.scale_base
                scale = self.scale.to(device=power.device) ** rearrange(
                    power,
                    "s -> s 1",
                )
                # We want the multiplication by scale to happen in fp32
                self._cos_cached = (torch.cos(freqs) * scale).to(dtype)
                self._sin_cached = (torch.sin(freqs) * scale).to(dtype)
                self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)
                self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)

    def get(
        self,
        x: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
        """Get the cached cosine and sine tensors for the input tensor x.

        Parameters
        ----------
        x : torch.Tensor
            Input tensor of shape
            (batch_size, seq_len, n_heads, attn_heads_projection_size), where seq_len
            is the sequence length and attn_heads_projection_size is the dimension each
            attention head projects to.

        Returns
        -------
        tuple[torch.Tensor, torch.Tensor]
            Tuple containing the cached cosine and sine tensors, both of shape
            (..., seq_len, attn_heads_projection_size), where seq_len is the sequence
            length and attn_heads_projection_size is the dimension each attention head
            projects to.

        """
        # Retrieve the sequence length from the input tensor x
        seq_len = x.shape[1]
        self._update_cos_sin_cache(seq_len, device=x.device, dtype=x.dtype)
        assert self._cos_cached is not None, (
            "Cosine caches should not be None after update."
        )
        assert self._sin_cached is not None, (
            "Sine caches should not be None after update."
        )

        return (
            self._cos_cached,
            self._sin_cached,
            self._cos_k_cached,
            self._sin_k_cached,
        )

    @staticmethod
    def apply_rotary_pos_emb(  # noqa: PLR0913, PLR0917
        q: torch.Tensor,
        k: torch.Tensor,
        sin: torch.Tensor,
        cos: torch.Tensor,
        sin_k: torch.Tensor | None,
        cos_k: torch.Tensor | None,
        offset: int,
        *,
        interleaved: bool,
        use_torch_apply_rotary_emb: bool = False,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Apply rotary positional encoding to query and key tensors.

        Parameters
        ----------
        q : torch.Tensor
            The query tensor which is expected to have shape
            (batch_size, seq_len, n_heads, attn_heads_projection_size).
        k : torch.Tensor
            The key tensor which is expected to have shape
            (batch_size, seq_len, n_heads, attn_heads_projection_size).
        sin : torch.Tensor
            Sine component of the rotary positional encoding, shape
            (seq_len, attn_heads_projection_size).
        cos : torch.Tensor
            Cosine component of the rotary positional encoding, shape
            (seq_len, attn_heads_projection_size).
        sin_k : torch.Tensor | None
            Sine component for the key tensor, shape
            (seq_len, attn_heads_projection_size) or None if not used.
        cos_k : torch.Tensor | None
            Cosine component for the key tensor, shape
            (seq_len, attn_heads_projection_size) or None if not used.
        offset : int
            Offset to apply to the sequence dimension, typically used for sliding window
            attention mechanisms.
        interleaved : bool
            If True, rotate pairs of even and odd dimensions (GPT-J style) instead of
            1st half and 2nd half (GPT-NeoX style).
        use_torch_apply_rotary_emb : bool, optional
            If True, use the torch implementation of apply_rotary_emb instead of the
            Flash Attention implementation.

        Returns
        -------
        tuple[torch.Tensor, torch.Tensor]
            Tuple containing the query and key tensors with the rotary positional
            encoding applied, both of shape
            (batch_size, seq_len, n_heads, attn_heads_projection_size).

        """
        rot_emb = (
            apply_rotary_emb
            if not use_torch_apply_rotary_emb
            else apply_rotary_emb_torch
        )
        # Use the efficient flash_attn apply_rotary_emb function
        q_rot = cast(
            "torch.Tensor",
            rot_emb(
                x=q,
                cos=cos,
                sin=sin,
                interleaved=interleaved,
                seqlen_offsets=offset,
                inplace=True,
            ),
        )
        if sin_k is None or cos_k is None:
            k_rot = cast(
                "torch.Tensor",
                rot_emb(
                    x=k,
                    cos=cos,
                    sin=sin,
                    interleaved=interleaved,
                    inplace=True,
                ),
            )
        else:
            k_rot = cast(
                "torch.Tensor",
                rot_emb(
                    x=k,
                    cos=sin_k,
                    sin=cos_k,
                    interleaved=interleaved,
                    seqlen_offsets=offset,
                    inplace=True,
                ),
            )

        return q_rot, k_rot

    def forward(
        self,
        q: torch.Tensor,
        k: torch.Tensor,
        pos_offset: int = 0,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Forward pass of the RotaryPositionalEncoding module.

        Parameters
        ----------
        q : torch.Tensor
            The query tensor which is expected to have shape
            (batch_size, seq_len, n_heads, attn_heads_projection_size).
        k : torch.Tensor
            The key tensor which is expected to have shape
            (batch_size, seq_len, n_heads, attn_heads_projection_size).
        pos_offset : int, optional
            Offset to apply to the sequence dimension, typically used for sliding window
            attention mechanisms (default is 0).

        Returns
        -------
        tuple[torch.Tensor, torch.Tensor]
            Tuple containing the query and key tensors with the rotary positional
            encoding applied, both of shape
            (batch_size, seq_len, n_heads, attn_heads_projection_size).

        """
        cos, sin, cos_k, sin_k = self.get(q)
        return self.apply_rotary_pos_emb(
            q=q,
            k=k,
            sin=sin,
            cos=cos,
            sin_k=sin_k,
            cos_k=cos_k,
            offset=pos_offset,
            interleaved=self.interleaved,
            use_torch_apply_rotary_emb=self.use_torch_apply_rotary_emb,
        )


def apply_rotary_emb_torch(  # noqa: PLR0913
    x: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    seqlen_offsets: int | torch.Tensor = 0,
    cu_seqlens: torch.Tensor | None = None,
    *,
    interleaved: bool = False,
    inplace: bool = False,
) -> torch.Tensor:
    """A CPU-compatible implementation of RoPE that matches the flash-attn API.

    Args:
        x: Input tensor. Shape is (batch, seqlen, nheads, headdim) if
           cu_seqlens is None, else (total_seqlen, nheads, headdim).
        cos: Cosine embeddings, shape (max_seqlen, rotary_dim // 2).
        sin: Sine embeddings, shape (max_seqlen, rotary_dim // 2).
        interleaved: If True, uses GPT-J style rotation.
        inplace: If True, performs the operation in-place.
        seqlen_offsets: An integer or a tensor of shape (batch_size,) for
                      shifting positions, used with a KV cache.
        cu_seqlens: Cumulative sequence lengths for packed sequences.
        max_seqlen: Not used in this CPU implementation, for API compatibility.

    Returns:
        The rotated tensor.

    """  # noqa: D401
    # Work on a copy if not inplace
    out = x if inplace else x.clone()

    # Determine the dimension to which RoPE is applied
    rotary_dim = cos.shape[-1] * 2

    # Split the tensor into the part to be rotated and the part to be passed through
    x_rot = out[..., :rotary_dim]

    # --- Prepare Position IDs and Gather Embeddings ---
    if cu_seqlens is not None:
        # Packed sequences logic
        seq_lens = cu_seqlens[1:] - cu_seqlens[:-1]
        position_ids = torch.cat([torch.arange(sl, device=x.device) for sl in seq_lens])  # type: ignore[reportCallIssue]
        # Offset is an integer for packed sequences
        position_ids += int(seqlen_offsets)

        cos_gathered = cos[position_ids].unsqueeze(1)
        sin_gathered = sin[position_ids].unsqueeze(1)
    else:
        # Non-packed sequences logic
        batch_size, seqlen = x.shape[0], x.shape[1]
        position_ids = torch.arange(seqlen, device=x.device).view(1, seqlen)

        # Add offsets. Handles both int and tensor offsets.
        position_ids += (
            seqlen_offsets.view(batch_size, 1)
            if isinstance(seqlen_offsets, torch.Tensor)
            else seqlen_offsets
        )

        cos_gathered = cos[position_ids].unsqueeze(2)
        sin_gathered = sin[position_ids].unsqueeze(2)

    # Reshape cos/sin to match the format of the input tensor for rotation
    if interleaved:
        cos_applied = rearrange(cos_gathered, "... d -> ... 1 (d 2)")
        sin_applied = rearrange(sin_gathered, "... d -> ... 1 (d 2)")
    else:
        cos_applied = repeat(cos_gathered, "... d -> ... (2 d)")
        sin_applied = repeat(sin_gathered, "... d -> ... (2 d)")

    # --- Apply Rotation ---
    x_rotated = rotate_half(x_rot, interleaved=interleaved)
    x_transformed = x_rot * cos_applied + x_rotated * sin_applied

    # --- Concatenate and Return ---
    out[..., :rotary_dim] = x_transformed
    return out
