# SPDX-License-Identifier: MIT
 
import argparse
import logging
import os
from multiprocessing import Pool, freeze_support, set_start_method
from typing_extensions import Optional

import torch
import torch.distributed as dist
from aiter import QuantType, dtypes, get_hip_quant

from aiter.awq_gemm_asm import asm_awq_gemm_a16w4, asm_awq_reorder_and_repack
from aiter.dist.communication_op import tensor_model_parallel_all_reduce
from aiter.dist.custom_gemm_allreduce import CustomGemmAllreduce
from aiter.dist.parallel_state import (
    destroy_distributed_environment,
    destroy_model_parallel,
    ensure_model_parallel_initialized,
    get_tp_group,
    graph_capture,
    init_distributed_environment,
    set_custom_all_reduce,
)
from aiter.dist.utils import get_distributed_init_method, get_ip, get_open_port
from aiter.ops.triton.gemm_a16w4 import awq_reorder_and_repack, reverse_awq_order
from aiter.test_common import benchmark, checkAllclose, perftest

logger = logging.getLogger("aiter")

set_start_method("spawn", force=True)


def awq_gemm_allreduce_custom(
    tp_size,
    pp_size,
    rankID,
    input,
    output_size, # N
    qweight,
    qzeros,
    scales,
    group_size,
    withGraph=False,
    distributed_init_method: Optional[str] = None,
):
    device = torch.device(f"cuda:{rankID}")
    torch.cuda.set_device(device)
    # init
    logger.info(f"RANK: {rankID} {tp_size} init_process_group...")
    set_custom_all_reduce(False)
    init_distributed_environment(
        world_size=tp_size,
        rank=rankID,
        distributed_init_method=distributed_init_method,
    )
    ensure_model_parallel_initialized(tp_size, pp_size)
    
    local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", min(tp_size, 8)))
    group = get_tp_group().device_group

    # awq_gemm_allreduce_Op = CustomGemmAllreduce(group, "awq", 256)
    
    input = input.to(device)
    qweight = qweight.to(device)
    qzeros = qzeros.to(device)
    scales = scales.to(device)

    # qweight layer is [K, N] before repack
    awq_gemm_allreduce_Op = CustomGemmAllreduce(group, local_world_size, "awq", output_size)
    
    # asm results
    asm_qweight, asm_qzeros = asm_awq_reorder_and_repack(qweight, qzeros)
    # output_asm = asm_awq_gemm_a16w4(input, asm_qweight, scales, asm_qzeros)
    # dist.all_reduce(output_asm, group=group)
    output_asm_parallel = asm_awq_gemm_a16w4(input, asm_qweight, scales, asm_qzeros)
    output_asm = tensor_model_parallel_all_reduce(output_asm_parallel)

    # prepare for triton kernel
    qweight_repack, qzeros_repack = awq_reorder_and_repack(qweight, qzeros)

    # warmup and align all gpu
    # dist.all_reduce(torch.zeros(1).cuda(), group=group)
    tensor_model_parallel_all_reduce(torch.zeros(1).cuda())
    torch.cuda.synchronize()

    if withGraph:
        graph = torch.cuda.CUDAGraph()
        with graph_capture() as gc:
            with torch.cuda.graph(graph, stream=gc.stream):
                out = awq_gemm_allreduce_Op.gemm_allreduce(input, qweight_repack, scales, qzeros_repack)
        out.fill_(0)

        @perftest()
        def run_ca():
            graph.replay()

        _, us = run_ca()
        out = (out, us)
    else:
        @perftest()
        def run_ca(input, qweight_repack, scales, qzeros_repack):
            return awq_gemm_allreduce_Op.gemm_allreduce(input, qweight_repack, scales, qzeros_repack)

        out = run_ca(input, qweight_repack, scales, qzeros_repack)

    # destroy
    if dist.is_initialized():
        destroy_model_parallel()
        destroy_distributed_environment()
        torch.cuda.empty_cache()
    return out, output_asm


@benchmark()
def test_awq_gemm_allreduce_custom(
    tp_size,
    pp_size,
    shape,
    dtype,
    withGraph=False,
    distributed_init_method: Optional[str] = None,
):
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "49373"
    pool = Pool(processes=tp_size)
    rets = []
    
    # test input data
    M, K, N, G = shape
    input_rows = M
    input_cols = K
    input_dtype = torch.float16
    qweight_rows = input_cols
    qweight_cols = N // 8
    scales_rows = qweight_rows // G
    scales_cols = N
    scales_dtype = torch.float16
    qzeros_rows = scales_rows
    qzeros_cols = qweight_cols
    
    for i in range(tp_size):
        # make sure each rank's data is different
        torch.manual_seed(i)
        input = torch.rand((input_rows, input_cols),
                            dtype=input_dtype) * 0.001
        qweight = torch.randint(0,
                                torch.iinfo(torch.int32).max,
                                (qweight_rows, qweight_cols))
        qzeros = torch.randint(0,
                            torch.iinfo(torch.int32).max,
                            (qzeros_rows, qzeros_cols))
        scales = torch.rand((scales_rows, scales_cols),
                            dtype=scales_dtype)
        
        rets.append(
            pool.apply_async(
                awq_gemm_allreduce_custom,
                args=(
                    tp_size,
                    pp_size,
                    i,
                    input,
                    N,
                    qweight,
                    qzeros,
                    scales,
                    G,
                    withGraph,
                    distributed_init_method,
                ),
            )
        )
        

    pool.close()
    pool.join()
    rets = [el.get() for el in rets]
    rets_triton = [el[0] for el in rets]
    # contains only ref results
    rets_ref = [el[1] for el in rets]

    for i, (out, us) in enumerate(rets_triton):
        msg = f"test_awq_gemm_allreduce_custom: {shape=} {dtype=} {withGraph=} {us:>8.2f} us"
        checkAllclose(rets_ref[i].cpu(), out.cpu(), rtol=1e-03, atol=1e-02, msg=msg)

AWQ_GEMM_TEST_CASES = [
  # K,N.G
  (256, 7168, 64),
]

l_dtype = ["fp16"]
l_shape = [
    # "M, K, N, G"
    (M, *KNG) for M in [1, 2, 16, 256] for KNG in AWQ_GEMM_TEST_CASES
]

parser = argparse.ArgumentParser(description="config input of test")
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,
    nargs="?",
    const=None,
    default=None,
    help="shape. e.g. -s 128,8192",
)

if __name__ == "__main__":
    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]
    for dtype in l_dtype:
        for shape in l_shape:
            test_awq_gemm_allreduce_custom(
                8,
                1,
                shape,
                dtype,
                withGraph=True,
                distributed_init_method=get_distributed_init_method(
                    get_ip(), get_open_port()
                ),
            )
