# Test for get_aiter_moe_config and aiter_moe with w8a8 channel-wise quantization

import torch
import pandas as pd

from aiter.fused_moe import fused_topk, torch_moe
from aiter import dtypes, ActivationType
from aiter.test_common import checkAllclose, perftest
from aiter.moe import (
    get_aiter_moe_config,
    aiter_moe,
    MoeSolutionType,
    MoeQuantType,
)
from aiter.ops.shuffle import moe_layout_shuffle_gemm1, moe_layout_shuffle_gemm2
import aiter


torch.set_default_device("cuda")


@perftest(num_warmup=1, num_iters=2)
def _run_torch_ref(hidden_states, w1, w2, topk_weights, topk_ids):
    return torch_moe(
        hidden_states,
        w1,
        w2,
        topk_weights,
        topk_ids,
    )


@perftest(num_warmup=10, num_iters=100, num_rotate_args=1)
def _run_aiter_moe_perf(
    hidden_states,
    w1,
    w2,
    topk_weights,
    topk_ids,
    moe_config,
    inplace,
    activation,
    w1_scale,
    w2_scale,
    w1_zp,
    w2_zp,
    a1_scale,
    a2_scale,
    block_shape,
    global_num_experts,
    expert_map,
):

    mortal_input = hidden_states.clone() # 保证inplace操作的正确性

    return aiter_moe(
        mortal_input,
        w1,
        w2,
        topk_weights,
        topk_ids,
        moe_config,
        inplace,
        activation,
        w1_scale,
        w2_scale,
        w1_zp,
        w2_zp,
        a1_scale,
        a2_scale,
        block_shape,
        global_num_experts,
        expert_map,
    )


def prepare_w8a8_channelwise_inputs(m, k, n, e, topk, dtype):
    """Prepare channel-wise quantized w8a8 inputs.

    Scale shape: (e, out_dim, 1) — one scale per output channel.
    block_shape is None for channel-wise.
    """
    torch.manual_seed(0)

    input_tensor = torch.randn((m, k), dtype=dtype, device="cuda") / 10
    w1_fp = torch.randn((e, 2 * n, k), dtype=dtype, device="cuda")
    w2_fp = torch.randn((e, k, n), dtype=dtype, device="cuda")

    # Channel-wise quantization: max per output channel (last dim of weight row)
    max_vals_w1 = torch.abs(w1_fp.to(torch.float32)).max(dim=-1, keepdim=True)[0]
    max_vals_w1 = max_vals_w1.clamp(min=1e-5)
    w1_scales = max_vals_w1 / 127.0  # (e, 2*n, 1)
    w1_qweight = (w1_fp / max_vals_w1 * 127.0).round().clamp(min=-128, max=127).to(torch.int8)

    max_vals_w2 = torch.abs(w2_fp.to(torch.float32)).max(dim=-1, keepdim=True)[0]
    max_vals_w2 = max_vals_w2.clamp(min=1e-5)
    w2_scales = max_vals_w2 / 127.0  # (e, k, 1)
    w2_qweight = (w2_fp / max_vals_w2 * 127.0).round().clamp(min=-128, max=127).to(torch.int8)

    score = torch.randn((m, e), dtype=dtype, device="cuda")
    topk_weights, topk_ids = fused_topk(input_tensor, score, topk, True)

    # moe_c backend needs layout-shuffled weights
    w1_qweight_shuffle = moe_layout_shuffle_gemm1(w1_qweight).view(*w1_qweight.shape)
    w2_qweight_shuffle = moe_layout_shuffle_gemm2(w2_qweight).view(*w2_qweight.shape)

    return {
        "input": input_tensor,
        "w1_ref": w1_fp,
        "w2_ref": w2_fp,
        "w1_qweight": w1_qweight,
        "w2_qweight": w2_qweight,
        "w1_qweight_shuffle": w1_qweight_shuffle,
        "w2_qweight_shuffle": w2_qweight_shuffle,
        "w1_scales": w1_scales,
        "w2_scales": w2_scales,
        "topk_weights": topk_weights,
        "topk_ids": topk_ids,
    }


def test_get_config(m, k, n, e, topk, dtype):
    """Test get_aiter_moe_config for channel-wise w8a8 (block_size=0)."""
    status, moe_cfg = get_aiter_moe_config(
        M=m,
        E=e,
        N1=2 * n,
        N2=k,
        K=k,
        top_k=topk,
        block_size=0,
        dtype=dtype,
        quant_type=MoeQuantType.W8A8,
    )

    if status:
        assert moe_cfg.quant_type == MoeQuantType.W8A8
        assert moe_cfg.solution_type in (
            MoeSolutionType.ASM,
            MoeSolutionType.MOE_C,
            MoeSolutionType.TRITON,
            MoeSolutionType.CK,
        )
        assert moe_cfg.config is not None
        aiter.logger.info(
            f"[get_config_w8a8_cw] {m=}, solution={moe_cfg.solution_type}, "
            f"config keys={list(moe_cfg.config.keys())}"
        )
    else:
        assert moe_cfg.solution_type is None
        assert moe_cfg.config is None
        aiter.logger.info(f"[get_config_w8a8_cw] {m=}, no solution found")

    return status, moe_cfg


def test_aiter_moe_w8a8_channelwise(m, k, n, e, topk, dtype):
    """End-to-end test of aiter_moe with channel-wise w8a8."""
    status, moe_cfg = get_aiter_moe_config(
        M=m,
        E=e,
        N1=2 * n,
        N2=k,
        K=k,
        top_k=topk,
        block_size=0,
        dtype=dtype,
        quant_type=MoeQuantType.W8A8,
    )

    if not status:
        aiter.logger.info(f"[aiter_moe_w8a8_cw] SKIP {m=}: no backend available")
        return None

    data = prepare_w8a8_channelwise_inputs(m, k, n, e, topk, dtype)

    # Torch reference uses original fp weights directly (no scales needed)
    ref_out, _ = _run_torch_ref(
        data["input"],
        data["w1_ref"],
        data["w2_ref"],
        data["topk_weights"],
        data["topk_ids"],
    )

    # Select shuffled or plain weights depending on backend
    if moe_cfg.solution_type == MoeSolutionType.MOE_C:
        w1_input = data["w1_qweight_shuffle"]
        w2_input = data["w2_qweight_shuffle"]
    else:
        w1_input = data["w1_qweight"]
        w2_input = data["w2_qweight"]

    aiter_out, aiter_us = _run_aiter_moe_perf(
        hidden_states=data["input"],
        w1=w1_input,
        w2=w2_input,
        topk_weights=data["topk_weights"],
        topk_ids=data["topk_ids"],
        moe_config=moe_cfg,
        inplace=True,
        activation="silu",
        w1_scale=data["w1_scales"],
        w2_scale=data["w2_scales"],
        w1_zp=None,
        w2_zp=None,
        a1_scale=None,
        a2_scale=None,
        block_shape=None,
        global_num_experts=e,
        expert_map=None,
    )


    print("aiter_out",aiter_out)
    print("ref_out",ref_out)


    msg = f"[aiter_moe_w8a8_cw] {m=}, backend={moe_cfg.solution_type}"
    checkAllclose(ref_out, aiter_out, rtol=0.01, atol=100, msg=msg)
    return {"m": m, "backend": moe_cfg.solution_type, "us": aiter_us}


if __name__ == "__main__":
    dtype = dtypes.bf16
    e = 256
    topk = 8
    k = 6144
    n = 256

    aiter.logger.info("=" * 60)
    aiter.logger.info("Part 1: Testing get_aiter_moe_config for w8a8 channel-wise")
    aiter.logger.info("=" * 60)
    test_tokens = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048 , 4096, 6144 , 8192 , 16384]
    for m in test_tokens:
        test_get_config(m, k, n, e, topk, dtype)

    aiter.logger.info("=" * 60)
    aiter.logger.info("Part 2: Testing aiter_moe end-to-end for w8a8 channel-wise")
    aiter.logger.info("=" * 60)
    df = []
    for m in test_tokens:
        ret = test_aiter_moe_w8a8_channelwise(m, k, n, e, topk, dtype)
        if ret is not None:
            df.append(ret)
    if df:
        df = pd.DataFrame(df)
        aiter.logger.info(f"summary:\n{df}")
