import sys
import os
from unit_test import UnitTest
import pytest
import torch
import torch_mlu
import btunittests

class ApplyRotary(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def rotate(self, x: torch.Tensor, interleaved: bool):
        if not interleaved:
            x1, x2 = x.chunk(2, dim=-1)
            return torch.cat((-x2, x1), dim=-1)
        else:
            y = torch.empty_like(x)
            x1, x2 = x[..., ::2], x[..., 1::2]
            y[..., ::2], y[..., 1::2] = -x2, x1
            return y

    def forward(self,
                output: torch.Tensor,       # [total_seqlen, num_heads, head_size]
                input: torch.Tensor,        # [total_seqlen, num_heads, head_size]
                sin_cache: torch.Tensor,    # [rope_seqlen, rotary_dim] / [batch, rope_seqlen, rotary_dim]
                cos_cache: torch.Tensor,    #
                seq_offsets: torch.Tensor,  # [batch] / [batch, max_seqlen]
                cu_seq_lens: torch.Tensor,    # [batch + 1]
                interleaved: bool,
                discrete: bool,
                dynamic: bool,
                max_seqlen: int,
                no_offset:bool,
                rope_2d:bool):
        packed = input.dim() == 3
        rope_dim = sin_cache.shape[-1]
        batch_size = cu_seq_lens.shape[0] - 1 if packed else input.shape[0]

        if not rope_2d:
            for i in range(batch_size):
                input_i = input[cu_seq_lens[i] : cu_seq_lens[i + 1]] if packed else input[i]
                ouput_i = output[cu_seq_lens[i] : cu_seq_lens[i + 1]] if packed else output[i]
                input_i = input_i[..., 0:rope_dim]
                ouput_i = ouput_i[..., 0:rope_dim]
                start_seq_idx = 0 if no_offset else seq_offsets[i]
                sin_cache_i = sin_cache[i] if dynamic else sin_cache
                cos_cache_i = cos_cache[i] if dynamic else cos_cache
                seq = input_i.shape[0]
                if discrete:
                    start_seq_idx = 0
                    token_offset = seq_offsets[cu_seq_lens[i]:cu_seq_lens[i]+seq]
                    sin_cache_i = sin_cache_i[token_offset]
                    cos_cache_i = cos_cache_i[token_offset]
                sin_cache_i = sin_cache_i[start_seq_idx:seq+start_seq_idx]
                cos_cache_i = cos_cache_i[start_seq_idx:seq+start_seq_idx]
                rot = self.rotate(input_i, interleaved)

                ouput_i[:] = rot * sin_cache_i.unsqueeze(1) + input_i * cos_cache_i.unsqueeze(1)
            output[..., rope_dim:] = input[..., rope_dim:]
        else:
            for i in range(batch_size):
                input_i_left = input[cu_seq_lens[i] : cu_seq_lens[i + 1]] if packed else input[i]
                ouput_i_left = output[cu_seq_lens[i] : cu_seq_lens[i + 1]] if packed else output[i]
                input_i_left = input_i_left[..., 0:rope_dim]
                ouput_i_left = ouput_i_left[..., 0:rope_dim]
                sin_cache_i = sin_cache[i] if dynamic else sin_cache
                cos_cache_i = cos_cache[i] if dynamic else cos_cache
                seq = input_i_left.shape[0]
                if discrete:
                    token_offset = seq_offsets[0][cu_seq_lens[i]:cu_seq_lens[i]+seq]
                    sin_cache_i = sin_cache_i[token_offset]
                    cos_cache_i = cos_cache_i[token_offset]
                sin_cache_i = sin_cache_i[:seq]
                cos_cache_i = cos_cache_i[:seq]
                rot = self.rotate(input_i_left, interleaved)

                ouput_i_left[:] = rot * sin_cache_i.unsqueeze(1) + input_i_left * cos_cache_i.unsqueeze(1)

                input_i_right = input[cu_seq_lens[i] : cu_seq_lens[i + 1]] if packed else input[i]
                ouput_i_right = output[cu_seq_lens[i] : cu_seq_lens[i + 1]] if packed else output[i]
                input_i_right = input_i_right[..., rope_dim:]
                ouput_i_right = ouput_i_right[..., rope_dim:]
                sin_cache_i = sin_cache[i] if dynamic else sin_cache
                cos_cache_i = cos_cache[i] if dynamic else cos_cache
                seq = input_i_right.shape[0]
                if discrete:
                    token_offset = seq_offsets[1][cu_seq_lens[i]:cu_seq_lens[i]+seq]
                    sin_cache_i = sin_cache_i[token_offset]
                    cos_cache_i = cos_cache_i[token_offset]
                sin_cache_i = sin_cache_i[:seq]
                cos_cache_i = cos_cache_i[:seq]
                rot = self.rotate(input_i_right, interleaved)

                ouput_i_right[:] = rot * sin_cache_i.unsqueeze(1) + input_i_right * cos_cache_i.unsqueeze(1)

class TestRotaryEmbeddingKernel(UnitTest):
    @pytest.mark.parametrize(
        "batch, max_seq_len, head_num_q, head_num_kv, head_size, discrete, rotary_dim, interleaved, rope_2d, dynamic_ntk, packed, no_offset, data_type",
        [
          (4, 16, 8, 4, 128, False, 128, False, False, True, False, True, torch.float32),
          (4, 16, 8, 4, 128, False, 128, False, False, True, True, True, torch.float32),
          (4, 16, 8, 4, 128, True, 128, False, False, True, True, True, torch.float32),
          (4, 16, 8, 4, 128, True, 128, False, False, False, True, True, torch.float32),
          (4, 16, 8, 4, 128, False, 128, False, False, True, False, True, torch.float16),
          (4, 16, 8, 4, 128, False, 128, False, False, True, True, True, torch.float16),
          (4, 16, 8, 4, 128, True, 128, False, False, True, True, True, torch.float16),
          (4, 16, 8, 4, 128, True, 128, True, False, True, True, True, torch.float16),
          (4, 16, 8, 4, 128, True, 128, False, False, False, True, True, torch.float16),
          (4, 16, 8, 4, 128, True, 128, True, False, False, True, True, torch.float16),
          (4, 16, 8, 4, 128, True, 128, True, True, False, True, True, torch.float16),
          (4, 16, 8, 4, 128, False, 128, False, False, True,  True, False, torch.float32),
          (4, 16, 8, 4, 128, True, 128, False, False, True,  True, False, torch.float32),
          (4, 16, 8, 4, 128, True, 128, False, True, True,  True, False, torch.float32),
         ])
    def testall(
        self,
        batch,
        max_seq_len,
        head_num_q,
        head_num_kv,
        head_size,
        discrete,
        rotary_dim,
        interleaved,
        rope_2d,
        dynamic_ntk,
        packed,
        no_offset,
        data_type):
        torch.manual_seed(0)
        if rope_2d:
            interleaved = False
            discrete = True
            dynamic_ntk = False
            no_offset = False
        if head_num_q != 2 * head_num_kv:
            head_num_q = 2 * head_num_kv
        if rotary_dim %2 != 0:
            rotary_dim -= 1
        if no_offset:
            discrete = False
        if rope_2d:
            head_size = int(head_size if (head_size % 4 == 0) else (head_size // 4) * 4)
            rotary_dim = int(head_size / 2);

        seq_lens = torch.randint(size=(batch, ), low=1, high=max_seq_len+1, dtype=torch.int32, device='mlu')
        total_seq_len = batch * max_seq_len
        if not packed:
            seq_lens = torch.randint(size=(batch, ), low=max_seq_len, high=max_seq_len+1, dtype=torch.int32, device='mlu')
        max_context_len = int(seq_lens.max())

        cu_seq_lens = None
        if packed:
            cu_seq_lens = torch.cumsum(seq_lens, dim=-1, dtype=torch.int32)
            cu_seq_lens = torch.nn.functional.pad(cu_seq_lens, (1,0), "constant", 0).mlu()
            total_seq_len = int(cu_seq_lens[-1])

        head_num_qkv = head_num_q + head_num_kv
        head_num_qk = head_num_q + head_num_kv

        context_shape = (total_seq_len, head_num_qkv + 1, head_size) if packed else  \
        (batch, max_seq_len, head_num_qkv + 1, head_size)
        context = torch.randn(size=context_shape, dtype=data_type).mlu()

        qkv = context[..., 0 : head_num_qkv, :]
        qk = context[..., 0 : head_num_qk, :]
        seq_offsets = None
        if rope_2d:
            seq_offsets = torch.randint(size=(2, total_seq_len), low=1, high=max_seq_len + 1, dtype=torch.int32, device='mlu')
        else:
            if discrete and not no_offset:
                seq_offsets = torch.randint(size=(total_seq_len,), low=1, high=max_seq_len + 1, dtype=torch.int32, device='mlu')
            elif not no_offset:
                seq_offsets = torch.randint(size=(batch,), low=1, high=max_seq_len + 1, dtype=torch.int32, device='mlu')

        cos_cache = None
        sin_cache = None
        if dynamic_ntk:
            cos_cache = torch.randn(size=(batch, max_seq_len * 2, rotary_dim), dtype=data_type).mlu()
            sin_cache = torch.randn(size=(batch, max_seq_len * 2, rotary_dim), dtype=data_type).mlu()
        else:
            cos_cache = torch.randn(size=(max_seq_len * 2, rotary_dim), dtype=data_type).mlu()
            sin_cache = torch.randn(size=(max_seq_len * 2, rotary_dim), dtype=data_type).mlu()

        base_output = torch.empty_like(qk)

        apply_rotary = ApplyRotary()
        apply_rotary(base_output, qkv, sin_cache, cos_cache, \
            seq_offsets, cu_seq_lens, interleaved, discrete, dynamic_ntk, max_context_len, no_offset, rope_2d)

        output = btunittests.rotary_embedding_test(
            qkv,
            sin_cache,
            cos_cache,
            seq_offsets,
            cu_seq_lens,
            interleaved,
            discrete,
            dynamic_ntk,
            rope_2d,
            max_context_len,
            no_offset)

        diff1 = super().diff1(output, base_output)
        diff2 = super().diff2(output, base_output)
        assert diff1 < 0.003 and diff2 < 0.003


if __name__ == "__main__":
    exit_code = pytest.main(
        [
            "-vs",
            os.path.abspath(__file__),
            "--junitxml="
            + os.path.splitext(os.path.basename(__file__))[0]
            + "_report.xml",
        ]
    )
    exit(exit_code)
