import torch
from torch_mlu import mlu
import unittest
import torch_mlu_ops as ops
from common_utils import *
from itertools import product
from torch.nn import functional as F
import time

def gen_tenosr(batch, head_num, head_size, block_seq_len, max_seq_len, dtype, pack):
    if not pack:
        out           = torch.randn(batch, max_seq_len, head_num, head_size, device="mlu", dtype=dtype)
        lse           = torch.randn(batch, head_num, max_seq_len, device="mlu", dtype=torch.float32)
        block_out     = torch.randn(batch, block_seq_len, head_num, head_size, device="mlu", dtype=dtype)
        block_lse     = torch.randn(batch, head_num, block_seq_len, device="mlu", dtype=torch.float32)
        seq_offset    = torch.randint(low=0, high=(max_seq_len - block_seq_len + 1), size=(batch, ), dtype=torch.int32, device="mlu")
        cu_seqs       = None
        block_cu_seqs = None
    else:
        seq_lens         = torch.randint(low=1, high=(max_seq_len + 1), size=(batch, ), dtype=torch.int32)
        block_seq_lens   = torch.randint(low=1, high=(block_seq_len + 1), size=(batch, ), dtype=torch.int32)
        block_seq_lens   = torch.minimum(seq_lens, block_seq_lens)
        seq_offset       = torch.zeros_like(seq_lens)
        for i in range(batch):
            seq_offset[i] = torch.randint(low=0, high=seq_lens[i]-block_seq_lens[i]+1, size=(1,), dtype=torch.int32)
        seq_offset       = seq_offset.mlu()
        cu_seqs          = torch.cat((torch.tensor([0]),torch.cumsum(seq_lens, dim=0))).to(torch.int32).mlu()
        block_cu_seqs    = torch.cat((torch.tensor([0]),torch.cumsum(block_seq_lens, dim=0))).to(torch.int32).mlu()
        total_seqs       = torch.sum(seq_lens)
        block_total_seqs = torch.sum(block_seq_lens)

        out           = torch.randn(total_seqs, head_num, head_size, device="mlu", dtype=dtype)
        lse           = torch.randn(batch, head_num, max_seq_len, device="mlu", dtype=torch.float32)
        block_out     = torch.randn(block_total_seqs, head_num, head_size, device="mlu", dtype=dtype)
        block_lse     = torch.randn(batch, head_num, block_seq_len, device="mlu", dtype=torch.float32)
    return (out, lse, block_out, block_lse, seq_offset, cu_seqs, block_cu_seqs)

class TestUpdateOutAndLse(BtTestCase):
    def op_impl_base(self, *args):
        out, lse, block_out, block_lse, seq_offsets, cu_seqs, block_cu_seqs = args
        return update_out_and_lse_torch(out, lse, block_out, block_lse, seq_offsets, cu_seqs, block_cu_seqs)

    def test_update_out_and_lse(self):
        test_case_num = 10
        dtype_choice = [torch.float16, torch.float32]
        random.seed(time.time())

        count = test_case_num
        while count > 0:
            batch         = random.randint(1, 16 + 1)
            head_num      = random.randint(1, 16 + 1)
            head_size     = random.randint(1, 512 + 1)
            block_seq_len = random.choice([1, random.randint(2, 2048 + 1)])
            max_seq_len   = 1 if block_seq_len == 1 else max(random.randint(2, 2048 + 1), block_seq_len)
            pack          = random.choice([True, False])

            # 避免测试出现mlu显存不够
            if batch * head_num * head_size * max_seq_len > 10 * 1024 * 1024 * 1024:
                continue
            if torch_mlu.mlu.is_bf16_supported():
                dtype_choice.append(torch.bfloat16)
            else:
                if batch * head_num * head_size * max_seq_len > 1 * 1024 * 1024 * 1024:
                    continue
            dtype         = random.choice(dtype_choice)

            print(f"test_update_out_and_lse] {batch}, {head_num}, {head_size}, {block_seq_len}, {max_seq_len}, {dtype}, {pack}")

            out, lse, block_out, block_lse, seq_offset, cu_seqs, block_cu_seqs =    \
                gen_tenosr(batch, head_num, head_size, block_seq_len, max_seq_len, dtype, pack)

            out_torch, lse_torch = self.op_impl_base(out, lse, block_out, block_lse, seq_offset, cu_seqs, block_cu_seqs)
            ops.update_out_and_lse(out, lse, block_out, block_lse, seq_offset, cu_seqs, block_cu_seqs)
            self.assertTensorsEqual(lse.cpu().float(), lse_torch.cpu().float(), 0.005, use_MSE=True)
            self.assertTensorsEqual(out.cpu().float(), out_torch.cpu().float(), 0.005, use_MSE=True)
            count -= 1
        print(f"[test_update_out_and_lse] {test_case_num} cases test pass")

    def test_combine_ring_attn(self):
        batch, head_num, head_size, block_seq_len, max_seq_len, dtype, pack =\
            1, 8, 128, 8192, 8192, torch.bfloat16, True

        if not torch_mlu.mlu.is_bf16_supported():
            dtype = torch.float16

        print(f"[test_combine_ring_attn] {batch}, {head_num}, {head_size}, {block_seq_len}, {max_seq_len}, {dtype}, {pack}")

        out, lse, block_out, block_lse, seq_offset, cu_seqs, block_cu_seqs =    \
            gen_tenosr(batch, head_num, head_size, block_seq_len, max_seq_len, dtype, pack)
        # total_seq和block_total_seq为所有batch的真实seq长度之和，假设所有seq的
        # out: [sum(out_seqs), 64, 128] block_out [sum(block_out_seqs), 64, 128]
        # lse: [1, 8, 8192]          block_lse [1, 8, 8192]
        # seq_offset: [128]
        # cu_seqs block_cu_seqs: [128 + 1]

        out_torch, lse_torch = self.op_impl_base(out, lse, block_out, block_lse, seq_offset, cu_seqs, block_cu_seqs)

        ops.update_out_and_lse(out, lse, block_out, block_lse, seq_offset, cu_seqs, block_cu_seqs)

        self.assertTensorsEqual(lse.cpu().float(), lse_torch.cpu().float(), 0.005, use_MSE=True)
        self.assertTensorsEqual(out.cpu().float(), out_torch.cpu().float(), 0.005, use_MSE=True)
        print("[test_combine_ring_attn] pass")

    def test_combine_decoder_attn(self):
        batch_list = [16, 128]
        head_num_list = [64]
        head_size_list = [128]
        block_seq_len_list = [1]
        max_seq_len_list = [1]
        dtype_list = [torch.float16]
        pack_list = [False]

        if torch_mlu.mlu.is_bf16_supported():
            dtype_list.append(torch.bfloat16)

        args = product(batch_list, head_num_list, head_size_list, block_seq_len_list, max_seq_len_list,
                       dtype_list, pack_list)

        for batch, head_num, head_size, block_seq_len, max_seq_len, dtype, pack in args:
            print(f"[test_combine_decoder_attn] {batch}, {head_num}, {head_size}, {block_seq_len}, {max_seq_len}, {dtype}, {pack}")

            out, lse, block_out, block_lse, seq_offset, cu_seqs, block_cu_seqs =    \
                gen_tenosr(batch, head_num, head_size, block_seq_len, max_seq_len, dtype, pack)
            # out: [128, 1, 64, 128] block_out [128, 1, 64, 128]
            # lse: [128, 64, 1]      block_lse [128, 1]
            # seq_offset cu_seqs block_cu_seqs: None

            out_torch, lse_torch = self.op_impl_base(out, lse, block_out, block_lse, seq_offset, cu_seqs, block_cu_seqs)

            ops.update_out_and_lse(out, lse, block_out, block_lse, seq_offset, cu_seqs, block_cu_seqs)

            self.assertTensorsEqual(lse.cpu().float(), lse_torch.cpu().float(), 0.005, use_MSE=True)
            self.assertTensorsEqual(out.cpu().float(), out_torch.cpu().float(), 0.005, use_MSE=True)
        print("[test_combine_decoder_attn] pass")

    def test_inductor(self):
        pack_list = [True, False]
        for pack in pack_list:
            batch, head_num, head_size, dtype = 16, 8, 128, torch.float16
            if pack:
                block_seq_len, max_seq_len = 1024, 2048
            else:
                block_seq_len, max_seq_len = 1, 1
            out, lse, block_out, block_lse, seq_offset, cu_seqs, block_cu_seqs = \
                gen_tenosr(batch, head_num, head_size, block_seq_len, max_seq_len, dtype, pack)
            self.base_opcheck(torch.ops.torch_mlu_ops.update_out_and_lse, (out, lse, block_out, block_lse, seq_offset, cu_seqs, block_cu_seqs))
        print("[test_update_out_and_lse] test_inductor check pass")

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