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

def round_func(tensor, mode=1):
    tensor = tensor.float()
    fractional_part = tensor - tensor.floor()
    even_mask = (tensor.floor() % 2 == 0)
    if mode == 1:
        rounded_tensor = torch.where(fractional_part == 0.5,
                                     torch.where(even_mask, tensor.floor(), tensor.ceil()),
                                     torch.round(tensor))
    elif mode == 2:
        rounded_tensor = torch.where(fractional_part == 0.5, tensor + 0.5, torch.round(tensor))
    elif mode == 3:
        rounded_tensor = torch.round(tensor)
    return rounded_tensor

def compute_with_tensor(src, batch, seq_len, head_num, head_size):
    src = src.view(batch * seq_len * head_num * 2, head_size)
    mask = torch.ones(src.size(0), dtype=bool)
    for i in range(head_num, src.size(0), head_num * 2):
        mask[i:i+head_num] = False
    src = src[mask]
    max_value_tensor, _ = torch.max(src.abs(), dim=1)
    scale = (max_value_tensor / 127.0).view(batch, seq_len, head_num)
    scale_ori_quant = 127 / max_value_tensor
    temp_tensor = src * scale_ori_quant.view(scale_ori_quant.size(0),1)
    dst = round_func(temp_tensor, 1).view(batch, seq_len, head_num, head_size)
    return scale,dst

class TestPerHeadQuantizeKernel(UnitTest):
    @pytest.mark.parametrize("batch, seq_len, head_num, head_size, src_dtype", [
        (17,1436,31,21, torch.float32),
        (17,1436,31,21, torch.float16),
        (random.randint(1,32),random.randint(16,2048),random.randint(1,32),random.randint(16,256),torch.float32),
        (random.randint(1,32),random.randint(16,2048),random.randint(1,32),random.randint(16,256),torch.float16),
        (2,3,5,5, torch.float32)
        ])
    def test_all(self, batch, seq_len, head_num, head_size, src_dtype):
        mlu_name = torch.mlu.get_device_name()
        if "MLU3" in mlu_name and src_dtype == torch.bfloat16:
            src_dtype = torch.half
            print("BFLOAT16 is not support on {}, use half instead".format(mlu_name))
        torch.manual_seed(0)
        random.seed(0)
        hidden = head_num * head_size
        seq_stride = 2 * hidden
        scale_dtype = torch.float32
        dst_dtype = torch.int8
        dst_bs_stride = seq_len * hidden
        dst_seq_stride = hidden
        dst_head_stride = head_size
        src_bs_stride = seq_len * seq_stride
        src_seq_stride = seq_stride
        src_head_stride = head_size

        src = torch.rand((batch, seq_len, head_num, 2 * head_size), dtype=src_dtype).mlu()
        scale, dst = compute_with_tensor(src, batch, seq_len, head_num, head_size)
        scale_mlu = torch.zeros((batch, seq_len, head_num), dtype=scale_dtype).mlu()
        dst_mlu = torch.zeros((batch, seq_len, head_num, head_size), dtype=dst_dtype).mlu()
        btunittests.per_head_quantize_test(dst_mlu, scale_mlu, src, batch, seq_len, head_num, head_size, dst_bs_stride,
                                           dst_seq_stride, dst_head_stride, src_bs_stride, src_seq_stride,
                                           src_head_stride)
        dst_diff1 = super().diff1(dst, dst_mlu)
        dst_diff2 = super().diff2(dst, dst_mlu)
        scale_diff1 = super().diff1(scale, scale_mlu)
        scale_diff2 = super().diff2(scale, scale_mlu)
        assert dst_diff1 < 0.003 and dst_diff2 < 0.003
        assert scale_diff1 < 0.003 and scale_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)
