import torch
import torch_mlu
import unittest
import torch_mlu_ops as ops
from torch.nn.parameter import Parameter
from torch.nn import functional as F
import torch.multiprocessing as mp
import sys
import os

work_dir = os.path.dirname(os.path.abspath(__file__))
sys.path.append(os.path.dirname(work_dir))
from common_utils import *

def compute_weight_only_scale(weight, quant_bit):
    int_max = float(2 ** (quant_bit - 1) - 1)
    weight_max = torch.max(torch.abs(weight), axis=1, keepdims=True)
    weight_scale = torch.div(int_max, weight_max[0])
    weight_int = torch.mul(weight, weight_scale)
    weight_int = weight_int.type(torch.int8)
    weight_scale_recip = torch.div(weight_max[0], int_max).type(torch.float).squeeze()
    return weight_int, weight_scale_recip

def quant_mm_split_k(rank, *args):
    torch.manual_seed(0)
    M, N, K, has_bias, has_res, dtype, block_m, world_size = args
    assert K % world_size == 0
    block_k = K // world_size
    start_k = rank * block_k
    end_k = start_k + block_k

    setup(rank, world_size)
    pg = get_default_group()
    cncl_comm = pg._get_backend(torch.device("mlu")).get_cncl_comm(rank)

    a = torch.randn(M, K, device="mlu", dtype=dtype)
    b = torch.randn(N, K, device="mlu", dtype=dtype)
    b_int, gemm_output_scale = compute_weight_only_scale(b, 8)
    c = torch.randn(M, N, device="mlu", dtype=dtype) if has_res else None
    bias = torch.randn(N, device="mlu", dtype=dtype) if has_bias else None
    torch_quant_matmul = QuantMatmul(b_int, None, None, None, None, gemm_output_scale, dtype)
    pt_output = torch_quant_matmul(a).detach()

    a = a[..., start_k:end_k].contiguous()
    b_int = b_int[..., start_k:end_k].contiguous()
    bias = bias if (bias is not None and rank == 0) else None
    c = c if (c is not None and rank == 0) else None

    param = [cncl_comm, a, None, None, b_int, None, None, bias, c, None,
                None, gemm_output_scale, None, "half", "weight_only", "quantize_none",
                "quantize_per_channel", 8, 1.0, 1.0, False, True, block_m]
    # beta = beta if (c is not None and rank == 0) else 0.
    output = ops._ops.quant_matmul_allreduce(*param)
    if rank == 0:
        assertTensorsEqual(output.cpu().float(), pt_output.cpu().float(), 0.006, use_MSE=True, use_RAE=True)
    cleanup()

def sq_mm_split_k(rank, *args):
    torch.manual_seed(0)
    M, N, K, has_bias, has_res, dtype, world_size, block_m = args
    assert K % world_size == 0
    block_k = K // world_size
    start_k = rank * block_k
    end_k = start_k + block_k

    setup(rank, world_size)
    pg = get_default_group()
    cncl_comm = pg._get_backend(torch.device("mlu")).get_cncl_comm(rank)

    a = torch.randint(-10, 10, (M, K), dtype=torch.int8).mlu()
    b = torch.randint(-10, 10, (N, K), dtype=torch.int8).mlu()
    c = torch.randn(M, N, device="mlu", dtype=dtype) if has_res else None
    bias = torch.randn(N, device="mlu", dtype=dtype) if has_bias else None
    a_scale = torch.randn(M, device="mlu", dtype=torch.float)
    b_scale = torch.randn(N, device="mlu", dtype=torch.float)
    torch_quant_matmul = QuantMatmul(b, bias, c, a_scale, b_scale, None, dtype)
    pt_output = torch_quant_matmul(a).detach()

    a = a[..., start_k:end_k].contiguous()
    b = b[..., start_k:end_k].contiguous()
    bias = bias if (bias is not None and rank == 0) else None
    c = c if (c is not None and rank == 0) else None
    # beta = beta if (c is not None and rank == 0) else 0.
    param = [cncl_comm, a, a_scale, b, b_scale, dtype, bias, c, 1.0, 1.0, block_m]
    output = ops.smooth_quant_matmul_allreduce(*param)
    if rank == 0:
        assertTensorsEqual(output.cpu().float(), pt_output.cpu().float(), 0.006, use_MSE=True, use_RAE=True)
    cleanup()

class TestGptQuantMatmulOp(BtTestCase):
    def op_impl_base(self, *args):
        return super().op_impl_base(*args)

    # weight only, no bias, no residual, no activation
    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_single_weight_only_matmul_allreduce due to ASan issues")
    def test_single_weight_only_matmul_allreduce(self):
        device_n = torch_mlu.mlu.device_count()
        if device_n < 2:
            print(f"device count is {device_n}, can not test communication between devices.", flush=True)
        block_m = 8
        M, K, N = 32, 256, 128
        has_bias = False
        has_res = False
        dtype = torch.half
        print("m={}, n={}, k={}, has_bias={}, has_res={}, dtype={}, tp_num={}, block_m={}, testing...".format(
            M, N, K, has_bias, has_res, dtype, device_n, block_m), flush=True)
        assert K % device_n == 0, f"K{K} must be divisible by tp_num{device_n}"
        param = [M, N, K, has_bias, has_res, dtype, block_m, device_n]
        mp.spawn(quant_mm_split_k, param, nprocs=device_n, join=True)

    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_single_sq_matmul_allreduce due to ASan issues")
    def test_single_sq_matmul_allreduce(self):
        device_n = torch_mlu.mlu.device_count()
        if device_n < 2:
            print(f"device count is {device_n}, can not test communication between devices.", flush=True)
        block_m = 1
        M, K, N = 2598, 1024, 1024
        has_bias = False
        has_res = False
        dtype = torch.half
        print("m={}, n={}, k={}, has_bias={}, has_res={}, dtype={}, tp_num={}, block_m={}, testing...".format(
            M, N, K, has_bias, has_res, dtype, device_n, block_m), flush=True)
        assert K % device_n == 0, f"K{K} must be divisible by tp_num{device_n}"
        param = [M, N, K, has_bias, has_res, dtype, device_n, block_m]
        mp.spawn(sq_mm_split_k, param, nprocs=device_n, join=True)

    @unittest.skip("not test")
    @unittest.skipIf('TMO_MEM_CHECK' in os.environ, "Skipping test_random_sq_mm_allreduce due to ASan issues")
    def test_random_sq_mm_allreduce(self):
        import random
        random.seed(0)
        device_n = torch_mlu.mlu.device_count()
        if device_n < 2:
            print(f"device count is {device_n}, can not test communication between devices.", flush=True)

        dtype_list = [torch.half, torch.bfloat16] if torch_mlu.mlu.is_bf16_supported() else [torch.half]
        for i in range(10):
            tp_num = random.randint(1, device_n)
            M = random.randint(1, 4096)
            N = random.choice([512, 1024, 2048, 4096])
            k_start = (1024 // tp_num) * tp_num
            K = random.randrange(k_start, 2048, tp_num)
            has_res = random.choice([False, True])
            has_bias = random.choice([False, True])
            dtype = random.choice(dtype_list)
            block_m = random.randint(1, 10)

            print("M={}, N={}, K={}, has_bias={}, has_res={}, dtype={}, tp_num={}, block_m={}, testing...".format(
                M, N, K, has_bias, has_res, dtype, tp_num, block_m), flush=True)
            param = [M, N, K, has_bias, has_res, dtype, tp_num, block_m]
            mp.spawn(sq_mm_split_k, param, nprocs=tp_num, join=True)

    def test_inductor(self):
        return super().test_inductor()

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