# SPDX-License-Identifier: Apache-2.0
"""Tests for the MOE layers.

Run `pytest tests/kernels/test_moe.py`.
"""
import pytest
import torch

import triton
import triton.language as tl
import time
import gc
import os
import sys
from contextlib import contextmanager

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '../')))

import aiter.ops.triton.fused_moe as fused_moe_module
from aiter.fused_moe import fused_topk
from aiter import per_token_quant_hip, per_block_quant_wrapper
from op_tests.utility.scalar_type import ScalarType, scalar_types
from op_tests.utility.utils import quantize_weights

try:
    from op_tests.triton_autotune.fused_moe.autotune_patches import patched_environment
except ImportError:

    @contextmanager
    def patched_environment():
        yield

# os.environ["VLLM_ENABLE_MOE_ALIGN_BLOCK_SIZE_TRITON"] = "1"
# os.environ["TRITON_PRINT_AUTOTUNING"] = "1"
os.environ["AMDGCN_USE_BUFFER_OPS"] = "1"

compile_only  = int(os.environ.get("TRITON_COMPILE_ONLY", "0"))
num_groups_compile = int(os.environ.get("NUM_GROUPS_COMPILE", "1"))
training_mode = int(os.environ.get("TRAINING_MODE", "0"))
benchmark_mode = int(os.environ.get("BENCHMARK_MODE", "0"))
cur_group = int(os.environ.get("CUR_GROUP", "0"))
fast_mode = int(os.environ.get("FAST_MODE", "0"))

if compile_only == 1:
    assert num_groups_compile <= 56
    # batch 1 and 64 is enough for us to gen all config kernels when compile_only
    batchsToCompile = [1, 64]
    batch_size_list = [batchsToCompile.copy() for _ in range(num_groups_compile)]
    groups_per_devices = 8
    device_ids = [device_id for _ in range(groups_per_devices)
                for device_id in range((num_groups_compile + groups_per_devices - 1) // groups_per_devices)]

    group_ids = {i: (device_ids[i], batch_size_list[i]) for i in range(num_groups_compile)}
    print(f"{group_ids=}")
else:
    # 7 groups
    if fast_mode == 0:
        group_ids = {
            # (device_id, [batch_size list])
            0: (1, [1, 2, 4, 8, 16, 24, 32, 64]),
            1: (2, [128, 256, 512]),
            2: (3, [1, 2, 4, 8, 16, 24, 32, 64]),
            3: (4, [128, 256, 512]),
            4: (5, [2048]),
            5: (6, [16384]),
            6: (7, [32768]),
            7: (4, [1, 2, 4, 8, 16, 24, 32, 64, 128, 256, 512]),
            8: (5, [1024, 2048, 4096, 8192]),
            9: (6, [16384]),
            10: (7, [32768]),
        }
    else:
        group_ids = {
            # (device_id, [batch_size list])
            0: (1, [1, 2, 4]),
            1: (2, [8, 16]),
            2: (3, [24]),
            3: (4, [32]),
            4: (5, [64]),
            5: (6, [128]),
            6: (7, [256]),
            7: (4, [1, 2, 4, 8, 16, 24, 32]),
            8: (5, [64]),
            9: (6, [128]),
            10: (7, [256]),
        }

# 定义全局配置变量
training_config = {
    "training_mode"     : training_mode,        # 0: 运行模式, 1: 多卡训练模式
    "cur_group"         : cur_group,            # 0: 第一组, 1: 第二组, 2: 第三组
    "default_device"    : 0,                 # 默认设备ID, or None
    "group_id"          : group_ids,
}

NUM_EXPERTS = [256]
HIDDEN_SIZES = [256]
FFN_HIDDEN_SIZES = [6144]
EP_SIZE = [1]
TOP_KS = [8]
DTYPES = [torch.float16]
GROUP_SIZES = [64]
HAS_ZPS = [True]
WEIGHT_BITS = [4]

print(f"[tune_moe_int4int8.py] Running with TRAINING_MODE={training_mode}, CUR_GROUP={cur_group}")


# 根据训练模式选择不同的批次大小参数
if training_config["training_mode"] == 0:
    # Preserve the original small-batch correctness sweep.
    # Larger batches may hit invalid default int4_w4a8 kernel configs when
    # no tuned config file is present for the current device.
    BATCH_SIZES = [4, 8]
else:
    # 多卡训练模式，根据当前分组选择批次大小
    if training_config["cur_group"] in training_config["group_id"]:
        BATCH_SIZES = training_config["group_id"][training_config["cur_group"]][1]
        print(f"!!!BATCH_SIZES: {BATCH_SIZES}")
    else:
        raise ValueError(f"Invalid training group: {training_config['cur_group']}")


case_counter_x = 1


def setup_device_from_training_config(skip_if_unavailable: bool = False) -> int:
    if not torch.cuda.is_available():
        if skip_if_unavailable:
            pytest.skip("CUDA is not available")
        raise RuntimeError("CUDA is not available")

    if training_config["training_mode"] == 0:
        device_id = training_config["default_device"]
    else:
        device_id = training_config["group_id"][training_config["cur_group"]][0]

    if device_id is not None:
        torch.cuda.set_device(device_id)
        print(f"!!!Set device to {device_id}")
    return device_id if device_id is not None else torch.cuda.current_device()


def get_quant_config(weight_bits: int, has_zp: bool) -> tuple[int, ScalarType]:
    if weight_bits == 4:
        return 2, scalar_types.uint4 if has_zp else scalar_types.uint4b8
    if weight_bits == 8:
        return 1, scalar_types.uint8 if has_zp else scalar_types.uint8b128
    raise ValueError(f"Unsupported weight_bits: {weight_bits}")


def allocate_quantized_weights(
    e: int,
    n: int,
    k: int,
    pack_factor: int,
    group_size: int,
    dtype: torch.dtype,
    device: str,
):
    w1_qweight = torch.empty((e, 2 * n, k // pack_factor), device=device, dtype=torch.uint8)
    w2_qweight = torch.empty((e, k, n // pack_factor), device=device, dtype=torch.uint8)
    w1_scales = torch.empty((e, 2 * n, k // group_size), device=device, dtype=dtype)
    w2_scales = torch.empty((e, k, n // group_size), device=device, dtype=dtype)
    w1_qzeros = torch.empty((e, 2 * n // pack_factor, k // group_size), device=device, dtype=torch.uint8)
    w2_qzeros = torch.empty((e, k // pack_factor, n // group_size), device=device, dtype=torch.uint8)
    return w1_qweight, w2_qweight, w1_scales, w2_scales, w1_qzeros, w2_qzeros


def quantize_weight_tensor(
    weight: torch.Tensor,
    quant_type: ScalarType,
    group_size: int,
    has_zp: bool,
    weight_bits: int,
):
    _, qweight, scales, qzeros = quantize_weights(weight.T, quant_type, group_size, has_zp, False)
    qweight = qweight.T.contiguous().to(torch.uint8)
    scales = scales.T
    if has_zp:
        qzeros = qzeros.T.contiguous().to(torch.uint8)
    if weight_bits == 4:
        qweight = qweight[:, 1::2] * 16 + qweight[:, ::2]
        if has_zp:
            qzeros = qzeros[1::2, :] * 16 + qzeros[::2, :]
    return qweight, scales, qzeros if has_zp else None


def quantize_dense_weights(
    w1: torch.Tensor,
    w2: torch.Tensor,
    quant_type: ScalarType,
    group_size: int,
    has_zp: bool,
    weight_bits: int,
    dtype: torch.dtype,
    device: str,
):
    e, _, k = w1.shape
    n = w2.shape[-1]
    pack_factor, _ = get_quant_config(weight_bits, has_zp)
    w1_qweight, w2_qweight, w1_scales, w2_scales, w1_qzeros, w2_qzeros = allocate_quantized_weights(
        e, n, k, pack_factor, group_size, dtype, device
    )

    for expert_id in range(e):
        qweight, scales, qzeros = quantize_weight_tensor(
            w1[expert_id], quant_type, group_size, has_zp, weight_bits
        )
        w1_qweight[expert_id] = qweight
        w1_scales[expert_id] = scales
        if has_zp:
            w1_qzeros[expert_id] = qzeros

        qweight, scales, qzeros = quantize_weight_tensor(
            w2[expert_id], quant_type, group_size, has_zp, weight_bits
        )
        w2_qweight[expert_id] = qweight
        w2_scales[expert_id] = scales
        if has_zp:
            w2_qzeros[expert_id] = qzeros

    return w1_qweight, w2_qweight, w1_scales, w2_scales, w1_qzeros, w2_qzeros


def quantize_streaming_weights(
    e: int,
    n: int,
    k: int,
    quant_type: ScalarType,
    group_size: int,
    has_zp: bool,
    weight_bits: int,
    dtype: torch.dtype,
    device: str,
):
    pack_factor, _ = get_quant_config(weight_bits, has_zp)
    w1_qweight, w2_qweight, w1_scales, w2_scales, w1_qzeros, w2_qzeros = allocate_quantized_weights(
        e, n, k, pack_factor, group_size, dtype, device
    )

    for expert_id in range(e):
        w1_expert = torch.randn((2 * n, k), device=device, dtype=dtype) / 10
        qweight, scales, qzeros = quantize_weight_tensor(
            w1_expert, quant_type, group_size, has_zp, weight_bits
        )
        w1_qweight[expert_id] = qweight
        w1_scales[expert_id] = scales
        if has_zp:
            w1_qzeros[expert_id] = qzeros
        del w1_expert, qweight, scales, qzeros

        w2_expert = torch.randn((k, n), device=device, dtype=dtype) / 10
        qweight, scales, qzeros = quantize_weight_tensor(
            w2_expert, quant_type, group_size, has_zp, weight_bits
        )
        w2_qweight[expert_id] = qweight
        w2_scales[expert_id] = scales
        if has_zp:
            w2_qzeros[expert_id] = qzeros
        del w2_expert, qweight, scales, qzeros

        torch.cuda.empty_cache()

    torch.cuda.synchronize()
    return w1_qweight, w2_qweight, w1_scales, w2_scales, w1_qzeros, w2_qzeros


def build_expert_map(
    e: int,
    ep_size: int,
    device: str,
    *expert_tensors: torch.Tensor,
):
    if ep_size <= 1:
        return None, expert_tensors

    local_e = e // ep_size
    e_ids = torch.randint(0, e, (local_e,), device=device, dtype=torch.int32)
    e_map = torch.full((e,), -1, device=device, dtype=torch.int32)
    e_map[e_ids] = torch.arange(local_e, device=device, dtype=torch.int32)
    return e_map, tuple(tensor[e_ids] for tensor in expert_tensors)


def input_helper(
    m: int,
    n: int,
    k: int,
    e: int,
    topk: int,
    ep_size: int,
    dtype: torch.dtype,
    group_size: int,
    has_zp: bool,
    weight_bits: int,
    device: str = "cuda",
    seed: int = 0,
    streaming_weights: bool = False,
    renormalize: bool = False,
):
    assert ep_size >= 1
    assert e % ep_size == 0, f"num_experts({e}) must be divisible by ep_size({ep_size})"

    torch.manual_seed(seed)
    a = torch.randn((m, k), device=device, dtype=dtype) / 10
    score_all = torch.randn((m, e), device=device, dtype=dtype)

    _, quant_type = get_quant_config(weight_bits, has_zp)

    if streaming_weights:
        quantized_weights = quantize_streaming_weights(
            e, n, k, quant_type, group_size, has_zp, weight_bits, dtype, device
        )
    else:
        w1 = torch.randn((e, 2 * n, k), device=device, dtype=dtype) / 10
        w2 = torch.randn((e, k, n), device=device, dtype=dtype) / 10
        quantized_weights = quantize_dense_weights(
            w1, w2, quant_type, group_size, has_zp, weight_bits, dtype, device
        )

    e_map, quantized_weights = build_expert_map(e, ep_size, device, *quantized_weights)
    topk_weights, topk_ids = fused_topk(a, score_all, topk, renormalize=renormalize)
    a_q, a_scale = per_block_quant_wrapper((1, group_size))(per_token_quant_hip)(a)
    w1_qweight, w2_qweight, w1_scales, w2_scales, w1_qzeros, w2_qzeros = quantized_weights
    return (
        a_q,
        w1_qweight,
        w2_qweight,
        w1_scales,
        w2_scales,
        score_all,
        e_map,
        topk_weights,
        topk_ids,
        a_scale,
        w1_qzeros,
        w2_qzeros,
    )

@pytest.mark.parametrize("m", BATCH_SIZES)
@pytest.mark.parametrize("n", HIDDEN_SIZES)
@pytest.mark.parametrize("k", FFN_HIDDEN_SIZES)
@pytest.mark.parametrize("e", NUM_EXPERTS)
@pytest.mark.parametrize("topk", TOP_KS)
@pytest.mark.parametrize("ep_size", EP_SIZE)
@pytest.mark.parametrize("dtype", DTYPES)
@pytest.mark.parametrize("group_size", GROUP_SIZES)
@pytest.mark.parametrize("has_zp", HAS_ZPS)
@pytest.mark.parametrize("weight_bits", WEIGHT_BITS)
@torch.inference_mode()
def test_fused_moe_int4int8(m: int, n: int, k: int, e: int, topk: int,
                            ep_size: int, dtype: torch.dtype, group_size: int,
                            has_zp: bool, weight_bits: int):
    assert ep_size >= 1
    if not torch.cuda.is_available():
        pytest.skip("CUDA is not available")

    setup_device_from_training_config(skip_if_unavailable=True)
    device = "cuda"

    ENABLE_TORCH_REF = False if m > 256 or training_config["training_mode"] == 1 else True

    global case_counter_x
    print(
        f"\nTest Count {case_counter_x}: m={m}, n={n}, k={k}, e={e}, topk={topk}, "
        f"ep_size={ep_size}, device={torch.cuda.current_device()}, dtype={dtype}, "
        f"group_size={group_size}, has_zp={has_zp}, weight_bits={weight_bits}, "
        f"ENABLE_TORCH_REF={ENABLE_TORCH_REF}!!!\n"
    )
    case_counter_x += 1

    M, N, K, E = m, n, k, e
    a_q, w1_qweight, w2_qweight, w1_scales, w2_scales, _score_all, e_map, topk_weights, topk_ids, a_scale, w1_qzeros, w2_qzeros = input_helper(
        m=M,
        n=N,
        k=K,
        e=E,
        topk=topk,
        ep_size=ep_size,
        dtype=dtype,
        group_size=group_size,
        has_zp=has_zp,
        weight_bits=weight_bits,
        device=device,
        seed=0,
        renormalize=False,
    )

    with patched_environment():
        fused_moe_module.fused_experts_impl(
            a_q,
            w1_qweight,
            w2_qweight,
            topk_weights,
            topk_ids,
            dtype,
            use_int4_w4a8=True,
            global_num_experts=E,
            expert_map=e_map,
            w1_scale=w1_scales,
            w2_scale=w2_scales,
            w1_zp=w1_qzeros if has_zp else None,
            w2_zp=w2_qzeros if has_zp else None,
            block_shape=[0, group_size],
            a1_scale=a_scale,
        )

# 定义MoE测试用例列表
moe_perf_model_cases_list = [
    (m, n, k, e, topk)
    for m in [1, 2, 4, 8, 16, 24, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768]
    for n in HIDDEN_SIZES
    for k in FFN_HIDDEN_SIZES
    for e in NUM_EXPERTS
    for topk in TOP_KS
]

# 配置MoE基准测试参数
moe_configs = [
    triton.testing.Benchmark(
        x_names=['BATCH_SIZE', 'HIDDEN_SIZE', 'FFN_HIDDEN_SIZE', 'NUM_EXPERTS', 'TOP_K'],
        x_vals=moe_perf_model_cases_list,
        line_arg='provider',  # 用于区分不同实现

        # line_vals=['triton', 'torch'],  # 测试的实现
        # line_names=['Triton', 'Torch'],  # 图例名称
        # styles=[('red', '-'), ('blue', '--')],
        line_vals=['triton'],  # 测试的实现
        line_names=['Triton'],  # 图例名称
        styles=[('red', '-')],

        ylabel='Time(ms)',
        xlabel='Batch Size',
        plot_name='MoE INT4INT8 Performance (Times ms)',  # 图表标题
        args={
            'dtype': torch.float16,
            'device': 'cuda',
            'group_size': 64,
            'weight_bits': 4,
            'has_zp': True
        }
    )
]


case_counter = 1
perf_profling = False

@triton.testing.perf_report(moe_configs)
def bench_fused_moe_int4int8(BATCH_SIZE, HIDDEN_SIZE, FFN_HIDDEN_SIZE, NUM_EXPERTS, TOP_K,
                             provider, dtype=torch.float16, device="cuda",
                             group_size=64, weight_bits=4, has_zp=True):
    """
    BATCH_SIZE: 批次大小 (m)
    HIDDEN_SIZE: 隐藏层大小 (n)
    FFN_HIDDEN_SIZE: FFN中间层大小 (k)
    NUM_EXPERTS: 专家数量 (e)
    TOP_K: 每个token选择的专家数量
    """
    warmup = 25
    rep = 100

    setup_device_from_training_config()
    device = "cuda"

    # 打印当前测试用例
    global case_counter  # 添加全局声明
    print("\nTest Count {}: BATCH_SIZE={}, HIDDEN_SIZE={}, FFN_HIDDEN_SIZE={}, NUM_EXPERTS={}, TOP_K={}\n".format(
        case_counter,
        BATCH_SIZE,
        HIDDEN_SIZE,
        FFN_HIDDEN_SIZE,
        NUM_EXPERTS,
        TOP_K
    ))
    case_counter += 1
    ep_size = EP_SIZE[0]
    e = NUM_EXPERTS

    if perf_profling:
        assert provider == "triton"
        print("Before quantization:")
        print(f"Allocated memory: {torch.cuda.memory_allocated() / (1024**2):.2f} MB")
        print(f"Reserved memory: {torch.cuda.memory_reserved() / (1024**2):.2f} MB")

    M, N, K, E = BATCH_SIZE, HIDDEN_SIZE, FFN_HIDDEN_SIZE, e
    a_q, w1_qweight, w2_qweight, w1_scales, w2_scales, _score_all, e_map, topk_weights, topk_ids, a_scale, w1_qzeros, w2_qzeros = input_helper(
        m=M,
        n=N,
        k=K,
        e=E,
        topk=TOP_K,
        ep_size=ep_size,
        dtype=dtype,
        group_size=group_size,
        has_zp=has_zp,
        weight_bits=weight_bits,
        device=device,
        seed=0,
        streaming_weights=perf_profling,
        renormalize=True,
    )

    if perf_profling:
        torch.cuda.empty_cache()
        gc.collect()
        print("After quantization:")
        print(f"Allocated memory: {torch.cuda.memory_allocated() / (1024**2):.2f} MB")
        print(f"Reserved memory: {torch.cuda.memory_reserved() / (1024**2):.2f} MB")

    if provider == "triton":
        with patched_environment():
            fn = lambda: fused_moe_module.fused_experts_impl(
                a_q,
                w1_qweight,
                w2_qweight,
                topk_weights,
                topk_ids,
                dtype,
                use_int4_w4a8=True,
                global_num_experts=E,
                expert_map=e_map,
                w1_scale=w1_scales,
                w2_scale=w2_scales,
                w1_zp=w1_qzeros if has_zp else None,
                w2_zp=w2_qzeros if has_zp else None,
                block_shape=[0, group_size],
                a1_scale=a_scale,
            )
            ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep, perf_profiling=perf_profling)
    else:
        raise ValueError(f"Unsupported provider: {provider}")
    return ms

# # # 运行基准测试
if benchmark_mode == 1:
    bench_fused_moe_int4int8.run(print_data=True)
