# SPDX-License-Identifier: Apache-2.0

# Adapted from https://github.com/sgl-project/sglang/pull/3730
import itertools
from contextlib import contextmanager

import pytest
import torch
import triton
import triton.testing

from aiter import silu_and_mul
from aiter.fused_moe import fused_topk
from aiter import per_token_quant_hip
from aiter import dtypes
import aiter.ops.triton.fused_moe as fused_moe_module

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

    @contextmanager
    def patched_environment():
        yield

from aiter.test_common import checkAllclose
from typing import List, Dict, Optional

import os
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, 32768]),
            10: (7, [65536]),
        }
    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,
}

print(f"[tune_moe_int8_channel.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]
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']}")


def setup_device_from_training_config() -> int:
    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 native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16):
    """Matrix multiplication function that supports per-token input quantization and per-column weight quantization"""
    A = A.to(torch.float32)
    B = B.to(torch.float32)

    assert A.shape[-1] == B.shape[-1], "Dimension mismatch"
    assert B.ndim == 2 and B.is_contiguous(), "B must be a 2D contiguous tensor"

    # Reshape input
    M = A.numel() // A.shape[-1]
    B = B.t()  # Transpose weight matrix
    N, K = B.shape
    origin_C_shape = A.shape[:-1] + (K,)
    A = A.reshape(M, N)

    # As is per-token [M, 1], Bs is per-column [1, K]
    C = torch.matmul(A, B)  # [M, K]
    C = As * C * Bs.view(1, -1)  # Broadcast per-column scale

    return C.reshape(origin_C_shape).to(output_dtype)


def torch_w8a8_per_column_moe_expert_mask(
    a,
    w1,
    w2,
    w1_s,
    w2_s,
    topk_weight,
    topk_ids,
    expert_map=None,
):
    """This function performs fused moe with per-column int8 quantization using native torch."""
    B, D = a.shape
    topk = topk_weight.shape[1]
    # Perform per-token quantization
    a_q, a_s = per_token_quant_hip(a)
    # Repeat tokens to match topk
    a_q = a_q.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
    # Also repeat the scale
    a_s = a_s.view(B, -1, 1).repeat(1, topk, 1).reshape(-1, 1)  # [B*topk, 1]

    out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)

    topk_weight = topk_weight.view(-1)
    topk_ids = topk_ids.view(-1)
    if expert_map is not None:
        topk_ids = expert_map[topk_ids]

    # Process each expert
    for i in range(w1.shape[0]):
        mask = topk_ids == i
        if mask.sum():
            # First MLP layer: note that a_s is now per-token
            inter_out = native_w8a8_per_token_matmul(
                a_q[mask], w1[i], a_s[mask], w1_s[i], output_dtype=a.dtype
            )
            # Activation function
            act_out = torch.empty((inter_out.shape[0], inter_out.shape[1] // 2),
                                      dtype=a.dtype, device=a.device)
            silu_and_mul(act_out, inter_out)
            # Quantize activation output with per-token
            act_out_q, act_out_s = per_token_quant_hip(act_out)

            # Second MLP layer
            out[mask] = native_w8a8_per_token_matmul(
                act_out_q, w2[i], act_out_s, w2_s[i], output_dtype=a.dtype
            )
    # Apply routing weights and sum
    return (
        out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
    ).sum(dim=1)


def input_helper(
    m: int,
    n: int,
    k: int,
    e: int,
    topk: int,
    ep_size: int,
    dtype: torch.dtype,
    device: str = "cuda",
    num_shared_experts: int = 0,
    seed: int = 0,
    ep_id: Optional[int] = None,
):
    """Helper function to generate input data for INT8 per-channel quantization MoE tests."""
    assert e % ep_size == 0, f"num_experts({e}) must be divisible by ep_size({ep_size})"

    torch.manual_seed(seed)
    factor_for_scale = 1e-2
    int8_max = 127
    int8_min = -128
    e_sum = e + num_shared_experts
    final_topk = topk + num_shared_experts

    a = torch.randn((m, k), dtype=dtype, device=device) / 10
    e_per_rank = e // ep_size

    # 生成全局权重（bf16 格式，后续量化为 int8）
    w1_global_fp32 = (torch.rand((e_sum, 2 * n, k), dtype=torch.float32, device=device) - 0.5) * 2
    w2_global_fp32 = (torch.rand((e_sum, k, n), dtype=torch.float32, device=device) - 0.5) * 2
    w1_global = (w1_global_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
    w2_global = (w2_global_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)

    # Per-channel scale
    w1_s_global = torch.rand(e_sum, 2 * n, device=device) * factor_for_scale
    w2_s_global = torch.rand(e_sum, k, device=device) * factor_for_scale

    # 情况1: ep_size == 1, 无共享专家
    if ep_size == 1 and num_shared_experts == 0:
        score_all = torch.randn((m, e), dtype=dtype, device=device)
        topk_weights, topk_ids = fused_topk(a, score_all, topk, renormalize=False)
        e_map = None
        w1, w2, w1_s, w2_s = w1_global[:e], w2_global[:e], w1_s_global[:e], w2_s_global[:e]

    # 情况2: ep_size > 1, 无共享专家
    elif ep_size > 1 and num_shared_experts == 0:
        if ep_id is None:
            ep_id = ep_size - 1
        score_all = torch.randn((m, e), dtype=dtype, device=device)
        topk_weights, topk_ids = fused_topk(a, score_all, topk, renormalize=False)

        expert_mask = torch.zeros((e,), dtype=dtypes.i32, device=device)
        start = ep_id * e_per_rank
        end = (ep_id + 1) * e_per_rank
        expert_mask[start:end] = 1
        indices = expert_mask.cumsum(0, dtype=dtypes.i32) - 1
        e_map = torch.where(
            expert_mask == 0,
            torch.full_like(expert_mask, -1, dtype=dtypes.i32),
            indices,
        )
        w1 = w1_global[start:end]
        w2 = w2_global[start:end]
        w1_s = w1_s_global[start:end]
        w2_s = w2_s_global[start:end]

    # 情况3: ep_size == 1, 有共享专家
    elif ep_size == 1 and num_shared_experts != 0:
        import aiter

        score_all = torch.randn((m, e_sum), dtype=dtype, device=device)
        score = score_all[:, :e]
        correction_bias = torch.randn((e,), device=device, dtype=dtype)
        w_sglang = torch.empty_strided((m, final_topk), (final_topk + 10, 1), device=device, dtype=dtypes.fp32)
        id_sglang = torch.empty_strided((m, final_topk), (final_topk + 10, 1), device=device, dtype=dtypes.i32)
        scale_factor = 1.2
        num_expert_group = 8
        topk_group = 4
        gate_out = aiter.moe_fused_gate(
            score,
            correction_bias,
            w_sglang,
            id_sglang,
            num_expert_group,
            topk_group,
            final_topk,
            num_shared_experts,
            scale_factor,
        )
        topk_weights = gate_out[0]
        topk_ids = gate_out[1]
        topk_ids, sorted_idx = torch.sort(topk_ids)
        topk_weights = topk_weights.gather(1, sorted_idx)

        e_map = None
        w1, w2, w1_s, w2_s = w1_global, w2_global, w1_s_global, w2_s_global

    # 情况4: ep_size > 1, 有共享专家
    elif ep_size > 1 and num_shared_experts != 0:
        import aiter

        if ep_id is None:
            ep_id = ep_size - 1

        # shared experts 路径沿用 bf16 逻辑：构造 fake expert 作为 mask 占位
        expert_mask = torch.zeros((e + num_shared_experts + 1,), dtype=dtypes.i32, device=device)
        start = ep_id * e_per_rank
        end = (ep_id + 1) * e_per_rank
        expert_mask[start:end] = 1
        expert_mask[e:e_sum] = 1
        expert_mask[-1] = 0

        indices = expert_mask.cumsum(0, dtype=dtypes.i32) - 1
        e_map = torch.where(
            expert_mask == 0,
            torch.full_like(expert_mask, -1, dtype=dtypes.i32),
            indices,
        )

        # 本地非共享专家 + 全局共享专家
        w1_local = w1_global[start:end]
        w2_local = w2_global[start:end]
        w1_s_local = w1_s_global[start:end]
        w2_s_local = w2_s_global[start:end]
        w1_shared = w1_global[e:e_sum]
        w2_shared = w2_global[e:e_sum]
        w1_s_shared = w1_s_global[e:e_sum]
        w2_s_shared = w2_s_global[e:e_sum]
        w1 = torch.cat([w1_local, w1_shared], dim=0)
        w2 = torch.cat([w2_local, w2_shared], dim=0)
        w1_s = torch.cat([w1_s_local, w1_s_shared], dim=0)
        w2_s = torch.cat([w2_s_local, w2_s_shared], dim=0)

        score = torch.randn((m, e), dtype=dtype, device=device)
        score_all = torch.randn((m, e_sum), dtype=dtype, device=device)
        correction_bias = torch.randn((e,), device=device, dtype=dtype)
        w_sglang = torch.empty_strided((m, final_topk), (final_topk + 10, 1), device=device, dtype=dtypes.fp32)
        id_sglang = torch.empty_strided((m, final_topk), (final_topk + 10, 1), device=device, dtype=dtypes.i32)
        scale_factor = 1.2
        num_expert_group = 8
        topk_group = 4
        gate_out = aiter.moe_fused_gate(
            score,
            correction_bias,
            w_sglang,
            id_sglang,
            num_expert_group,
            topk_group,
            final_topk,
            num_shared_experts,
            scale_factor,
        )
        topk_weights = gate_out[0]
        topk_ids = gate_out[1]
        topk_ids, sorted_idx = torch.sort(topk_ids)
        topk_weights = topk_weights.gather(1, sorted_idx)

    else:
        raise AssertionError("Invalid input combination")

    return a, w1, w2, w1_s, w2_s, score_all, e_map, topk_weights, topk_ids

#deepseek model int8 channel
NUM_SHARED_EXPERTS = [0]
NUM_EXPERTS = [256]
HIDDEN_SIZES = [128]
FFN_HIDDEN_SIZES = [7168]
EP_SIZE = [ep_size]
TOP_KS = [8]
DTYPES = [torch.half]

#minimax m2.5
NUM_SHARED_EXPERTS = [0]
NUM_EXPERTS = [256]
HIDDEN_SIZES = [384] # N1/2 = K2 = HIDDEN_SIZES, N1 = 2*HIDDEN_SIZES = 768
FFN_HIDDEN_SIZES = [3072] # K1, N2
EP_SIZE = [ep_size]
TOP_KS = [8]
DTYPES = [torch.half]

case_counter_x = 1

@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("num_shared_experts", NUM_SHARED_EXPERTS)
@pytest.mark.parametrize("dtype", DTYPES)
@torch.inference_mode()
def test_w8a8_int8_channel_fused_moe(m, n, k, e, topk, ep_size, num_shared_experts, dtype):
    # 设置CUDA设备
    if not torch.cuda.is_available():
        pytest.skip("CUDA is not available")

    setup_device_from_training_config()

    # 获取当前设备
    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"num_shared_experts={num_shared_experts}, ep_size={ep_size}, "
        f"device={torch.cuda.current_device()}, dtype={dtype}, ENABLE_TORCH_REF={ENABLE_TORCH_REF}!!!\n"
    )
    case_counter_x += 1

    M, N, K, E = m, n, k, e
    a, w1, w2, w1_s, w2_s, _score_all, e_map, topk_weights, topk_ids = input_helper(
        m=M,
        n=N,
        k=K,
        e=E,
        topk=topk,
        ep_size=ep_size,
        dtype=dtype,
        device=device,
        num_shared_experts=num_shared_experts,
        seed=0,
    )

    with patched_environment():
        out = fused_moe_module.fused_experts_impl(
            a,
            w1,
            w2,
            topk_weights,
            topk_ids,
            a.dtype,
            use_int8_w8a8=True,
            per_channel_quant=True,
            global_num_experts=E + num_shared_experts,
            expert_map=e_map,
            w1_scale=w1_s,
            w2_scale=w2_s,
        )

    if ENABLE_TORCH_REF:
        ref_out = torch_w8a8_per_column_moe_expert_mask(
                        a, w1, w2, w1_s, w2_s, topk_weights, topk_ids,
                        expert_map=e_map
                    )
        rel_diff = (
            torch.mean(torch.abs(out.to(torch.float32) - ref_out.to(torch.float32)))
            / torch.mean(torch.abs(ref_out.to(torch.float32)))
        )
        assert rel_diff < 0.05, (
            f"[TRITON_check] {M=}, {K=}, {N=}, {E=}, {topk=}, "
            f"{num_shared_experts=}, {ep_size=}, dtype={dtype}, rel_diff={rel_diff}"
        )


# 定义MoE测试用例列表
moe_perf_model_cases_list = [
    (m, n, k, e, topk, num_shared_experts)
    for m in BATCH_SIZES
    for n in HIDDEN_SIZES
    for k in FFN_HIDDEN_SIZES
    for e in NUM_EXPERTS
    for topk in TOP_KS
    for num_shared_experts in NUM_SHARED_EXPERTS
]

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

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

        ylabel='Time(ms)',
        xlabel='Batch Size',
        plot_name='MoE INT8 Channel Performance (Times ms)',  # 图表标题
        args={
            'dtype': torch.half,
            'device': 'cuda',
        }
    )
]


case_counter = 1
perf_profling = False

@triton.testing.perf_report(moe_configs)
def bench_fused_moe_int8_channel(BATCH_SIZE, HIDDEN_SIZE, FFN_HIDDEN_SIZE, NUM_EXPERTS, TOP_K, NUM_SHARED_EXPERTS,
                   provider, dtype=torch.half, device="cuda"):
    """
    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={}, 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

    M, N, K, E = m, n, k, e
    a, w1, w2, w1_s, w2_s, _score_all, e_map, topk_weights, topk_ids = input_helper(
        m=M,
        n=N,
        k=K,
        e=E,
        topk=TOP_K,
        ep_size=ep_size,
        dtype=dtype,
        device=device,
        num_shared_experts=NUM_SHARED_EXPERTS,
        seed=0,
    )

    ms = None
    if provider == "triton":
        with patched_environment():
            fn = lambda: fused_moe_module.fused_experts_impl(
                a,
                w1,
                w2,
                topk_weights,
                topk_ids,
                a.dtype,
                use_int8_w8a8=True,
                per_channel_quant=True,
                global_num_experts=E + NUM_SHARED_EXPERTS,
                expert_map=e_map,
                w1_scale=w1_s,
                w2_scale=w2_s,
            )
            ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
    else:
        raise ValueError(f"Unsupported provider: {provider}")

    return ms

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

if __name__ == "__main__":
    pytest.main([__file__, "-s"])