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

import torch
import torch.distributed as dist
# from aiter import get_hip_quant, QuantType
from aiter import dtypes
from aiter import pertoken_quant
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.test_common import benchmark, checkAllclose, perftest
from aiter.tuned_gemm import tgemm
from einops import rearrange
from vllm.model_executor.layers.quantization.utils.int8_utils import (
    per_token_group_quant_int8,
)
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
    per_token_group_quant_fp8,
)

logger = logging.getLogger("aiter")

set_start_method("spawn", force=True)

def _build_inputs(
    rank: int,
    shape,
    dtype,
    quant_dtype,
    block_shape,
):
    """Create per-rank input/weight/scales locally to avoid cross-host pickling."""
    m, n, k = shape
    block_shape_n, block_shape_k = block_shape
    torch.manual_seed(rank)

    x = torch.randn((m, k), dtype=torch.float16, device="cuda") * 0.001
    if quant_dtype == dtypes.fp8:
        x, x_scale = per_token_group_quant_fp8(x, block_shape_k)
    elif quant_dtype == dtypes.i8:
        x, x_scale = per_token_group_quant_int8(x, block_shape_k)
    else:
        raise ValueError(f"Unsupported quant dtype: {quant_dtype}")

    block_shape_n_aligned = math.ceil(n / block_shape_n) * block_shape_n
    w2 = torch.rand((1, block_shape_n_aligned, k), dtype=dtype, device="cuda")
    tmp = rearrange(
        w2.view(
            -1,
            w2.shape[1] // block_shape_n,
            block_shape_n,
            math.ceil(w2.shape[2] / block_shape_k),
            block_shape_k,
        ),
        "e num_blk_n blk_n num_blk_k blk_k -> e num_blk_n num_blk_k (blk_n blk_k)",
    ).contiguous()
    w2_qweight, w2_scales = pertoken_quant(tmp, quant_dtype=quant_dtype)
    weight = rearrange(
        w2_qweight.view(
            -1,
            w2.shape[1] // block_shape_n,
            w2.shape[2] // block_shape_k,
            block_shape_n,
            block_shape_k,
        ),
        "e num_blk_n num_blk_k blk_n blk_k -> e (num_blk_n blk_n) (num_blk_k blk_k)",
    ).contiguous()

    weight = weight.view(weight.shape[1], weight.shape[2])
    weight = weight[:n]
    w_scale = w2_scales.view(w2_scales.shape[1], w2_scales.shape[2])
    return x, weight, x_scale, w_scale


# x, weight, x_scale, w_scale, withGraph
def w8a8_gemm_allreduce_custom(
    tp_size,
    pp_size,
    rankID,
    shape,
    dtype,
    quant_dtype,
    block_size,
    withGraph=False,
    distributed_init_method: Optional[str] = None,
):
    world_rank = rankID if rankID is not None else int(os.environ.get("RANK", 0))
    world_size = int(os.environ.get("WORLD_SIZE", tp_size))
    local_rank = int(os.environ.get("LOCAL_RANK", world_rank))
    local_world_size = int(os.environ.get("LOCAL_WORLD_SIZE", min(tp_size, 8))) # assume 1 node has 8 ranks

    device = torch.device(f"cuda:{local_rank}")
    torch.cuda.set_device(device)
    # init
    logger.info(
        f"RANK: {world_rank}/{world_size} local_rank={local_rank} init_process_group..."
    )
    set_custom_all_reduce(False)
    init_distributed_environment(
        world_size=world_size,
        rank=world_rank,
        local_rank=local_rank,
        distributed_init_method=distributed_init_method,
    )
    ensure_model_parallel_initialized(tp_size, pp_size)

    group = get_tp_group().device_group

    x, weight, x_scale, w_scale = _build_inputs(
        world_rank, shape, dtype, quant_dtype, block_size
    )
    x = x.to(device)
    weight = weight.to(device)
    x_scale = x_scale.to(device)
    w_scale = w_scale.to(device)

    # qweight layer is [K, N] before repack
    N, _ = weight.shape
    w8a8_gemm_allreduce_Op = CustomGemmAllreduce(group, local_world_size, "blockwise_int8", N)

    # asm results
    bias = None
    output_asm_parallel = tgemm.scale_mm(
        x, weight, bias, dtype, x_scale, w_scale, scale_type=2
    )
    # dist.all_reduce(output_asm, group=group)
    output_asm = tensor_model_parallel_all_reduce(output_asm_parallel)

    # for debug
    # print(f"{rankID=}")
    # print(f"output_asm_parallel: {output_asm_parallel[:1, ::384]}")

    # 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 = w8a8_gemm_allreduce_Op.gemm_allreduce(
                    x, weight, x_scale, w_scale, block_size, dtype
                )
        out.fill_(0)

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

        _, us = run_ca()
        out = (out, us)
    else:
        @perftest()
        def run_ca(input, weight, x_scale, w_scale, block_size, dtype):
            return w8a8_gemm_allreduce_Op.gemm_allreduce(
                input, weight, x_scale, w_scale, block_size, dtype
            )

        # out = run_ca(x, weight, x_scale, w_scale, block_size, dtype)
        out = w8a8_gemm_allreduce_Op.gemm_allreduce(x, weight, x_scale, w_scale, block_size, dtype)
        # for debug
        # out = (out, 555)

    # print(f"{x=} {weight=} {out=}")

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


@benchmark()
def test_w8a8_gemm_allreduce_custom(
    tp_size,
    pp_size,
    shape,
    dtype,
    quant_dtype=dtypes.i8,
    withGraph=False,
    distributed_init_method: Optional[str] = None,
):
    block_shape = (128, 128)

    # If launched via torchrun/mpirun with env:// rendezvous, run rank locally.
    if "WORLD_SIZE" in os.environ:
        world_size = int(os.environ["WORLD_SIZE"])
        cur_rank = int(os.environ["LOCAL_RANK"])
        assert (
            world_size == tp_size
        ), f"Expected tp_size {tp_size} to match WORLD_SIZE {world_size} in multi-node run."
        out, ref = w8a8_gemm_allreduce_custom(
            tp_size,
            pp_size,
            cur_rank,
            shape,
            dtype,
            quant_dtype,
            block_shape,
            withGraph=withGraph,
            # torchrun provides env:// rendezvous (MASTER_ADDR/PORT/RANK/WORLD_SIZE)
            distributed_init_method="env://",
        )
        msg = f"test_w8a8_gemm_allreduce_custom: {shape=} {dtype=} {quant_dtype=} {withGraph=} {out[1]:>8.2f} us"
        # checkAllclose(ref.cpu(), out[0].cpu(), rtol=1e-02, atol=1e-02, msg=msg)
        N_per_rank = ref.shape[1] // world_size
        # if cur_rank == 1:
        #     print(f"{out[0][:ref.shape[0], N_per_rank*cur_rank:] }")
        #     # checkAllclose(ref[:, N_per_rank*cur_rank:N_per_rank*(cur_rank+1)].cpu(), out[0][:ref.shape[0], :].cpu().to(torch.float16), rtol=1e-02, atol=1e-02, msg=msg)
        # if cur_rank != 1:
        #     return
        checkAllclose(ref[:, N_per_rank*cur_rank:N_per_rank*(cur_rank+1)].cpu(), out[0][:ref.shape[0], N_per_rank*cur_rank:N_per_rank*(cur_rank+1)].cpu().to(torch.float16), rtol=1e-02, atol=1e-02, msg=msg)
        return

    # Fallback: single-node spawn (original behavior).
    os.environ["MASTER_ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "49373"
    pool = Pool(processes=tp_size)
    rets = []
    
    block_shape = (128, 128)

    # test input data
    m, n, k = shape

    block_shape_n, block_shape_k = block_shape
    scale_n = (n + block_shape_n - 1) // block_shape_n
    scale_k = (k + block_shape_k - 1) // block_shape_k

    block_shape = (128, 128)
    
    # test input data
    m, n, k = shape
    
    block_shape_n, block_shape_k = block_shape
    scale_n = (n + block_shape_n - 1) // block_shape_n
    scale_k = (k + block_shape_k - 1) // block_shape_k
    
    for i in range(tp_size):
        rets.append(
            pool.apply_async(
                w8a8_gemm_allreduce_custom,
                args=(
                    tp_size,
                    pp_size,
                    i,
                    shape,
                    dtype,
                    quant_dtype,
                    block_shape,
                    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_w8a8_gemm_allreduce_custom: {shape=} {dtype=} {quant_dtype=} {withGraph=} {us:>8.2f} us"
        checkAllclose(rets_ref[i].cpu(), out.cpu(), rtol=1e-02, atol=1e-02, msg=msg)

W8A8_BLOCK_GEMM_TEST_CASES = [
  # n, k
  (1536, 7168),
#   (3072, 1536),
#   (576, 7168),
#   (7168, 256),
#   (7168, 2048),
#   (4608, 7168),
#   (7168, 2304),
#   (512, 7168),
#   (4096, 512),
]

l_dtype = ["fp16"]
l_quant_dtype = ["i8"]
l_shape = [
    # (M, *NK) for M in [16, 32, 64, 128, 256] for NK in W8A8_BLOCK_GEMM_TEST_CASES
    (M, *NK) for M in [256] for NK in W8A8_BLOCK_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(
    "-q",
    "--quant-dtype",
    type=str,
    choices=l_quant_dtype,
    nargs="?",
    const=None,
    default=None,
    help="quantization data type",
)

parser.add_argument(
    "-s",
    "--shape",
    type=dtypes.str2tuple,
    nargs="?",
    const=None,
    default=None,
    help="shape. e.g. -s 128,8192",
)
parser.add_argument(
    "--tp-size",
    type=int,
    default=None,
    help="tensor-parallel size (defaults to WORLD_SIZE if set, else 8)",
)
parser.add_argument(
    "--pp-size",
    type=int,
    default=1,
    help="pipeline-parallel size (defaults to 1)",
)

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.quant_dtype is None:
        l_quant_dtype = [dtypes.d_dtypes[key] for key in l_quant_dtype]
    else:
        l_quant_dtype = [dtypes.d_dtypes[args.quant_dtype]]
    if args.shape is not None:
        l_shape = [args.shape]
    tp_default = int(os.environ.get("WORLD_SIZE", 8))
    tp_size_arg = args.tp_size if args.tp_size is not None else tp_default
    for dtype in l_dtype:
        for shape in l_shape:
            for quant_dtype in l_quant_dtype:
                test_w8a8_gemm_allreduce_custom(
                    tp_size_arg,
                    args.pp_size,
                    shape,
                    dtype,
                    quant_dtype,
                    withGraph=True,
                    distributed_init_method=get_distributed_init_method(
                        get_ip(), get_open_port()
                    ),
                )
