# Test for get_aiter_moe_config and aiter_moe with W16A16 non-gated ReLU²
# (Nemotron-style MOE: N1 = intermediate_size, activation = relu2)

import torch
import pandas as pd
from typing import Optional, List

from aiter.fused_moe import fused_topk
from aiter import dtypes
from aiter.test_common import checkAllclose, perftest
from aiter.moe import (
    get_aiter_moe_config,
    aiter_moe,
    MoeSolutionType,
    MoeQuantType,
)
from aiter.fused_moe_asm_wna16 import fused_experts_asm_impl
from aiter.ops.shuffle import asm_shuffle_weight_b8
import aiter

torch.set_default_device("cuda")


# ---------------------------------------------------------------------------
# Torch reference for non-gated ReLU² MOE
# ---------------------------------------------------------------------------
def torch_moe_relu2(hidden_states, w1, w2, topk_weights, topk_ids):
    """Reference implementation for non-gated ReLU² MOE.

    w1: [E, inter_dim, model_dim]   (NOT 2*inter_dim)
    w2: [E, model_dim, inter_dim]
    """
    computeType = torch.float32
    dtype = hidden_states.dtype
    hidden_states = hidden_states.to(computeType)
    w1 = w1.to(computeType)
    w2 = w2.to(computeType)

    B, D = hidden_states.shape
    topk = topk_weights.shape[1]

    hidden_states = hidden_states.view(B, -1, D).repeat(1, topk, 1)
    out = torch.zeros(
        (B, topk, D),
        dtype=computeType,
        device=hidden_states.device,
    )

    for E_id in range(w1.shape[0]):
        mask = topk_ids == E_id
        if mask.sum():
            sub_tokens = hidden_states[mask]
            # GEMM1
            h = sub_tokens @ w1[E_id].T
            # ReLU²
            h = torch.relu(h) ** 2
            # GEMM2
            out[mask] = h @ w2[E_id].T

    return (out * topk_weights.view(B, -1, 1)).sum(dim=1).to(dtype)


# ---------------------------------------------------------------------------
# Weight preparation helpers (W16A16 non-gated)
# ---------------------------------------------------------------------------
def prepare_w16a16_nogate_inputs(m, k, n, e, topk, dtype):
    """Build all tensors needed to run a non-gated w16a16 MOE test.

    Key difference from gated: w1 shape is [E, n, k] instead of [E, 2*n, k].
    """
    torch.manual_seed(0)
    input_tensor = torch.randn((m, k), device="cuda", dtype=dtype) / 10
    w1 = torch.randn((e, n, k), device="cuda", dtype=dtype) / 2
    w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 2
    score = torch.randn((m, e), device="cuda", dtype=dtype)

    w1_shuffle = asm_shuffle_weight_b8(w1, stage=1)
    w2_shuffle = asm_shuffle_weight_b8(w2, stage=2)

    topk_weights, topk_ids = fused_topk(input_tensor, score, topk, True)

    return {
        "input": input_tensor,
        "w1": w1,
        "w2": w2,
        "w1_shuffle": w1_shuffle,
        "w2_shuffle": w2_shuffle,
        "topk_weights": topk_weights,
        "topk_ids": topk_ids,
        "score": score,
    }


# ---------------------------------------------------------------------------
# Test: get_aiter_moe_config (w16a16 non-gated relu2)
# ---------------------------------------------------------------------------
def test_get_config(m, k, n, e, topk, dtype):
    """Test that get_aiter_moe_config returns a valid w16a16 config with
    activation='relu2' or gracefully reports no-solution."""
    N1 = n         # non-gated: N1 = intermediate_size (NOT 2 * intermediate_size)
    N2 = k         # down / hidden_size
    K = k          # model dimension

    status, moe_cfg = get_aiter_moe_config(
        M=m, E=e, N1=N1, N2=N2, K=K,
        top_k=topk, block_size=0, dtype=dtype,
        quant_type=MoeQuantType.W16A16,
        activation="relu2",
        gated=False,
    )

    if status:
        assert moe_cfg.solution_type is not None, \
            "status=True but solution_type is None"
        assert moe_cfg.config is not None, \
            "status=True but config is None"
        assert moe_cfg.solution_type in (
            MoeSolutionType.ASM,
            MoeSolutionType.TRITON,
        ), f"Unexpected solution_type: {moe_cfg.solution_type}"
        assert moe_cfg.quant_type == MoeQuantType.W16A16
        aiter.logger.info(
            f"[get_config_w16a16_nogate] {m=}, {k=}, {n=}, {e=}, {topk=}, "
            f"solution={moe_cfg.solution_type}, "
            f"config keys={list(moe_cfg.config.keys())}"
        )
    else:
        assert moe_cfg.solution_type is None, \
            "status=False but solution_type is not None"
        assert moe_cfg.config is None, \
            "status=False but config is not None"
        aiter.logger.info(
            f"[get_config_w16a16_nogate] {m=}, {k=}, {n=}, {e=}, {topk=}, "
            f"no solution found (expected on unsupported configs)"
        )

    return status, moe_cfg


# ---------------------------------------------------------------------------
# Test: aiter_moe end-to-end for w16a16 non-gated relu2
# ---------------------------------------------------------------------------
@perftest(num_warmup=1, num_iters=2)
def _run_torch_ref(hidden_states, w1, w2, topk_weights, topk_ids):
    return torch_moe_relu2(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,
        routed_scaling_factor,
        ):
    return aiter_moe(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, routed_scaling_factor)


def test_aiter_moe_w16a16_nogate(m, k, n, e, topk, dtype, inplace, routed_scaling_factor):
    """End-to-end: get config -> run aiter_moe with relu2 -> compare with
    torch reference."""
    N1 = n       # non-gated: N1 = intermediate_size
    N2 = k
    K = k

    status, moe_cfg = get_aiter_moe_config(
        M=m, E=e, N1=N1, N2=N2, K=K,
        top_k=topk, block_size=0, dtype=dtype,
        quant_type=MoeQuantType.W16A16,
        activation="relu2",
        gated=False,
    )

    if not status:
        aiter.logger.info(
            f"[aiter_moe_w16a16_nogate] SKIP {m=}, {N1=}, {N2=}, {K=}, {e=}, {topk=}: "
            f"no backend available"
        )
        return None

    backend = moe_cfg.solution_type
    aiter.logger.info(
        f"[aiter_moe_w16a16_nogate] {m=}, {N1=}, {N2=}, {K=}, {e=}, {topk=}, "
        f"backend={backend}"
    )

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

    # Torch reference
    ref_out, _ = _run_torch_ref(
        data["input"], data["w1"], data["w2"],
        data["topk_weights"], data["topk_ids"],
    )

    # aiter_moe dispatch with relu2 activation
    aiter_us = 1.0
    aiter_out, aiter_us = _run_aiter_moe_perf(
        hidden_states=data["input"],
        w1=data["w1"],
        w2=data["w2"],
        topk_weights=data["topk_weights"],
        topk_ids=data["topk_ids"],
        moe_config=moe_cfg,
        inplace=inplace,
        activation="relu2",
        w1_scale=None,
        w2_scale=None,
        w1_zp=None,
        w2_zp=None,
        a1_scale=None,
        a2_scale=None,
        block_shape=None,
        global_num_experts=e,
        expert_map=None,
        routed_scaling_factor=routed_scaling_factor,
    )

    msg = (f"[aiter_moe_w16a16_nogate] {m=}, {N1=}, {N2=}, {K=}, {e=}, {topk=}, "
           f"backend={backend}")
    checkAllclose(ref_out, aiter_out, rtol=0.01, atol=0.5, msg=msg)
    return {"m": m, "backend": backend, "us": aiter_us}


# ---------------------------------------------------------------------------
# Test: aiter_moe w16a16 non-gated ASM shuffle vs non-shuffle
# ---------------------------------------------------------------------------
@perftest(num_warmup=10, num_iters=100, num_rotate_args=1)
def _run_asm_perf(hidden_states, w1, w2, topk_weights, topk_ids,
                  dtype, global_num_experts, expert_map):
    return fused_experts_asm_impl(
        hidden_states, w1, w2, topk_weights, topk_ids, dtype,
        activation="relu2",
        global_num_experts=global_num_experts,
        expert_map=expert_map)


@perftest(num_warmup=10, num_iters=100, num_rotate_args=1)
def _run_asm_shuffle_perf(hidden_states, w1, w2, topk_weights, topk_ids,
                         dtype, global_num_experts, expert_map):
    return fused_experts_asm_impl(
        hidden_states, w1, w2, topk_weights, topk_ids, dtype,
        activation="relu2",
        global_num_experts=global_num_experts,
        expert_map=expert_map,
        use_shuffle=1)


def test_aiter_moe_w16a16_nogate_shuffle(m, k, n, e, topk, dtype):
    """Test w16a16 non-gated ASM with shuffled weights vs non-shuffled ASM."""
    data = prepare_w16a16_nogate_inputs(m, k, n, e, topk, dtype)

    try:
        asm_out, asm_us = _run_asm_perf(
            data["input"], data["w1"], data["w2"],
            data["topk_weights"], data["topk_ids"],
            dtype, e, None)
    except Exception as exc:
        aiter.logger.info(
            f"[w16a16_nogate_shuffle] SKIP {m=}: ASM not available ({exc})")
        return None

    shuffle_out, shuffle_us = _run_asm_shuffle_perf(
        data["input"], data["w1_shuffle"], data["w2_shuffle"],
        data["topk_weights"], data["topk_ids"],
        dtype, e, None)

    msg = (f"[w16a16_nogate_shuffle] {m=}, {k=}, {n=}, {e=}, {topk=}, "
           f"asm_us={asm_us:.2f}, shuffle_us={shuffle_us:.2f}")
    checkAllclose(asm_out, shuffle_out, rtol=0.01, atol=0.01, msg=msg)

    uplift = asm_us / shuffle_us - 1 if shuffle_us > 0 else 0
    return {
        "m": m,
        "asm_us": asm_us,
        "shuffle_us": shuffle_us,
        "shuffle_uplift": f"{uplift:.1%}",
    }


if __name__ == "__main__":

    dtype = dtypes.bf16

    # Nemotron-style MoE parameters (non-gated, ReLU²)
    e = 256
    topk = 8
    k = 3072       # model_dim / hidden_size
    n = 128        # intermediate_size (NOT multiplied by 2)
    inplace = False
    routed_scaling_factor = 1.0

    # --- Part 1: test get_aiter_moe_config (w16a16 non-gated relu2) ---
    aiter.logger.info("=" * 60)
    aiter.logger.info("Part 1: Testing get_aiter_moe_config for w16a16 non-gated relu2")
    aiter.logger.info("=" * 60)

    test_tokens = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096]
    for m in test_tokens:
        test_get_config(m, k, n, e, topk, dtype)

    # --- Part 2: test aiter_moe end-to-end (w16a16 non-gated relu2) ---
    aiter.logger.info("=" * 60)
    aiter.logger.info("Part 2: Testing aiter_moe end-to-end for w16a16 non-gated relu2")
    aiter.logger.info("=" * 60)

    df = []
    for m in test_tokens:
        ret = test_aiter_moe_w16a16_nogate(m, k, n, e, topk, dtype, inplace, routed_scaling_factor)
        if ret is not None:
            df.append(ret)
    if df:
        df = pd.DataFrame(df)
        aiter.logger.info(f"aiter_moe non-gated relu2 summary:\n{df}")

    # --- Part 3: test ASM shuffle vs non-shuffle (w16a16 non-gated relu2) ---
    aiter.logger.info("=" * 60)
    aiter.logger.info("Part 3: Testing ASM shuffle vs non-shuffle for w16a16 non-gated relu2")
    aiter.logger.info("=" * 60)

    df_shuffle = []
    for m in test_tokens:
        ret = test_aiter_moe_w16a16_nogate_shuffle(m, k, n, e, topk, dtype)
        if ret is not None:
            df_shuffle.append(ret)
    if df_shuffle:
        df_shuffle = pd.DataFrame(df_shuffle)
        aiter.logger.info(f"shuffle summary (non-gated relu2):\n{df_shuffle}")
