import torch
from einops import rearrange, repeat
from torch import einsum, Tensor
from torch.nn import Module


def rotate_half(x: Tensor, interleaved=True) -> Tensor:
    if not interleaved:
        # x_half1, x_half2 = x.chunk(2, dim=-1)
        # Using torch.split instead of chunk for ONNX export compatibility.
        x1, x2 = torch.split(x, x.size(-1) // 2, dim=-1)
        return torch.cat((-x2, x1), dim=-1)
    else:
        x = rearrange(x, '... (d r) -> ... d r', r=2)
        x1, x2 = x.unbind(dim=-1)
        x = torch.stack((-x2, x1), dim=-1)
        return rearrange(x, '... d r -> ... (d r)')


def apply_rotary_emb(freqs: Tensor, t: Tensor, interleaved=True) -> Tensor:
    rot_dim = freqs.shape[-1]
    t_to_rotate = t[..., :rot_dim]
    t_pass_through = t[..., rot_dim:]

    t_rotated = (t_to_rotate * freqs.cos()) + (rotate_half(t_to_rotate, interleaved) * freqs.sin())

    return torch.cat((t_rotated, t_pass_through), dim=-1)


class RotaryEmbedding(Module):
    def __init__(
            self,
            dim,
            theta=10000,
            max_seq_len=8192,
            interleaved: bool = True
    ):
        super().__init__()
        self.interleaved = interleaved
        self.cached_freqs_seq_len = max_seq_len
        inv_freq = 1. / (theta ** (torch.arange(0, dim, 2).float() / dim))
        self.register_buffer('inv_freq', inv_freq, persistent=False)
        self.register_buffer('cached_freqs', self._precompute_cache(max_seq_len), persistent=False)

    def _precompute_cache(self, seq_len: int):
        seq = torch.arange(seq_len, device=self.inv_freq.device, dtype=self.inv_freq.dtype)
        freqs = einsum('i, j -> i j', seq, self.inv_freq)
        if self.interleaved:
            freqs = repeat(freqs, '... n -> ... (n r)', r=2)
        else:
            freqs = torch.cat((freqs, freqs), dim=-1)
        return freqs

    def forward(self, seq_len: int) -> Tensor:
        if seq_len > self.cached_freqs_seq_len:
            raise RuntimeError("sequence exceeds RoPE max_seq_len!")
        return self.cached_freqs[0: seq_len].detach()

    def rotate_queries_or_keys(self, t: Tensor) -> Tensor:
        device, dtype, seq_len = t.device, t.dtype, t.shape[-2]
        freqs = self.forward(seq_len=seq_len)

        return apply_rotary_emb(freqs.to(device=device, dtype=dtype), t, self.interleaved)