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

import pytest
import torch
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

import os

# 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"))
ep_size = int(os.environ.get("EP_SIZE", "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]
EP_SIZE = [ep_size]
TOP_KS = [8]
N = 2048 if ep_size > 1 else 256

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


# 根据训练模式选择不同的批次大小参数
if training_config["training_mode"] == 0:
    # 常规运行模式，使用所有批次大小
    BATCH_SIZES = [1,2,4,8,16,24,32,64,128,256,512,1024,2048,4096,8192,16384,32768]
    # BATCH_SIZES = [1,2,4,8,16,24,32,64,128,256,512]
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

@pytest.mark.parametrize("m", BATCH_SIZES)
# @pytest.mark.parametrize("m", [4])
@pytest.mark.parametrize("n", [N])
@pytest.mark.parametrize("k", [7168])
@pytest.mark.parametrize("e", NUM_EXPERTS)
@pytest.mark.parametrize("topk", TOP_KS)
@pytest.mark.parametrize("ep_size", EP_SIZE)
@pytest.mark.parametrize("dtype", [torch.float16])
@pytest.mark.parametrize("group_size", [64])
@pytest.mark.parametrize("has_zp", [True])
@pytest.mark.parametrize("weight_bits", [4])
def test_fused_moe_wn16(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
    print(m, n, k, e, ep_size, topk, dtype, group_size, has_zp, weight_bits)

    # 设置CUDA设备
    if not torch.cuda.is_available():
        pytest.skip("CUDA is not available")

    # 根据配置设置设备
    if training_config["training_mode"] == 0:
        device_id = training_config["default_device"]
        if device_id is not None:
            torch.cuda.set_device(device_id)
            print(f"!!!Set device to {device_id}")
    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}")

    # 获取当前设备
    device = "cuda"
    # 打印当前测试用例
    global case_counter_x  # 添加全局声明
    print("\nTest Count {}: BATCH_SIZE={}, HIDDEN_SIZE={}, FFN_HIDDEN_SIZE={}, NUM_EXPERTS={}, TOP_K={}, EP_SIZE={}\n".format(
        case_counter_x,
        m,
        n,
        k,
        e,
        topk,
        ep_size
    ))
    case_counter_x += 1  # 递增计数器

    torch.manual_seed(0)
    a = torch.randn((m, k), device=device, dtype=dtype) / 10
    score = torch.randn((m, e), device=device, dtype=dtype)
    e_ep = e // ep_size

    if weight_bits == 4:
        pack_factor = 2
        quant_type = scalar_types.uint4 if has_zp else scalar_types.uint4b8
    elif weight_bits == 8:
        pack_factor = 1
        quant_type = scalar_types.uint8 if has_zp else scalar_types.uint8b128

    # 预分配量化后的张量（保持不变）
    w1_qweight = torch.empty((e_ep, 2 * n, k // pack_factor),
                            device=device, dtype=torch.uint8)
    w2_qweight = torch.empty((e_ep, k, n // pack_factor),
                            device=device, dtype=torch.uint8)
    w1_scales = torch.empty((e_ep, 2 * n, k // group_size),
                            device=device, dtype=dtype)
    w2_scales = torch.empty((e_ep, k, n // group_size),
                            device=device, dtype=dtype)
    w1_qzeros = torch.empty((e_ep, 2 * n // pack_factor, k // group_size),
                            device=device, dtype=torch.uint8)
    w2_qzeros = torch.empty((e_ep, k // pack_factor, n // group_size),
                            device=device, dtype=torch.uint8)

    # 关键优化：逐个专家处理，避免同时保存所有原始权重
    for expert_id in range(e_ep):
        # 1. 生成并处理w1专家权重
        w1_expert = torch.randn((2 * n, k), device=device, dtype=dtype) / 10

        weight, qweight, scales, qzeros = quantize_weights(
            w1_expert.T, quant_type, group_size, has_zp, False)
        del w1_expert  # 立即释放原始权重

        weight = weight.T
        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, :]

        w1_qweight[expert_id] = qweight
        w1_scales[expert_id] = scales
        if has_zp:
            w1_qzeros[expert_id] = qzeros

        # 释放量化过程中的临时变量
        del weight, qweight, scales, qzeros
        torch.cuda.empty_cache()

        # 2. 生成并处理w2专家权重
        w2_expert = torch.randn((k, n), device=device, dtype=dtype) / 10

        weight, qweight, scales, qzeros = quantize_weights(
            w2_expert.T, quant_type, group_size, has_zp, False)
        del w2_expert  # 立即释放原始权重

        weight = weight.T
        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, :]

        w2_qweight[expert_id] = qweight
        w2_scales[expert_id] = scales
        if has_zp:
            w2_qzeros[expert_id] = qzeros

        # 释放量化过程中的临时变量
        del weight, qweight, scales, qzeros
        torch.cuda.empty_cache()
    torch.cuda.synchronize()

    with patched_environment():
        topk_weights, topk_ids = fused_topk(a, score, topk, renormalize=True)
        if ep_size > 1:
            # to make sure at least one topk is in local experts
            local_e = e // ep_size
            begin_e_id = (topk_ids[0, 2] // local_e * local_e).item()
            e_ids = torch.arange(begin_e_id, begin_e_id + local_e, device="cuda", 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)
        else:
            e_map = None
        triton_output = fused_moe_module.fused_experts_impl(a,
                                w1_qweight,
                                w2_qweight,
                                topk_weights,
                                topk_ids,
                                a.dtype,
                                use_int4_w4a16=weight_bits == 4,
                                use_int8_w8a16=weight_bits == 8,
                                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])


# 定义MoE测试用例列表
moe_perf_model_cases_list = [
    # # 'BATCH_SIZE, HIDDEN_SIZE, FFN_HIDDEN_SIZE, NUM_EXPERTS, TOP_K'
    (1,     N, 7168, 256, 8),
    (2,     N, 7168, 256, 8),
    (3,     N, 7168, 256, 8),
    (4,     N, 7168, 256, 8),
    (8,     N, 7168, 256, 8),
    (5,     N, 7168, 256, 8),
    (6,     N, 7168, 256, 8),
    (7,     N, 7168, 256, 8),
    (8,     N, 7168, 256, 8),
    (9,     N, 7168, 256, 8),
    (10,    N, 7168, 256, 8),
    (11,    N, 7168, 256, 8),
    (12,    N, 7168, 256, 8),
    (13,    N, 7168, 256, 8),
    (14,    N, 7168, 256, 8),
    (15,    N, 7168, 256, 8),
    (16,    N, 7168, 256, 8),
    (17,    N, 7168, 256, 8),
    (18,    N, 7168, 256, 8),
    (19,    N, 7168, 256, 8),
    (20,    N, 7168, 256, 8),
    (21,    N, 7168, 256, 8),
    (22,    N, 7168, 256, 8),
    (23,    N, 7168, 256, 8),
    (24,    N, 7168, 256, 8),
    (25,    N, 7168, 256, 8),
    (26,    N, 7168, 256, 8),
    (27,    N, 7168, 256, 8),
    (28,    N, 7168, 256, 8),
    (29,    N, 7168, 256, 8),
    (30,    N, 7168, 256, 8),
    (31,    N, 7168, 256, 8),
    (32,    N, 7168, 256, 8),
    (64,    N, 7168, 256, 8),
    (128,   N, 7168, 256, 8),
    (256,   N, 7168, 256, 8),
    (512,   N, 7168, 256, 8),
    (1024,  N, 7168, 256, 8),
    (2048,  N, 7168, 256, 8),
    (4096,  N, 7168, 256, 8),
    (8192,  N, 7168, 256, 8),
    (16384, N, 7168, 256, 8),
    (32768, N, 7168, 256, 8),
]

# 配置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 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_wn16(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

    # 根据配置设置设备
    if training_config["training_mode"] == 0:
        device_id = training_config["default_device"]
        if device_id is not None:
            torch.cuda.set_device(device_id)
            print(f"!!!Set device to {device_id}")
    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}")

    # 获取当前设备
    device = "cuda"

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

    # 准备输入数据
    e = NUM_EXPERTS
    m = BATCH_SIZE
    n = HIDDEN_SIZE
    k = FFN_HIDDEN_SIZE

    torch.manual_seed(0)
    a = torch.randn((m, k), device=device, dtype=dtype) / 10
    score = torch.randn((m, e), device=device, dtype=dtype)
    e_ep = e // ep_size

    if weight_bits == 4:
        pack_factor = 2
        quant_type = scalar_types.uint4 if has_zp else scalar_types.uint4b8
    elif weight_bits == 8:
        pack_factor = 1
        quant_type = scalar_types.uint8 if has_zp else scalar_types.uint8b128

    # 预分配量化后的张量（保持不变）
    w1_qweight = torch.empty((e_ep, 2 * n, k // pack_factor),
                            device=device, dtype=torch.uint8)
    w2_qweight = torch.empty((e_ep, k, n // pack_factor),
                            device=device, dtype=torch.uint8)
    w1_scales = torch.empty((e_ep, 2 * n, k // group_size),
                            device=device, dtype=dtype)
    w2_scales = torch.empty((e_ep, k, n // group_size),
                            device=device, dtype=dtype)
    w1_qzeros = torch.empty((e_ep, 2 * n // pack_factor, k // group_size),
                            device=device, dtype=torch.uint8)
    w2_qzeros = torch.empty((e_ep, k // pack_factor, n // group_size),
                            device=device, dtype=torch.uint8)

    # 关键优化：逐个专家处理，避免同时保存所有原始权重
    for expert_id in range(e_ep):
        # 1. 生成并处理w1专家权重
        w1_expert = torch.randn((2 * n, k), device=device, dtype=dtype) / 10

        weight, qweight, scales, qzeros = quantize_weights(
            w1_expert.T, quant_type, group_size, has_zp, False)
        del w1_expert  # 立即释放原始权重

        weight = weight.T
        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, :]

        w1_qweight[expert_id] = qweight
        w1_scales[expert_id] = scales
        if has_zp:
            w1_qzeros[expert_id] = qzeros

        # 释放量化过程中的临时变量
        del weight, qweight, scales, qzeros
        torch.cuda.empty_cache()

        # 2. 生成并处理w2专家权重
        w2_expert = torch.randn((k, n), device=device, dtype=dtype) / 10

        weight, qweight, scales, qzeros = quantize_weights(
            w2_expert.T, quant_type, group_size, has_zp, False)
        del w2_expert  # 立即释放原始权重

        weight = weight.T
        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, :]

        w2_qweight[expert_id] = qweight
        w2_scales[expert_id] = scales
        if has_zp:
            w2_qzeros[expert_id] = qzeros

        # 释放量化过程中的临时变量
        del weight, qweight, scales, qzeros
        torch.cuda.empty_cache()
    torch.cuda.synchronize()

    if provider == "triton":
        with patched_environment():
            topk_weights, topk_ids = fused_topk(a, score, TOP_K, renormalize=True)
            if ep_size > 1:
                # to make sure at least one topk is in local experts
                local_e = e // ep_size
                begin_e_id = (topk_ids[0, 2] // local_e * local_e).item()
                e_ids = torch.arange(begin_e_id, begin_e_id + local_e, device="cuda", dtype=torch.int32)
                e_map = torch.full((e, ), -1, 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)
            else:
                e_map = None
            fn = lambda: fused_moe_module.fused_experts_impl(a,
                                  w1_qweight,
                                  w2_qweight,
                                  topk_weights,
                                  topk_ids,
                                  a.dtype,
                                  use_int4_w4a16=weight_bits == 4,
                                  use_int8_w8a16=weight_bits == 8,
                                  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])
            ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep, perf_profiling=perf_profling)

    return ms

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