# SPDX-License-Identifier: MIT
# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.

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
import numpy as np

from aiter import dtypes
from aiter.dist.communication_op import tensor_model_parallel_all_reduce
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

logger = logging.getLogger("aiter")

set_start_method("spawn", force=True)


def allreduce_custom(
    tp_size,
    pp_size,
    rankID,
    x,
    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(True)
    init_distributed_environment(
        world_size=tp_size,
        rank=rankID,
        distributed_init_method=distributed_init_method,
    )
    ensure_model_parallel_initialized(tp_size, pp_size)
    x = x.to(device)
    # dist.barrier(device_ids=[i for i in range(tp_size)])

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

    if withGraph:
        graph = torch.cuda.CUDAGraph()
        with graph_capture() as gc:
            with torch.cuda.graph(graph, stream=gc.stream):
                out = tensor_model_parallel_all_reduce(x)
        out.fill_(0)

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

        _, us = run_ca()
        out = (out, us)
    else:

        @perftest()
        def run_ca(x):
            return tensor_model_parallel_all_reduce(x)

        out = run_ca(x)

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


@benchmark()
def test_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)
    ref = torch.zeros(shape, dtype=dtype)
    element_size = ref.element_size()
    total_elements = np.prod(shape)
    bytes_torch = element_size * total_elements
    rets = []
    for i in range(tp_size):
        x = torch.randn(shape, dtype=dtype)
        ref += x
        rets.append(
            pool.apply_async(
                allreduce_custom,
                args=(tp_size, pp_size, i, x, withGraph, distributed_init_method),
            )
        )
    pool.close()
    pool.join()
    rets = [el.get() for el in rets]
    uss = []
    for out, us in rets:
        uss.append(int(us))
        bw = bytes_torch/1.0E3/us
        msg = f"test_allreduce_custom: {shape=} {dtype=} {withGraph=} {bytes_torch=} {us:>8.2f}us {bw:>8.2f}GB/s"
        checkAllclose(ref, out.to(ref), msg=msg)
    return {
        "uss": uss, 
        "avg_time": np.mean([us for _, us in rets]),
        "min_time": np.min([us for _, us in rets]),
        "max_time": np.max([us for _, us in rets]),
        "shape": shape,
        "dtype": str(dtype),
        "withGraph": withGraph,
        "tp_size": tp_size,
        "pp_size": pp_size
    }

l_dtype = ["fp32"] #, "bf16"
l_shape = [(1, 256), (1, 512), (1, 1024), (1, 2048), (1, 4096), (1,8192),(2, 8192),(4, 8192),(8, 8192),(16, 8192),(32, 8192),(64, 8192),(128, 8192),
           (512, 8192), (1024, 8192), (2048, 8192)]
l_withGraph = [True]

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",
)
parser.add_argument(
    "-g",
    "--withGraph",
    type=bool,
    nargs="?",
    const=None,
    default=None,
    help="whether with graph mode",
)


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]
    if args.withGraph is not None:
        l_withGraph = [args.withGraph]
    from prettytable import PrettyTable
    allreduce_table = PrettyTable()
    field_names = [
        "size(B)",
        "sharp",
        "dtype",
        "withGraph",
        "avgLat(us)",
        "avgBW(GB/s)",
        "bestLat(us)",
        "bestBW(GB/s)",
        "worstLat(us)",
        "worstBW(GB/s)",
    ]
    allreduce_table.title = "Custom All Reduce Performance"
    allreduce_table.field_names = field_names
    for withGraph in l_withGraph:
        for dtype in l_dtype:
            for shape in l_shape:
                ref = torch.zeros(1, dtype=dtype)
                element_size = ref.element_size()
                total_elements = np.prod(shape)
                bytes_torch = element_size * total_elements
                result  = test_allreduce_custom(
                      8,
                      1,
                      shape,
                      dtype,
                      withGraph,
                      distributed_init_method=get_distributed_init_method(
                          get_ip(), get_open_port()
                      ),
                  )
                min_time = result.get('min_time', 0)
                max_time = result.get('max_time', 0)
                avg_time = result.get('avg_time', 0)
                bw_best = bytes_torch/1.0E3/min_time
                bw_worst = bytes_torch/1.0E3/max_time
                bw_average = bytes_torch/1.0E3/avg_time
                allreduce_table.add_row(
                    [
                        bytes_torch,
                        shape,
                        dtype,
                        "✓" if withGraph else "✗",
                        f"{avg_time:.2f}",
                        f"{bw_average:.2f}",
                        f"{min_time:.2f}",
                        f"{bw_best:.2f}",
                        f"{max_time:.2f}",
                        f"{bw_worst:.2f}",
                    ]
                )
    print(allreduce_table)