import torch
import torch_mlu
import unittest
import torch_mlu_ops as ops
from common_utils import *
from itertools import product

def gen_args(bs, seq_len, q_heads, kv_heads, head_size, rope_dim, interleaved, discrete, dynamic_ntk, packed, dtype):
    cu_context_lens = None
    total_seq_len = bs * seq_len
    max_context_len = seq_len
    if packed:
        context_lens = torch.randint(size=(bs, ), low=1, high=seq_len+1, dtype=torch.int32, device='mlu')
        total_seq_len = context_lens.sum().item()
        max_context_len = context_lens.max().item()
        cu_context_lens = torch.cumsum(context_lens, dim=-1, dtype=torch.int32)
        cu_context_lens = torch.nn.functional.pad(cu_context_lens, (1,0), "constant", 0)

    context_shape = (total_seq_len, q_heads + kv_heads + 1, head_size) if packed else  \
    (bs, seq_len, q_heads + kv_heads + 1, head_size)
    context = torch.randn(size=context_shape, dtype=dtype).mlu()
    qk = context[..., 0 : q_heads + kv_heads, :]

    position_id = None
    if discrete:
        position_id = torch.randint(0, max_context_len, size=(total_seq_len,), dtype=torch.int32, device="mlu")
    else:
        position_id = torch.randint(0, max_context_len, size=(bs,), dtype=torch.int32, device="mlu")

    rope_seqlen = seq_len * 2
    cos_shape = (bs, rope_seqlen, rope_dim) if dynamic_ntk else (rope_seqlen, rope_dim)
    cos_cache = torch.randn(size=cos_shape, dtype=dtype, device="mlu")
    sin_cache = torch.randn(size=cos_shape, dtype=dtype, device="mlu")

    return (qk, sin_cache, cos_cache, position_id, cu_context_lens, \
                interleaved, discrete, dynamic_ntk, max_context_len)

class TestApplyRotaryOp(BtTestCase):
    def op_impl_base(self, *args):
        def rotate(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
        input, sin_cache, cos_cache, position_ids, cu_seqlen, interleaved, discrete, dynamic_ntk, max_seqlen = args
        packed = input.dim() == 3
        rope_dim = sin_cache.shape[-1]
        batch_size = cu_seqlen.shape[0] - 1 if packed else input.shape[0]
        sin_cache_float = sin_cache.float()
        cos_cache_float = cos_cache.float()
        cu_seqlen_cpu = cu_seqlen.cpu() if cu_seqlen is not None else None
        position_ids_cpu = position_ids.cpu() if position_ids is not None else None
        for i in range(batch_size):
            input_i = input[cu_seqlen_cpu[i] : cu_seqlen_cpu[i + 1]] if packed else input[i]
            input_i = input_i[..., 0:rope_dim]
            sin_cache_i = sin_cache_float[i] if dynamic_ntk else sin_cache_float
            cos_cache_i = cos_cache_float[i] if dynamic_ntk else cos_cache_float
            seq = input_i.shape[0]
            if discrete:
                if packed:
                    position_id_i = position_ids_cpu[cu_seqlen_cpu[i] : cu_seqlen_cpu[i + 1]]
                else:
                    position_id_i = position_ids_cpu.view(batch_size, -1)[i]
                sin_cache_i = sin_cache_i[position_id_i]
                cos_cache_i = cos_cache_i[position_id_i]
            else:
                if position_ids_cpu is None:
                    sin_cache_i = sin_cache_i[:seq]
                    cos_cache_i = cos_cache_i[:seq]
                else:
                    pos_id = position_ids_cpu[i].item()
                    sin_cache_i = sin_cache_i[pos_id : seq + pos_id]
                    cos_cache_i = cos_cache_i[pos_id : seq + pos_id]
            rot = rotate(input_i.float(), interleaved)
            output_i = rot * sin_cache_i.unsqueeze(1) + input_i * cos_cache_i.unsqueeze(1)
            input_i[:] = output_i.to(input.dtype)
        return input

    def test_apply_rotary(self):
        bs_list = [1, 8]
        seq_len_list = [1, 128, 1024]
        q_heads_list = [8, 32]
        kv_heads_list = [1]
        head_size_list = [128, 256]
        rope_dim_list = [256, 128, 64, 24]
        is_interleaved_list = [True, False]
        discrete_list = [True, False]
        dynamic_ntk_list = [True, False]
        packed_list = [True, False]
        dtype_list = [torch.half, torch.float]
        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)
        args = product(bs_list, seq_len_list, q_heads_list, kv_heads_list, head_size_list, \
                       rope_dim_list, is_interleaved_list, discrete_list, dynamic_ntk_list, packed_list, dtype_list)
        for bs, seq_len, q_heads, kv_heads, head_size, rope_dim, interleaved, discrete, dynamic_ntk, packed, dtype in args:
            print("bs: {}, seq_len: {}, q_heads: {}, kv_heads: {}, head_size: {}, "
                  "rope_dim: {}, interleaved: {}, discrete: {}, dynamic_ntk: {}, packed: {}, dtype: {} testing...".format(  \
                      bs, seq_len, q_heads, kv_heads, head_size, rope_dim, interleaved, discrete, dynamic_ntk, packed, dtype), flush=True)
            if rope_dim > head_size:
                print("rope_dim = {}, head_size = {},rope_dim should less than head_size".format(rope_dim, head_size))
                continue

            qk, sin_cache, cos_cache, position_id, cu_context_lens, \
                interleaved, _, _, max_context_len = gen_args(bs, seq_len, q_heads, kv_heads, head_size,
                                                            rope_dim, interleaved, discrete, dynamic_ntk, packed, dtype)

            qk_base  = qk.clone()
            self.op_impl_base(qk_base, sin_cache, cos_cache, position_id, cu_context_lens, \
                interleaved, discrete, dynamic_ntk, max_context_len)

            qk_out = ops.apply_rotary(qk, sin_cache, cos_cache, position_id, cu_context_lens, \
                interleaved, discrete, dynamic_ntk, max_context_len)

            self.assertTensorsEqual(qk_out.cpu().float(), qk_base.cpu().float(), 0.003, use_MSE=True, use_RAE=True)

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_prevent due to ASan issues")
    def test_prevent(self):
        bs, seq, head_num, head_size, dynamic_ntk, rotary_seqlen, rotary_dim, discrete = 1, 1024, 8, 128, False, 512, 128, True

        q = torch.randn(bs, seq, head_num, head_size, dtype=torch.half, device="mlu")
        sin = torch.randn(rotary_seqlen, rotary_dim, dtype=torch.half, device="mlu")
        cos = torch.randn(rotary_seqlen, rotary_dim, dtype=torch.half, device="mlu")
        self.assertException("discrete must be false if position ids is null.", ops.apply_rotary,
                                q, sin, cos, None, None, False, discrete, dynamic_ntk, seq)

        discrete = False
        self.assertException("max_seqlen must less than or equal to rope_seqlen.", ops.apply_rotary,
                                q, sin, cos, None, None, False, discrete, dynamic_ntk, seq)
        position_ids = torch.zeros(bs, dtype=torch.int32, device="mlu")
        self.assertException("max_seqlen must less than or equal to rope_seqlen.", ops.apply_rotary,
                                q, sin, cos, position_ids, None, False, discrete, dynamic_ntk, seq)

    def test_inductor(self):
        is_interleaved_list = [True, False]
        discrete_list = [True, False]
        dynamic_ntk_list = [True, False]
        packed_list = [True, False]
        params = product(is_interleaved_list, discrete_list, dynamic_ntk_list, packed_list)
        bs, seq_len, q_heads, kv_heads, head_size, rope_dim, dtype = 8, 1024, 8, 1, 256, 24, torch.float16
        for interleaved, discrete, dynamic_ntk, packed in params:
            print(f"==== check apply_rotary interleaved: {interleaved}, discrete: {discrete}, dynamic_ntk: {dynamic_ntk}, packed: {packed} ====")
            args = gen_args(bs, seq_len, q_heads, kv_heads, head_size,
                            rope_dim, interleaved, discrete, dynamic_ntk, packed, dtype)
            self.base_opcheck(torch.ops.torch_mlu_ops.apply_rotary, args)


if __name__ == '__main__':
    exit(run_unittest(TestApplyRotaryOp))
