import sys
import os
current_dir = os.path.dirname(os.path.abspath(__file__))
parent_dir = os.path.dirname(current_dir)
sys.path.append(parent_dir)
from unit_test import UnitTest
import pytest
import torch
import torch_mlu
import btunittests
import random

class TestSliceCuSeqlensKernel(UnitTest):
    @pytest.mark.parametrize("batch_size, parallel_num", [
        (8,3),
        (16,16),
        (32,16),
        (64,16),
        (128,16),
        (256,16),
        (512,16),
        (1024,16),
        (2048,16),
        (4096,16),
        ])
    def test_all(self, batch_size, parallel_num):
        torch.manual_seed(0)
        random.seed(0)
        if batch_size < parallel_num:
            return
        cu_seq_lens = torch.randint(low=0, high=1024*1024, size=(batch_size+1,), dtype=torch.int32).mlu()
        every = (batch_size + parallel_num - 1) // parallel_num
        repeat = batch_size // every
        remain = batch_size % every
        loop = repeat + (remain != 0)
        result = []
        for i in range(loop):
            elem_num = 1 + (i == loop - 1 and remain != 0 and remain or every)
            start_idx = i * every
            end_idx = start_idx + elem_num
            slice_tensor = cu_seq_lens[start_idx:end_idx]
            adjusted_tensor = slice_tensor - slice_tensor[0]
            result.append(adjusted_tensor)
        sliced_cu_seq_lens = torch.cat(result)
        sliced_cu_seq_lens_mlu = torch.zeros(batch_size + loop, dtype=torch.int32).mlu()
        btunittests.slice_cu_seq_lens_test(cu_seq_lens, sliced_cu_seq_lens_mlu, batch_size, parallel_num)
        sliced_cu_seq_lens_diff1 = super().diff1(sliced_cu_seq_lens, sliced_cu_seq_lens_mlu)
        sliced_cu_seq_lens_diff2 = super().diff2(sliced_cu_seq_lens, sliced_cu_seq_lens_mlu)
        assert sliced_cu_seq_lens_diff1 == 0 and sliced_cu_seq_lens_diff2 == 0

class TestGenerateCuSeqlensKernel(UnitTest):
    @pytest.mark.parametrize("seq_len, parallel_num, is_causal_mask, is_kv_seq_len", [
        (2154, 4, False, False),
        (3412, 3, True, False),
        (996, 7, False, True),
        (872, 4, False, False),
        (634, 12, True, False),
        (486, 4, False, True),
        (125, 6, False, False),
        ])
    def test_all(self, seq_len, parallel_num, is_causal_mask, is_kv_seq_len):
        every = (seq_len + parallel_num - 1) // parallel_num
        repeat = seq_len // every
        remain = seq_len % every
        loop = repeat + (remain != 0)

        generate_cu_seq_lens = torch.zeros(2 * loop, dtype=torch.int32).mlu()

        rep = seq_len // loop
        rem = seq_len % loop
        base = rep + (rem != 0)

        if is_causal_mask and is_kv_seq_len:
            for i in range(loop):
                generate_cu_seq_lens[2 * i + 1] = seq_len if i == loop - 1 else (i + 1) * base
        elif not is_kv_seq_len:
            for i in range(loop):
                generate_cu_seq_lens[2 * i + 1] = seq_len - (loop - 1) * base if i == loop - 1 else base
        elif not is_causal_mask and is_kv_seq_len:
            for i in range(loop):
                generate_cu_seq_lens[2 * i + 1] = seq_len

        generate_cu_seq_lens_mlu = torch.zeros(2 * loop, dtype=torch.int32).mlu()
        btunittests.generate_cu_seq_lens_test(generate_cu_seq_lens_mlu, seq_len, parallel_num, is_causal_mask, is_kv_seq_len)
        generate_cu_seq_lens_diff1 = super().diff1(generate_cu_seq_lens, generate_cu_seq_lens_mlu)
        generate_cu_seq_lens_diff2 = super().diff2(generate_cu_seq_lens, generate_cu_seq_lens_mlu)
        assert generate_cu_seq_lens_diff1 == 0 and generate_cu_seq_lens_diff2 == 0


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)
