# SPDX-License-Identifier: MIT
 
import os
import torch
import argparse
import aiter
from aiter import dtypes
from aiter.fused_moe import fused_topk, torch_moe
from op_tests.utility.utils import quantize_weights
from aiter.ops.triton.fused_moe import fused_experts_impl
from op_tests.utility.scalar_type import ScalarType, scalar_types
from aiter.ops.communication import init_dist_env,destroy_dist_env
from aiter.ops.shuffle import shuffle_weight
import mori
import multiprocessing as mp
from aiter import get_hip_quant
from aiter.test_common import checkAllclose, run_perftest
import torch.distributed as dist
torch.cuda.manual_seed_all(40)
def run_ref(
    world_size,
    E,
    tokens,
    topk_weights,
    topk_ids,
    w1_ep,
    w2_ep,
    w1_scale_ep,
    w2_scale_ep,
    quant_type,
):
    ref = torch.zeros(tokens.shape, dtype=tokens.dtype, device=tokens.device)
    # out_list = []
    for i in range(world_size):
        mask = (topk_ids >= i * E // world_size) & (
            topk_ids < (i + 1) * E // world_size
        )
        if not mask.any():
            continue
        topk_ids_ep = topk_ids[mask.any(1)]
        topk_weights_ep = topk_weights[mask.any(1)]
        tokens_ep = tokens[mask.any(1)]
        expert_mask = torch.zeros((E,), dtype=dtypes.i32, device=w1_ep[i].device)
        expert_mask[E // world_size * i : E // world_size * (i + 1)] = 1
        num_local_tokens = torch.tensor(
            [tokens_ep.shape[0]], dtype=dtypes.i32, device=tokens_ep.device
        )
        quant_func = get_hip_quant(
            quant_type
            if quant_type != aiter.QuantType.per_128x128
            else aiter.QuantType.per_1x128
        )
        tokens_ep_qt, scale = quant_func(tokens_ep, quant_dtype=dtypes.fp8)
        out, us = run_perftest(
            fused_moe,
            tokens_ep_qt,
            w1_ep[i],
            w2_ep[i],
            topk_weights_ep,
            topk_ids_ep,
            expert_mask,
            num_local_tokens=num_local_tokens,
            w1_scale=w1_scale_ep[i],
            w2_scale=w2_scale_ep[i],
            quant_type=quant_type,
            a1_scale=scale,
            dtype=dtypes.bf16,
        )
        print(f"rank {i} us={us:.4f}")
        # out_list.append(out)
        # return out_list
        ref[mask.any(1)] += out
        # ref[mask.any(1)] += out.to(dtypes.fp32)
    return ref.to(tokens)


def run_mori(
    rankID,
    world_size,
    E,
    tokens,
    topk_weights,
    topk_ids,
    w1,
    w2,
    w1_scale,
    w2_scale,
    w1_zeros,
    w2_zeros,
    quant_type,
    ref_out,
):
    token_num = tokens.shape[0]
    hdim = tokens.shape[-1]
    topk = topk_weights.shape[-1]
    dtype = tokens.dtype
    device = torch.device(f"cuda:{rankID}")
    torch.cuda.set_device(device)
    init_dist_env(world_size, rankID)
    tokens = tokens.to(device)
    topk_weights = topk_weights.to(device)
    topk_ids = topk_ids.to(device)
    w1 = w1.to(device)
    w2 = w2.to(device)
    w1_scale = w1_scale.to(device) if w1_scale is not None else None
    w2_scale = w2_scale.to(device) if w2_scale is not None else None
    w1_zeros =  w1_zeros.to(device) if w1_zeros is not None else None
    w2_zeros = w2_zeros.to(device) if w2_zeros is not None else None
    scale = None


    print(f"=============RANK:{device} tokes_before_dispatch===============\n {tokens}")
    # init dist
    world_group = torch.distributed.group.WORLD
    assert world_group is not None
    torch._C._distributed_c10d._register_process_group("default", world_group)
    mori.shmem.shmem_torch_process_group_init("default")
    mori_config = mori.ops.EpDispatchCombineConfig(
        data_type=tokens.dtype,
        rank=rankID,
        world_size=world_size,
        hidden_dim=hdim,
        scale_dim=scale.shape[-1] if scale is not None else 0,
        scale_type_size=scale.dtype.itemsize if scale is not None else 0,
        max_token_type_size=dtype.itemsize,
        max_num_inp_token_per_rank=2
        * 8192
        * 1024
        // tokens.dtype.itemsize
        // hdim
        * 2,
        num_experts_per_rank=E // world_size,
        num_experts_per_token=topk,
    )
    mori_op = mori.ops.EpDispatchCombineOp(mori_config)
    dist.barrier()
    (
        dispatch_output,
        dispatch_weights,
        dispatch_scale,
        dispatch_ids,
        dispatch_recv_token_num,
    ) = mori_op.dispatch(tokens, topk_weights, None, topk_ids)
    torch.cuda.synchronize()
    # src_token_pos = mori_op.get_dispatch_src_token_pos().cpu()
    # src_token_num = src_token_pos.shape[0]
    # src_token_order = torch.sort(src_token_pos)[1].cpu()
    # print(
    #     f"{rankID=} {src_token_pos=} {src_token_order=}"
    # )
    # dispatch_ids = dispatch_ids[: src_token_num].to(dtypes.i32)[src_token_order]
    # dispatch_output = dispatch_output[: src_token_num][src_token_order]
    # dispatch_weights = dispatch_weights[: src_token_num][src_token_order]
    # dispatch_scale = dispatch_scale[: src_token_num][src_token_order]

    expert_mask = torch.zeros((E,), dtype=dtypes.i32, device=device)
    expert_mask[E // world_size * rankID : E // world_size * (rankID + 1)] = 1
    #triton version
    indices = expert_mask.cumsum(0, dtype=dtypes.i32) - 1
    e_map = torch.where(expert_mask == 0, torch.tensor(-1, dtype=dtypes.i32, device="cuda"), expert_mask)
    e_map = torch.where(e_map == 1, indices, e_map)
    # print(f"rank:{device}  {e_map=}")
    print(f"######### rank:{device} dispatch done############\n  recv_token:{dispatch_output} \n")
    # print(f"-------------rank:{device} dispatch_output_shape:{dispatch_output.shape}   ori_tokens:{tokens.shape}, {dispatch_recv_token_num=} {dispatch_weights.shape} {dispatch_ids.shape}")
    moe_input = dispatch_output[:dispatch_recv_token_num]
    moe_weights = dispatch_weights[:dispatch_recv_token_num]
    moe_ids = dispatch_ids[:dispatch_recv_token_num]
    
    # print(f"2222222222  {dispatch_output.shape} {dispatch_weights.shape} {dispatch_ids.shape}")
    # print(f"++++++++++++++++\n check_rank:{device}, w1:{w1}, w1_scale:{w1_scale}, w1_zeros:{w1_zeros} ")
    out= fused_experts_impl(
        moe_input,
        w1,
        w2,
        moe_weights,
        moe_ids,
        False,
        "silu",
        False,
        False,
        False,
        True,
        False,
        False,
        E,
        e_map,
        w1_scale,
        w2_scale,
        w1_zeros,
        w2_zeros,
        None,
        None,
        [0, 64],
    )
    print(f"######### rank:{device} moe done############  \n output:{out}")
    # print(f"rank {rankID} us={us:.4f}")

    # aiter.destroy_dist_env()
    # return out[:src_token_num].cpu()
    dist.barrier()
    combine_output,combine_weight = mori_op.combine(out, None, topk_ids)
    dist.barrier()
    # print(f"{rankID=} {combine_output.shape=} {combine_output.dtype=} {out.dtype=}")
    destroy_dist_env()
    print(f"rank:{device} combine:{combine_output} ref_out:{ref_out.to(device)} ")
    checkAllclose(ref_out.to(device), combine_output[:token_num], msg=f"rank:{device}")
    # print(f"33333333333333 {device=}   Combine Done      {combine_weight[:token_num]}")
    return combine_output[:token_num].cpu()


def weight_per_128x128_quant(weight, quant_dtype):
    E, dim1, dim2 = weight.shape
    weight_blocks = weight.view(
        E, dim1 // 128, 128, dim2 // 128, 128
    )  # [E, num_blocks_dim1, 128, num_blocks_dim2, 128]
    weight_blocks = weight_blocks.permute(
        0, 1, 3, 2, 4
    ).contiguous()  # [E, num_blocks_dim1, num_blocks_dim2, 128, 128]
    weight_blocks = weight_blocks.view(E, -1, 128 * 128)  # [E, num_blocks, 128*128]
    weight_qt, weight_scale = aiter.pertoken_quant(
        weight_blocks, quant_dtype=quant_dtype
    )
    weight_qt = weight_qt.view(
        E, dim1 // 128, dim2 // 128, 128, 128
    )  # [E, num_blocks_dim1, num_blocks_dim2, 128, 128]
    weight_qt = weight_qt.permute(
        0, 1, 3, 2, 4
    ).contiguous()  # [E, num_blocks_dim1, 128, num_blocks_dim2, 128]
    weight_qt = weight_qt.view(E, dim1, dim2)  # [E, dim1, dim2]
    weight_scale = weight_scale.view(
        E, dim1 // 128, dim2 // 128
    )  # [E, num_blocks_dim1, num_blocks_dim2]
    return weight_qt, weight_scale

def awq_weight_quant(w1,w2, w1_ref, w2_ref, group_size, pack_factor,has_zp, dtype):
    # print(f"w1_shape:{w1.shape} w2_shape:{w2.shape}")
    weight_bits=4
    quant_type = scalar_types.uint4 if has_zp else scalar_types.uint4b8
    E = w1.shape[0]
    hdim = w2.shape[1]
    idim = w2.shape[2]
    w1_qweight = torch.empty((E, 2 * idim, hdim // pack_factor),
                             device="cuda",
                             dtype=torch.uint8)
    w2_qweight = torch.empty((E, hdim, idim // pack_factor),
                             device="cuda",
                             dtype=torch.uint8)
    w1_scales = torch.empty((E, 2 * idim, hdim // group_size),
                            device="cuda",
                            dtype=dtype)
    w2_scales = torch.empty((E, hdim, idim // group_size),
                            device="cuda",
                            dtype=dtype)

    # weight quant part
    w1_qzeros = torch.empty((E , 2 * idim// pack_factor, hdim// group_size ),
                            device="cuda",
                            dtype=torch.uint8)
    w2_qzeros = torch.empty((E, hdim // pack_factor, idim // group_size),
                            device="cuda",
                            dtype=torch.uint8)
    for i in range(E * 2):
        expert_id = i % E
        if i // E == 0:
            w, w_ref, w_qweight, w_scales, w_qzeros = \
                w1, w1_ref, w1_qweight, w1_scales, w1_qzeros
        else:
            w, w_ref, w_qweight, w_scales, w_qzeros = \
                w2, w2_ref, w2_qweight, w2_scales, w2_qzeros
        weight, qweight, scales, qzeros = quantize_weights(
            w[expert_id].T, quant_type, group_size, has_zp, False)
        weight = weight.T
        qweight = qweight.T.contiguous().to(torch.uint8)
        scales = scales.T
        if has_zp:
            qzeros = qzeros.T.contiguous().to(torch.uint8)
        if weight_bits == 4:
            qweight = qweight[:, 1::2] * 16 + qweight[:, ::2]   # 偶数列存储低4位，奇数列存储高4位
            if has_zp:
                qzeros = qzeros[1::2, :] * 16 + qzeros[::2, :]
        w_ref[expert_id] = weight
        w_qweight[expert_id] = qweight
        w_scales[expert_id] = scales
        if has_zp:
            w_qzeros[expert_id] = qzeros
    return w1_qweight, w1_scales, w1_qzeros, w1_ref, w2_qweight,w2_scales,w2_qzeros,w2_ref


def test_dispatch_combine(
    world_size, shape, dtype, E, topk, quant_type=aiter.QuantType.No
):
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "49373"
    mp.set_start_method("spawn", force=True)
    pool = mp.Pool(processes=world_size)

    quant_func = (
        aiter.get_torch_quant(quant_type)
        if quant_type != aiter.QuantType.per_128x128
        else weight_per_128x128_quant
    )

    tokenNum, hdim, idim = shape
    tokens = torch.randn((tokenNum, hdim), dtype=dtype, device="cuda") /10
    score = torch.randn((tokenNum, E), device="cuda", dtype=dtype)
    topk_weights, topk_ids = fused_topk(tokens, score, topk, True)
    print(f"=========   Genearte Topkweights&ids {topk_weights=} {topk_ids=}")
    # w1 = torch.randn(())
    w1 = torch.randn((E, 2 * idim, hdim), dtype=dtype, device="cuda") / 10
    w2 = torch.randn((E, hdim, idim), dtype=dtype, device="cuda") / 10
    w1_ref = w1.clone()
    w2_ref = w2.clone()

    # weight quant part
    w1_qweight, w1_scales, w1_qzeros, w1_ref, w2_qweight,w2_scales,w2_qzeros,w2_ref = awq_weight_quant(w1,w2, w1_ref, w2_ref,64, 2, True, dtype)
    # print(f"{w1_qweight=}, {w1_scales=}, {w1_qzeros=}, {w1_ref.T}")
    # if quant_type == aiter.QuantType.per_128x128:
    #     weight_per_128x128_quant(w1, quant_dtype=dtypes.fp8)

    # w1_qt, w1_scale = quant_func(w1, quant_dtype=dtypes.fp8)
    # w2_qt, w2_scale = quant_func(w2, quant_dtype=dtypes.fp8)
    # w1_qt = shuffle_weight(w1_qt)
    # w2_qt = shuffle_weight(w2_qt)

    tokens_dp = tokens.chunk(world_size)
    topk_weights_dp = topk_weights.chunk(world_size)
    topk_ids_dp = topk_ids.chunk(world_size)
    w1_ep = w1_qweight.chunk(world_size)
    w2_ep = w2_qweight.chunk(world_size)
    w1_scale_ep = (
        w1_scales.chunk(world_size) if w1_scales   is not None else [None] * world_size
    )
    w2_scale_ep = (
        w2_scales.chunk(world_size) if w2_scales is not None else [None] * world_size
    )
    
    w1_zeros_ep = (
        w1_qzeros.chunk(world_size) if w1_qzeros is not None else [None] * world_size
    )
    w2_zeros_ep = (
        w2_qzeros.chunk(world_size) if w2_qzeros is not None else [None] * world_size
    )
    # print(f"full_token:{tokens_dp}")
    ref_noep = torch_moe(
        tokens,
        w1_ref,
        w2_ref,
        topk_weights,
        topk_ids,
        expert_mask = None
    )
    ref_dp = ref_noep.chunk(world_size)
    print(f"full_token:{tokens_dp} \n ref_result:{ref_dp}")
    rets = []
    for i in range(world_size):
        rets.append(
            pool.apply_async(
                run_mori,
                args=(
                    i,
                    world_size,
                    E,
                    tokens_dp[i],
                    topk_weights_dp[i],
                    topk_ids_dp[i],
                    w1_ep[i],
                    w2_ep[i],
                    w1_scale_ep[i],
                    w2_scale_ep[i],
                    w1_zeros_ep[i],
                    w2_zeros_ep[i],
                    quant_type,
                    ref_dp[i]
                ),
            )
        )
    pool.close()
    pool.join()
    rets = [el.get() for el in rets]


    # ret_out = torch.cat(rets, dim=0)
    # checkAllclose(ref_noep, ret_out.to(ref), msg="total tokens:")

    # for i in range(world_size):
    #     print(f"rank:{i}, reference:{ref_dp[i]}, recv:{rets[i]}")
    #     checkAllclose(ref_dp[i], rets[i].to(ref_dp[i]), msg=f"rank:{i}")


l_dtype = ["bf16"]
l_shape = [(16, 7168, 2048)]
quant_types = [
    aiter.QuantType.No,
    aiter.QuantType.per_Token,
    aiter.QuantType.per_128x128,
][-2:]

parser = argparse.ArgumentParser(description="config input of test")
parser.add_argument(
    "-q",
    "--quant_type",
    type=str,
    choices=["No", "per_Token", "per_128x128"],
    nargs="?",
    const=None,
    default=None,
    help="quantization type",
)
parser.add_argument(
    "-d",
    "--dtype",
    type=str,
    choices=l_dtype,
    nargs="?",
    const=None,
    default=None,
    help="data type",
)
parser.add_argument(
    "-s",
    "--shape",
    type=dtypes.str2tuple,
    choices=l_shape,
    nargs="?",
    const=None,
    default=None,
    help="shape",
)


if __name__ == "__main__":
    mp.freeze_support()
    args = parser.parse_args()
    if args.dtype is None:
        l_dtype = [dtypes.d_dtypes[key] for key in l_dtype]
    else:
        l_dtype = [dtypes.d_dtypes[args.dtype]]
    if args.shape is not None:
        l_shape = [args.shape]
    if args.quant_type is not None:
        quant_types = [eval(f"aiter.QuantType.{args.quant_type}")]

    # for quant_type in quant_types:
        # for dtype in l_dtype:
        # for shape in l_shape:
    test_dispatch_combine(8, (16, 7168, 2048), dtypes.d_dtypes["bf16"], 16, 2, aiter.QuantType.No)
