# SPDX-License-Identifier: MIT
import torch
import itertools
import aiter
from aiter import dtypes
from aiter.test_common import (
    checkAllclose,
    benchmark,
    run_perftest,
)
from aiter.int4_utils import *

from aiter.fused_moe import (
    fused_topk,
    moe_sorting,
    fused_moe,
    torch_moe_stage1,
    torch_moe_stage2,
    get_block_size_M,
)
from aiter.fused_moe_ck import (
    ck_moe_stage_1,
    ck_moe_stage_2,
)


from aiter.ops.shuffle import shuffle_weight
from aiter import ActivationType

torch.int4 = getattr(torch, "int4", torch.uint32)
torch.set_default_device("cuda")
torch.set_printoptions(precision=4, sci_mode=False)

BLOCK_SIZE_M = 32


def print_full_high_dim(label, tensor, max_full_dim=12):
    if tensor.ndim == 0:
        print(f"{label}: {tensor}")
        return
    dim0 = tensor.shape[0]
    if dim0 > max_full_dim:
        print(f"{label} has dim0={dim0}, showing first {max_full_dim} entries only")
        indices = range(max_full_dim)
    else:
        indices = range(dim0)
    for idx in indices:
        print(f"{label}[{idx}]: {tensor[idx]}")
    if dim0 > max_full_dim:
        print(f"... {label} truncated; dim0={dim0} ...")


def print_full_leading_dims(label, tensor, max_entries=None):
    if tensor.ndim == 0:
        print(f"{label}: {tensor}")
        return
    leading_shape = tensor.shape[:-1]
    if not leading_shape:
        print(f"{label}: {tensor}")
        return
    ranges = []
    truncated = False
    for axis, size in enumerate(leading_shape):
        if max_entries is not None and size > max_entries:
            ranges.append(range(max_entries))
            truncated = True
            print(f"{label} dim{axis} size={size}, showing first {max_entries} entries")
        else:
            ranges.append(range(size))
    for idx in itertools.product(*ranges):
        idx_str = ",".join(str(i) for i in idx)
        print(f"{label}[{idx_str}]: {tensor[idx]}")
    if truncated:
        print(f"... {label} truncated across leading dims ...")


@benchmark()
def test_fmoe(
    dtype,
    token,
    model_dim,
    inter_dim,
    E,
    topk,
    actType,
    qType,
    AQDType,
    WQDType,
    use_g1u1=False,
    doweight_stage1=False,
):
    torch_quant = aiter.get_torch_quant(qType)
    torch_act = aiter.get_torch_act(actType)
    input = torch.randn((token, model_dim), dtype=dtype) / 10
    # input = torch.full((token, model_dim), 0.01, dtype=dtype)
    # token_scales = (
    #     torch.arange(1, token + 1, device="cuda", dtype=torch.float32).view(token, 1) * 0.01
    # )
    # input = token_scales.to(dtype=dtype).repeat(1, model_dim)

    if use_g1u1:
        w1 = torch.randn((E, inter_dim * 2, model_dim), dtype=dtype) / 10
        # w1 = torch.full((E, inter_dim * 2, model_dim), 0.01, dtype=dtype)
        # expert_scales = (
        #     torch.arange(1, E + 1, device=input.device, dtype=torch.float32)
        #     .view(E, 1, 1)
        #     * 0.01
        # )
        # w1 = expert_scales.to(dtype=dtype).repeat(1, inter_dim * 2, model_dim)
    else:
        w1 = torch.randn((E, inter_dim, model_dim), dtype=dtype)
    w2 = torch.randn((E, model_dim, inter_dim), dtype=dtype) / 10
    # w2 = torch.full((E, model_dim, inter_dim), 0.02, dtype=dtype)

    score = torch.randn((token, E), dtype=dtype)
    topk_weights, topk_ids = fused_topk(input, score, topk, True)

    M, _ = topk_ids.shape

    sorted_ids, sorted_weights, sorted_expert_ids, num_valid_ids, tokens_positions_per_expert, moe_buf = moe_sorting(
        topk_ids, topk_weights, E, model_dim, torch.float32, BLOCK_SIZE_M
    )

    # print(f"########### token_per_expert: {tokens_positions_per_expert}")


    if qType == aiter.QuantType.per_Tensor:
        w1_qt, w1_scale = aiter.pertoken_quant(w1.view(E, -1), quant_dtype=WQDType)
        w2_qt, w2_scale = aiter.pertoken_quant(w2.view(E, -1), quant_dtype=WQDType)
    elif qType == aiter.QuantType.per_Token and WQDType == torch.int4:  # int4 w quant
        w1_qt, w1_scale = aiter.pertoken_quant(w1, quant_dtype=dtypes.i8, dtypeMax=7)
        w2_qt, w2_scale = aiter.pertoken_quant(w2, quant_dtype=dtypes.i8, dtypeMax=7)
    else:
        w1_qt, w1_scale = torch_quant(w1, quant_dtype=WQDType)
        w2_qt, w2_scale = torch_quant(w2, quant_dtype=WQDType)

    w1_qt = w1_qt_aiter = w1_qt.view(w1.shape)
    w2_qt = w2_qt_aiter = w2_qt.view(w2.shape)

    a1_qt, a1_scale = torch_quant(input, quant_dtype=AQDType)

    out1_ref = torch_moe_stage1(
        a1_qt,
        w1_qt,
        w2_qt,
        topk_weights,
        topk_ids,
        dtype=dtype,
        activation=actType,
        quant_type=qType,
        a1_scale=a1_scale,
        w1_scale=w1_scale,
        doweight=doweight_stage1,
        group_by_expert=True,      #是否以expert为单位对token排序
    )

    if WQDType == torch.int4:  # int4 w quant
        w1_qt_aiter = rearrange_4bit_elements(
            convert_int8_to_uint32_int4(
                shuffle_weight(w1_qt_aiter, (16, 16), use_int4=True)
            )
        )
        w2_qt_aiter = rearrange_4bit_elements(
            convert_int8_to_uint32_int4(
                shuffle_weight(w2_qt_aiter, (16, 16), use_int4=True)
            )
        )
    else:
        pass
        # w1_qt_aiter = shuffle_weight(w1_qt_aiter, layout=(16, 16))
        # w2_qt_aiter = shuffle_weight(w2_qt_aiter, layout=(16, 16))

    # ######################## stage 1 start ###########
    out1_ck, us = run_perftest(
        ck_moe_stage_1,
        a1_qt,
        w1_qt_aiter,
        w2_qt_aiter,
        sorted_ids,
        sorted_expert_ids,
        tokens_positions_per_expert,
        num_valid_ids,
        False,
        False,
        w1_scale,
        a1_scale,
        dtype,
        topk,
        BLOCK_SIZE_M,
        actType,
        sorted_weights=sorted_weights if doweight_stage1 else None,
        num_iters=100,
        num_warmup=20,
    )

    # group_by_expert=True时，torch_moe_stage1的输出符合ck_moe_stage1的输出，可进行以下对比
    checkAllclose(
        out1_ref.view(-1),
        out1_ck.view(-1),
        rtol=1e-2,
        atol=1e-2,
        msg=f"[perf]  ck_moe_stage_1:{us:>8.2f} us, {token*model_dim*inter_dim*2*topk*2/us/1000/1000:>8.2f} tflops......(quant:{AQDType})",
    )

    # ######################## stage 1 end ###########



    # get the input for stage 2, group_by_expert=False; but the above out1_ref  group_by_expert=True
    out1_torch = torch_moe_stage1(
        a1_qt,
        w1_qt,
        w2_qt,
        topk_weights,
        topk_ids,
        dtype=dtype,
        activation=actType,
        quant_type=qType,
        a1_scale=a1_scale,
        w1_scale=w1_scale,
        doweight=doweight_stage1,
        group_by_expert=False,      #是否以expert为单位对token排序
    )

    # ######################## stage 2 start ###########
    if qType == aiter.QuantType.per_Token:
        out1_torch = out1_torch.view(token, -1)
    a2_qt, a2_scale = torch_quant(out1_torch, quant_dtype=AQDType)


    out2_ref = torch_moe_stage2(
        a2_qt,
        w1_qt,  # E, inter_dim*2, model_dim
        w2_qt,  # E, model_dim, inter_dim
        topk_weights,
        topk_ids,
        dtype=dtype,
        quant_type=qType,
        w2_scale=w2_scale,
        a2_scale=a2_scale,
        doweight=not doweight_stage1,
    )

    out2_ck, us = run_perftest(
        ck_moe_stage_2,
        out1_ck,
        w1_qt_aiter,
        w2_qt_aiter,
        sorted_ids,
        sorted_expert_ids,
        tokens_positions_per_expert,
        num_valid_ids,
        False,
        False,
        w2_scale,
        a2_scale,
        dtype,
        topk,
        BLOCK_SIZE_M,
        sorted_weights=sorted_weights if not doweight_stage1 else None,
        moe_buf=None,
        num_iters=100,
        num_warmup=20,
    )

    checkAllclose(
        out2_ref.view(-1),
        out2_ck.view(-1),
        rtol=1e-2,
        atol=1e-2,
        msg=f"[perf]  ck_moe_stage_2:{us:>8.2f} us, {token*model_dim*inter_dim*topk*2/us/1000/1000:>8.2f} tflops......(quant:{AQDType})",
    )


list_dtype = [dtypes.bf16, dtypes.fp16][1:2]
list_dim = [(2048, 1024), (7168, 256)][1:]
list_tokenNum = [
    1, 16, 32, 64, 256, 512, 1024, 2048
]
list_quant = [
    (aiter.QuantType.No, None, None),  # a16w16
    # (aiter.QuantType.per_Tensor, dtypes.fp8, dtypes.fp8),  # a8w8
    # (aiter.QuantType.per_Token, dtypes.fp8, dtypes.fp8),  # a8w8
    # (aiter.QuantType.per_Token, dtypes.fp8, torch.int4),  # a8w4
    # (aiter.QuantType.per_128x128, dtypes.fp8, dtypes.fp8),  # a8w8 TODO add test
]
list_act = [aiter.ActivationType.Silu, aiter.ActivationType.Gelu][0:1]
list_doweight_stage1 = [False][:]       # Do the results multiply with topk_weight
expert, topk = 256, 8

import pandas as pd

for (
    dtype,
    act_type,
    (quant_type, aq_dtype, wq_dtype),
    (model_dim, inter_dim),
    doweight_stage1,
) in itertools.product(
    list_dtype, list_act, list_quant, list_dim, list_doweight_stage1
):
    df = []
    for m in list_tokenNum:
        ret = test_fmoe(
            dtype,
            m,
            model_dim,      # hidden_size
            inter_dim,      # intermediate_size
            expert,
            topk,
            act_type,
            quant_type,
            aq_dtype,
            wq_dtype,
            use_g1u1=True,
            doweight_stage1=doweight_stage1,
        )
        df.append(ret)
    df = pd.DataFrame(df)
    aiter.logger.info(f"summary:\n{df}")
    print(f"###############  Done test item. #################\n")
