# SPDX-License-Identifier: Apache-2.0
"""Tests for the channel-wise W4A8 MoE path."""

import gc
import os
from contextlib import contextmanager
from typing import Optional

import pytest
import torch
import torch.nn.functional as F
import triton

import aiter.ops.triton.fused_moe as fused_moe_module
from aiter import per_token_quant_hip
from op_tests.test_aiter_moe_with_config_w4a8_perchannel import _run_triton_ref

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

    @contextmanager
    def patched_environment():
        yield


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
    batchsToCompile = [1, 16]
    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:
    if fast_mode == 0:
        group_ids = {
            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 = {
            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,
    "cur_group": cur_group,
    "default_device": 0,
    "group_id": group_ids,
}

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

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

if training_config["training_mode"] == 0:
    BATCH_SIZES = [1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384]
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 allocate_quantized_weights(
    e: int,
    n: int,
    k: int,
    device: str,
):
    w1_qweight = torch.empty((e, 2 * n, k // 2), device=device, dtype=torch.uint8)
    w2_qweight = torch.empty((e, k, n // 2), device=device, dtype=torch.uint8)
    w1_scales = torch.empty((e, 2 * n, 1), device=device, dtype=torch.float32)
    w2_scales = torch.empty((e, k, 1), device=device, dtype=torch.float32)
    return w1_qweight, w2_qweight, w1_scales, w2_scales


def generate_channelwise_w4a8_weight(
    n: int,
    k: int,
    device: str,
):
    weight_u4 = torch.randint(0, 15, (n, k), device=device, dtype=torch.int8)
    weight_ref = torch.where(weight_u4 > 7, weight_u4 - 16, weight_u4).contiguous()
    # Match the Triton kernel's unpack convention: even-k nibble in low 4 bits,
    # odd-k nibble in high 4 bits.
    qweight = (
        weight_u4[:, 1::2].to(torch.uint8) * 16
        + weight_u4[:, ::2].to(torch.uint8)
    ).contiguous()
    scales = torch.randn((n, 1), device=device, dtype=torch.float32)
    return weight_ref, qweight, scales


def quantize_dense_weights(
    e: int,
    n: int,
    k: int,
    device: str,
    return_reference: bool = True,
):
    w1_qweight, w2_qweight, w1_scales, w2_scales = allocate_quantized_weights(
        e, n, k, device
    )

    w1_ref = torch.empty((e, 2 * n, k), device=device, dtype=torch.int8) if return_reference else None
    w2_ref = torch.empty((e, k, n), device=device, dtype=torch.int8) if return_reference else None

    for expert_id in range(e):
        ref, qweight, scales = generate_channelwise_w4a8_weight(2 * n, k, device)
        if return_reference:
            w1_ref[expert_id] = ref
        w1_qweight[expert_id] = qweight
        w1_scales[expert_id] = scales

        ref, qweight, scales = generate_channelwise_w4a8_weight(k, n, device)
        if return_reference:
            w2_ref[expert_id] = ref
        w2_qweight[expert_id] = qweight
        w2_scales[expert_id] = scales

    return w1_ref, w2_ref, w1_qweight, w2_qweight, w1_scales, w2_scales


def quantize_streaming_weights(
    e: int,
    n: int,
    k: int,
    device: str,
):
    w1_qweight, w2_qweight, w1_scales, w2_scales = allocate_quantized_weights(
        e, n, k, device
    )

    for expert_id in range(e):
        _ref, qweight, scales = generate_channelwise_w4a8_weight(2 * n, k, device)
        w1_qweight[expert_id] = qweight
        w1_scales[expert_id] = scales
        del qweight, scales, _ref

        _ref, qweight, scales = generate_channelwise_w4a8_weight(k, n, device)
        w2_qweight[expert_id] = qweight
        w2_scales[expert_id] = scales
        del qweight, scales, _ref

        torch.cuda.empty_cache()

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


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

    local_e = e // ep_size
    local_ids = torch.arange(local_e, device=device, dtype=torch.int32)
    e_map = torch.full((e,), -1, device=device, dtype=torch.int32)
    e_map[local_ids] = torch.arange(local_e, device=device, dtype=torch.int32)

    local_tensors = []
    for tensor in expert_tensors:
        local_tensors.append(None if tensor is None else tensor[local_ids])
    return e_map, tuple(local_tensors)


def generate_unique_int_tensor(
    m: int,
    topk: int,
    high: int,
    device: str,
):
    if topk > high:
        raise ValueError(f"topk ({topk}) cannot exceed num_experts ({high})")

    topk_ids = torch.empty((m, topk), device=device, dtype=torch.int32)
    for row in range(m):
        topk_ids[row] = torch.randperm(high, device=device, dtype=torch.int32)[:topk]
    return topk_ids


def input_helper(
    m: int,
    n: int,
    k: int,
    e: int,
    topk: int,
    ep_size: int,
    dtype: torch.dtype,
    device: str = "cuda",
    seed: int = 0,
    streaming_weights: bool = False,
    renormalize: bool = False,
    return_reference: bool = True,
):
    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.rand((m, k), device=device, dtype=dtype) / 10000
    topk_logits = torch.rand((m, topk), device=device, dtype=torch.float32)
    topk_weights = F.softmax(topk_logits, dim=1)
    topk_ids = generate_unique_int_tensor(m, topk, e, device)
    score_all = None

    if streaming_weights:
        quantized_weights = quantize_streaming_weights(e, n, k, device)
    else:
        quantized_weights = quantize_dense_weights(
            e,
            n,
            k,
            device,
            return_reference=return_reference,
        )

    e_map, quantized_weights = build_expert_map(e, ep_size, device, *quantized_weights)
    w1_ref, w2_ref, w1_qweight, w2_qweight, w1_scales, w2_scales = quantized_weights
    return (
        a,
        w1_ref,
        w2_ref,
        w1_qweight,
        w2_qweight,
        w1_scales,
        w2_scales,
        score_all,
        e_map,
        topk_weights,
        topk_ids,
    )


def prepare_case(
    m: int,
    n: int,
    k: int,
    e: int,
    topk: int,
    ep_size: int,
    dtype: torch.dtype,
    device: str,
):
    a, w1_ref, w2_ref, w1_qweight, w2_qweight, w1_scale, w2_scale, _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,
        seed=0,
        renormalize=False,
        return_reference=True,
    )
    return (
        a,
        w1_ref,
        w2_ref,
        w1_qweight,
        w2_qweight,
        w1_scale,
        w2_scale,
        topk_weights,
        topk_ids,
        e_map,
    )


def torch_w4a8_channelwise_reference_fallback(
    a: torch.Tensor,
    w1_ref: torch.Tensor,
    w2_ref: torch.Tensor,
    w1_scale: torch.Tensor,
    w2_scale: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    expert_map: Optional[torch.Tensor],
):
    num_tokens, hidden = a.shape
    topk = topk_weights.shape[1]
    a_q, a_scale = per_token_quant_hip(a)

    a_q = a_q.view(num_tokens, 1, hidden).repeat(1, topk, 1).reshape(-1, hidden)
    a_scale = a_scale.view(num_tokens, 1, 1).repeat(1, topk, 1).reshape(-1, 1)

    routed_ids = topk_ids.reshape(-1)
    routed_weights = topk_weights.reshape(-1)
    if expert_map is not None:
        routed_ids = expert_map[routed_ids]

    out = torch.zeros((num_tokens * topk, w2_ref.shape[1]), device=a.device, dtype=torch.float32)

    for expert_id in range(w1_ref.shape[0]):
        mask = routed_ids == expert_id
        if not mask.any():
            continue

        inter = torch.matmul(
            a_q[mask].to(torch.float32),
            w1_ref[expert_id].to(torch.float32).T,
        ) * a_scale[mask].to(torch.float32) * w1_scale[expert_id].view(1, -1)
        half = inter.shape[1] // 2
        act = (F.silu(inter[:, :half]) * inter[:, half:]).to(a.dtype)

        act_q, act_scale = per_token_quant_hip(act)
        out[mask] = torch.matmul(
            act_q.to(torch.float32),
            w2_ref[expert_id].to(torch.float32).T,
        ) * act_scale.to(torch.float32) * w2_scale[expert_id].view(1, -1)

    out = out.view(num_tokens, topk, -1)
    out = out * routed_weights.view(num_tokens, topk, 1).to(out.dtype)
    return out.sum(dim=1).to(a.dtype)


def torch_w4a8_channelwise_reference(
    a: torch.Tensor,
    w1_ref: torch.Tensor,
    w2_ref: torch.Tensor,
    w1_scale: torch.Tensor,
    w2_scale: torch.Tensor,
    topk_weights: torch.Tensor,
    topk_ids: torch.Tensor,
    expert_map: Optional[torch.Tensor],
):
    if expert_map is not None:
        return torch_w4a8_channelwise_reference_fallback(
            a,
            w1_ref,
            w2_ref,
            w1_scale,
            w2_scale,
            topk_weights,
            topk_ids,
            expert_map,
        )

    return _run_triton_ref(
        a,
        w1_ref,
        w2_ref,
        topk_weights,
        topk_ids,
        a.dtype,
        None,
        w1_scale,
        w2_scale,
        w1_ref.shape[0],
    )


@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)
@torch.inference_mode()
def test_int4_w4a8_channelwise_fused_moe(
    m: int,
    n: int,
    k: int,
    e: int,
    topk: int,
    ep_size: int,
    dtype: torch.dtype,
):
    setup_device_from_training_config(skip_if_unavailable=True)
    device = "cuda"

    ENABLE_TORCH_REF = training_config["training_mode"] == 0 and m <= 256

    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"ENABLE_TORCH_REF={ENABLE_TORCH_REF}!!!\n"
    )
    case_counter_x += 1

    a, w1_ref, w2_ref, w1_qweight, w2_qweight, w1_scale, w2_scale, _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,
        seed=0,
        renormalize=False,
        return_reference=ENABLE_TORCH_REF,
    )

    with patched_environment():
        out = fused_moe_module.fused_experts_impl(
            a,
            w1_qweight,
            w2_qweight,
            topk_weights,
            topk_ids,
            dtype,
            use_int4_w4a8=True,
            per_channel_quant=True,
            global_num_experts=e,
            expert_map=e_map,
            w1_scale=w1_scale,
            w2_scale=w2_scale,
        )

    if ENABLE_TORCH_REF:
        ref = torch_w4a8_channelwise_reference(
            a,
            w1_ref,
            w2_ref,
            w1_scale,
            w2_scale,
            topk_weights,
            topk_ids,
            e_map,
        )
        rel_diff = (
            torch.mean(torch.abs(out.to(torch.float32) - ref.to(torch.float32)))
            / torch.clamp(torch.mean(torch.abs(ref.to(torch.float32))), min=1e-6)
        )
        assert rel_diff < 0.05, (
            f"{m=}, {n=}, {k=}, {e=}, {topk=}, {ep_size=}, {dtype=}, {rel_diff=}"
        )


# 定义MoE测试用例列表
moe_perf_model_cases_list = [
    (m, n, k, e, topk, ep_size)
    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 ep_size in EP_SIZE
]

# 配置MoE基准测试参数
moe_configs = [
    triton.testing.Benchmark(
        x_names=["BATCH_SIZE", "HIDDEN_SIZE", "FFN_HIDDEN_SIZE", "NUM_EXPERTS", "TOP_K", "EP_SIZE"],
        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 INT4INT8 Channelwise Performance (Times ms)",
        args={
            "dtype": torch.float16,
            "device": "cuda",
        },
    )
]


case_counter = 1
perf_profling = False


@triton.testing.perf_report(moe_configs)
def bench_fused_moe_int4int8_channelwise(
    BATCH_SIZE,
    HIDDEN_SIZE,
    FFN_HIDDEN_SIZE,
    NUM_EXPERTS,
    TOP_K,
    EP_SIZE,
    provider,
    dtype=torch.float16,
    device="cuda",
):
    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

    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")

    a, _w1_ref, _w2_ref, w1_qweight, w2_qweight, w1_scale, w2_scale, _score_all, e_map, topk_weights, topk_ids = input_helper(
        m=BATCH_SIZE,
        n=HIDDEN_SIZE,
        k=FFN_HIDDEN_SIZE,
        e=NUM_EXPERTS,
        topk=TOP_K,
        ep_size=EP_SIZE,
        dtype=dtype,
        device=device,
        seed=0,
        streaming_weights=perf_profling,
        renormalize=True,
        return_reference=False,
    )

    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,
                w1_qweight,
                w2_qweight,
                topk_weights,
                topk_ids,
                dtype,
                use_int4_w4a8=True,
                per_channel_quant=True,
                global_num_experts=NUM_EXPERTS,
                expert_map=e_map,
                w1_scale=w1_scale,
                w2_scale=w2_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_channelwise.run(print_data=True)
