import os
os.environ['TMO_GEN_CASE'] = '0'

import sys
sys_args = sys.argv
sys.argv = [sys_args.pop(0)]  # prevent unittest printing help info

import copy
import argparse
import torch
import torch_mlu
import importlib
import random
from common_utils import create_op_param

def assert_tensor_equal(a: torch.Tensor, b: torch.Tensor, threshold):
    assert a.size() == b.size()
    a_ = a.cpu().reshape(-1).double()
    b_ = b.cpu().reshape(-1).double()
    nan_mask = a_ != a_
    assert torch.equal(nan_mask, b_ != b_), "tensor a and tensor b have different number of nan"
    diff = a_ - b_
    diff[nan_mask] = 0
    a_[nan_mask] = 0
    b_[nan_mask] = 0
    eps = 1e-10
    diff1 = diff.abs().sum() / (a_.abs().sum() + eps)
    diff2 = torch.sqrt((diff**2).sum() / ((a_**2).sum() + eps))
    print(f"[torch_mlu_ops] diff1: {diff1}, diff2: {diff2}")
    assert diff1 <= threshold, f"diff1: {diff1} <= threshold: {threshold}"
    assert diff2 <= threshold, f"diff2: {diff2} <= threshold: {threshold}"

def check_value_equal(x, y, threshold):
    assert type(x) == type(y)
    if type(x) is torch.Tensor:
        assert_tensor_equal(x, y, threshold)
    elif type(x) is list or type(x) is tuple:
        for i in range(len(x)):
            check_value_equal(x[i], y[i], threshold)
    elif x is not None:
        assert x == y

def check_equal(a, b, threshold):
    assert type(a) == type(b)
    if type(a) is tuple or type(a) is list:
        assert len(a) == len(b)
        for x, y in zip(a, b):
            check_value_equal(x, y, threshold)
    else:
        check_value_equal(a, b, threshold)

def get_base_obj(module_name, class_name):
    try:
        mod_ = importlib.import_module(module_name)
        cls_ = getattr(mod_, class_name)
        return cls_()
    except ImportError as e:
        print(f"Failed to import class '{class_name}' from module '{module_name}': {e}")
        return None
    except AttributeError as e:
        print(f"Module '{module_name}' does not have a class named '{class_name}': {e}")
        return None

def get_tmo_func(func_name):
    import torch_mlu_ops as tmo
    return getattr(tmo, func_name)


op_map = {
    "active": ["test_active", "TestActive", 0.004],
    "apply_rotary": ["test_apply_rotary", "TestApplyRotaryOp", 0.003],
    "attention_project": ["test_attn_proj", "TestAttnProjOp", 0.003],
    "batch_matmul": ["test_batch_matmul", "TestBatchMatMulOp", 0.004],
    "copy_blocks": ["test_copy_blocks", "TestCopyBlocksOp", 0],
    "dequant_from_linear_cache": ["test_dequant_from_linear_cache", "TestDequantFromLinearCache", 0.001],
    "dequant_from_paged_cache": ["test_dequant_from_paged_cache", "TestDequantFromPagedCache", 0.001],
    "ffn": ["test_ffn", "TestFFNOp", 0.005],
    "flash_attention": ["test_flash_attention", "TestFlashAttnOp", 0.005],
    # "flash_attn_sq_mm_allreduce": ["test_flash_attn_sq_mm_allreduce", "TestFlashAttnSqMMAllreduce", 0.003],
    "fused_layer_norm": ["test_fused_layernorm", "TestFuseLayerNormOp", 0.003],
    "fused_moe": ["test_moe", "TestFusedMOEOp", 0.006],
    "fused_norm_attention_project": ["test_fused_attn_proj", "TestFusedNormAttnProjOp", 0.003],
    "fused_norm_residual_ffn": ["test_fused_ffn", "TestFusedNormResidualFFNoP", 0.003],
    "fused_rms_norm": ["test_fused_rmsnorm", "TestFuseRmsNormOp", 0.0032],
    "fused_rope": ["test_fused_rope", "TestFusedRopeOp", 0.003],
    "group_gemm": ["test_group_gemm", "TestGroupGemmOp", 0.006],
    "matmul": ["test_matmul", "TestMatMulOp", 0.003],
    # "matmul_allreduce": ["test_matmul_all_reduce", "TestMatMulAllReduceOp", 0.006],
    "moe_active": ["test_moe_add_bias_activation", "TestMoeActiveKernel", 0.003],
    "moe_cast_gating": ["test_moe_cast_gating", "TestMoeCastGating", 0.0001],
    "moe_combine_result": ["test_moe_combine_result", "TestCombineResult", 0.003],
    "moe_expand_input": ["test_moe_expand_input", "TestExpandInput", 0],
    "moe_gen_idx": ["test_moe_gen_idx", "TestGenIdx", 0],
    "moe_quantize": ["test_smooth_quant", "TestSmoothQuantOp", 0.01],
    "moe_softmax_topk": ["test_moe_softmax_topk", "TestSoftmaxTopkOp", 0.003],
    "offline_quant_to_linear_cache": ["test_offline_quant_to_linear_cache", "TestOfflineQuantToLinearCache", 0.03],
    "offline_quant_to_paged_cache": ["test_offline_quant_to_paged_cache", "TestOfflineQuantToPagedCache", 0.03],
    "per_token_smooth_quantize": ["test_per_token_smooth_quantize", "TestPerTokenSmoothQuantizeOp", 0.003],
    # "preload": ["test_preload", "TestPreloadOp", 1],
    "quant_to_linear_cache": ["test_quant_to_linear_cache", "TestQuantToLinearCache", 0.003],
    "quant_to_paged_cache": ["test_quant_to_paged_cache", "TestQuantToPagedCache", 0.009],
    "quantize": ["test_quantize", "TestQuantizeOp", 0.003],
    "reshape_linear_cache": ["test_reshape_linear_cache", "TestReshapeLinearCache", 0],
    "reshape_paged_cache": ["test_reshape_paged_cache", "TestReshapePagedCacheOp", 0],
    "single_query_cached_kv_attn": ["test_single_query_cached_kv_attn", "TestSingleQueryAttnOp", 0.003],
    "single_query_mixed_cached_kv_attn": ["test_single_query_mixed_cached_kv_attn", "TestSingleQueryMixedKVAttnOp", 0.003],
    "smooth_quant_group_gemm": ["test_smooth_quant_group_gemm", "TestSmoothQuantGroupGemmOp", 0.006],
    "smooth_quant_matmul": ["test_smooth_quant_matmul", "TestSmoothQuantMatmulOp", 0.006],
    # "smooth_quant_matmul_allreduce": ["test_quant_matmul_all_reduce", "TestGptQuantMatmulOp", 0.006],
    "swap_blocks": ["test_swap_blocks", "TestSwapBlocksOp", 0],
    "update_out_and_lse": ["test_update_out_and_lse", "TestUpdateOutAndLse", 0.005],
    "weight_only_quant_matmul": ["test_weight_only_quant_matmul", "TestWeightOnlyQuantMatmulOp", 0.004],
}


def run_case(pt_case):
    op_name = pt_case.pop('op')
    op_obj = get_base_obj(op_map[op_name][0], op_map[op_name][1])
    if hasattr(op_obj, "run_gen_case"):
        op_obj.run_gen_case(pt_case)
    else:
        dump_data = pt_case.pop('dump_data')
        if dump_data:
            params = pt_case
        else:
            params = dict()
            for k,v in pt_case.items():
                params[k] = create_op_param(v)
        params_bak = copy.deepcopy(params)
        result_tmo = get_tmo_func(op_name)(**params)
        result_base = op_obj.op_impl_base(*params_bak.values())
        check_equal(result_tmo, result_base, op_map[op_name][-1])

def main():
    random.seed(0)
    torch.manual_seed(0)
    parser = argparse.ArgumentParser()
    parser.add_argument('--case_path', required=True, type=str, help='specify the case path')
    parser.add_argument('--detail', action="store_true", help="show content of pt file")
    args = parser.parse_args(args=sys_args)
    pt_case = torch.load(args.case_path)
    if args.detail:
        for k,v in pt_case.items():
            print(f"{k}: {v}")
        exit(0)
    print(f"[torch_mlu_ops] run {args.case_path} ...")
    run_case(pt_case)
    print(f"[torch_mlu_ops] run {args.case_path} successfully")

if __name__ == "__main__":
    main()
