# SPDX-License-Identifier: MIT
 
import os
from typing import Optional
import aiter
import torch
import torch.nn.functional as F
import torch.distributed as dist
import argparse
import itertools
from aiter import dtypes

from aiter.dist.parallel_state import (
    ensure_model_parallel_initialized,
    init_distributed_environment,
    set_custom_all_reduce,
    get_tp_group,
    graph_capture,
    destroy_model_parallel,
    destroy_distributed_environment,
)
from aiter.dist.utils import get_open_port, get_distributed_init_method, get_ip
from aiter.dist.communication_op import (
    tensor_model_parallel_all_reduce,
    tensor_model_parallel_fused_allreduce_rmsnorm,
)
from aiter.test_common import (
    checkAllclose,
    perftest,
    benchmark,
)
from multiprocessing import set_start_method, Pool, freeze_support
import logging

logger = logging.getLogger("aiter")

set_start_method("spawn", force=True)


def prebuild_rmsnorm_module(dtype):
    # Trigger JIT build in the parent process once to avoid multi-process
    # contention on file-baton lock in workers.
    if not torch.cuda.is_available():
        return
    dev = torch.device("cuda:0")
    x = torch.randn((1, 128), dtype=dtype, device=dev)
    ar_out = torch.randn_like(x)
    out = torch.empty_like(x)
    residual_out = torch.empty_like(x)
    weight = torch.randn((x.shape[-1],), dtype=dtype, device=dev)
    try:
        print(f"Prebuilding rmsnorm module for {dtype=}.")
        aiter.rmsnorm2d_fwd_with_add(
            out,
            ar_out,
            x,
            residual_out,
            weight,
            1e-6,
        )
        torch.cuda.synchronize(dev)
    except Exception as e:
        raise RuntimeError(
            f"Failed to prebuild rmsnorm module for {dtype=}. "
            "Please check aiter JIT build environment before running multi-process tests."
        ) from e
    print(f"Prebuilding rmsnorm module for {dtype=} done.")

def fused_ar_rmsnorm(
    tp_size,
    pp_size,
    rankID,
    x,
    weight,
    eps,
    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)
    weight = weight.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):
                res_out, out = tensor_model_parallel_fused_allreduce_rmsnorm(
                    x, x, weight, eps
                )
        out.fill_(0)
        res_out.fill_(0)

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

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

        @perftest()
        def run_ca(x):
            res_out, out = tensor_model_parallel_fused_allreduce_rmsnorm(
                x, x, weight, eps
            )
            return out

        out = run_ca(x)

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


def get_acc_value_with_cudagraph(
    tp_size,
    pp_size,
    rankID,
    x,
    weight,
    eps,
    loop_time=1,
    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)
    weight = weight.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()

    # out = torch.empty_like(x)
    graph = torch.cuda.CUDAGraph()
    with graph_capture() as gc:
        with torch.cuda.graph(graph, stream=gc.stream):
            # out = torch.empty_like(x)
            res_out, out = tensor_model_parallel_fused_allreduce_rmsnorm(
                x, x, weight, eps
            )
    out.fill_(0)

    def run_ca():
        graph.replay()
        rslt = out.clone()
        out.fill_(0)
        return rslt

    for i in range(loop_time):
        out = run_ca()

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


def get_acc_value_only(
    tp_size,
    pp_size,
    rankID,
    x,
    weight,
    eps,
    loop_time=1,
    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)
    weight = weight.to(device)
    # dist.barrier(device_ids=[i for i in range(tp_size)])

    # warmup and align all gpu
    group = get_tp_group().device_group
    torch.cuda.synchronize()

    for i in range(loop_time):
        res, out = tensor_model_parallel_fused_allreduce_rmsnorm(x, x, weight, eps)

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


def split_ar_rmsnorm(
    tp_size,
    pp_size,
    rankID,
    x,
    weight,
    eps,
    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)
    weight = weight.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):
                ar_out = tensor_model_parallel_all_reduce(x)
                # out = aiter.rms_norm(ar_out, weight, eps, 0)
                out = torch.empty_like(ar_out)
                residual_out = torch.empty_like(ar_out)
                aiter.rmsnorm2d_fwd_with_add(
                    out,
                    ar_out,
                    x,
                    residual_out,
                    weight,
                    eps)
        out.fill_(0)

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

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

        @perftest()
        def run_ca(x):
            ar_out = tensor_model_parallel_all_reduce(x)
            out = torch.empty_like(ar_out)
            residual_out = torch.empty_like(ar_out)
            aiter.rmsnorm2d_fwd_with_add(
                out,
                ar_out,
                x,
                residual_out,
                weight,
                eps
            )
            return out

        out = run_ca(x)

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


def split_vs_fused_ar_rmsnorm(
    tp_size,
    pp_size,
    rankID,
    x,
    weight,
    eps,
    withGraph=False,
    distributed_init_method: Optional[str] = None,
):
    device = torch.device(f"cuda:{rankID}")
    torch.cuda.set_device(device)
    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)
    weight = weight.to(device)

    group = get_tp_group().device_group
    dist.all_reduce(torch.zeros(1).cuda(), group=group)
    torch.cuda.synchronize()

    # For split-vs-fused numerical comparison, eager mode is sufficient and
    # avoids graph-capture + collective interaction risks in one worker lifecycle.
    @perftest()
    def run_split(inp):
        ar_out = tensor_model_parallel_all_reduce(inp)
        split_out = torch.empty_like(ar_out)
        residual_out = torch.empty_like(ar_out)
        aiter.rmsnorm2d_fwd_with_add(
            split_out,
            ar_out,
            inp,
            residual_out,
            weight,
            eps,
        )
        return split_out

    split_out, split_us = run_split(x)

    @perftest()
    def run_fused(inp):
        _, fused_out = tensor_model_parallel_fused_allreduce_rmsnorm(
            inp, inp, weight, eps
        )
        return fused_out

    fused_out, fused_us = run_fused(x)

    if dist.is_initialized():
        destroy_model_parallel()
        destroy_distributed_environment()
        torch.cuda.empty_cache()
    split_out_cpu = split_out.detach().cpu()
    fused_out_cpu = fused_out.detach().cpu()
    return rankID, split_out_cpu, fused_out_cpu, split_us, fused_us


@benchmark()
def test_split_ar_rmsnorm(
    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)
    rets = []
    cpu_rslt = []
    weight_list = []
    res_inp = []
    # print(type(shape[0]), shape[1], ref.device)
    m = shape[0]
    n = shape[1]
    eps = 1e-6
    for i in range(tp_size):
        x = torch.randn(shape, dtype=dtype)
        res_inp.append(x)
        ref += x
        weight = torch.randn((n,), dtype=dtype)
        weight_list.append(weight)
        rets.append(
            pool.apply_async(
                split_ar_rmsnorm,
                args=(
                    tp_size,
                    pp_size,
                    i,
                    x,
                    weight,
                    eps,
                    withGraph,
                    distributed_init_method,
                ),
            )
        )
    pool.close()
    pool.join()
    for i in range(tp_size):
        host_rslt = F.rms_norm(
            input=(ref + res_inp[i]),
            normalized_shape=(ref.shape[-1],),
            weight=weight_list[i],
            eps=eps,
        )
        cpu_rslt.append(host_rslt)
    rets = [el.get() for el in rets]
    for out, us in rets:
        msg = f"test_split_ar_rmsnorm: {shape=} {dtype=} {withGraph=} {us:>8.2f}"
        # print(cpu_rslt[out.device.index])
        checkAllclose(cpu_rslt[out.device.index], out.to(ref),rtol=0.03, atol=0.03, msg=msg)


@benchmark()
def test_fused_ar_rmsnorm(
    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)
    rets = []
    cpu_rslt = []
    weight_list = []
    res_inp = []
    # print(type(shape[0]), shape[1], ref.device)
    m = shape[0]
    n = shape[1]
    eps = 1e-6
    for i in range(tp_size):
        x = torch.randn(shape, dtype=dtype)
        # x = torch.ones(shape, dtype=dtype)
        res_inp.append(x)
        # print(f"device {i}, x[0][0] = {x[0][0]}")
        ref += x
        weight = torch.randn((n,), dtype=dtype)
        weight_list.append(weight)
        rets.append(
            pool.apply_async(
                fused_ar_rmsnorm,
                args=(
                    tp_size,
                    pp_size,
                    i,
                    x,
                    weight,
                    eps,
                    withGraph,
                    distributed_init_method,
                ),
            )
        )
    pool.close()
    pool.join()
    print(f"rslt[0][0] = {ref[0][0]}")

    for i in range(tp_size):
        host_rslt = F.rms_norm(
            input=(ref + res_inp[i]),
            normalized_shape=(ref.shape[-1],),
            weight=weight_list[i],
            eps=eps,
        )
        # host_rslt = ref + res_inp[i]
        cpu_rslt.append(host_rslt)

    rets = [el.get() for el in rets]
    for out, us in rets:
        msg = f"test_fused_ar_rmsnorm: {shape=} {dtype=} {withGraph=} {us:>8.2f}"
        # print(cpu_rslt[out.device.index])
        checkAllclose(cpu_rslt[out.device.index], out.to(ref),rtol=0.03, atol=0.03 ,msg=msg)
        # checkAllclose(ref, out.to(ref), msg=msg)


@benchmark()
def test_split_vs_fused_ar_rmsnorm(
    tp_size,
    pp_size,
    shape,
    dtype,
    withGraph=False,
    distributed_init_method: Optional[str] = None,
):
    eps = 1e-6
    n = shape[1]

    x_list = [torch.randn(shape, dtype=dtype) for _ in range(tp_size)]
    weight_list = [torch.randn((n,), dtype=dtype) for _ in range(tp_size)]

    compare_dist_init_method = distributed_init_method
    if compare_dist_init_method is None:
        compare_dist_init_method = get_distributed_init_method(get_ip(), get_open_port())

    pool = Pool(processes=tp_size)
    rets = []
    for i in range(tp_size):
        rets.append(
            pool.apply_async(
                split_vs_fused_ar_rmsnorm,
                args=(
                    tp_size,
                    pp_size,
                    i,
                    x_list[i].clone(),
                    weight_list[i].clone(),
                    eps,
                    withGraph,
                    compare_dist_init_method,
                ),
            )
        )
    pool.close()
    rank_rets = [el.get() for el in rets]
    pool.join()

    for rank_id, split_out, fused_out, split_us, fused_us in rank_rets:
        diff = (split_out - fused_out).abs()
        max_abs = diff.max().item()
        rel = diff / fused_out.abs().clamp_min(1e-6)
        rel_flat = rel.reshape(-1)
        max_rel, max_rel_idx = torch.max(rel_flat, dim=0)
        max_rel = max_rel.item()
        max_rel_idx = max_rel_idx.item()
        split_at_max_rel = split_out.reshape(-1)[max_rel_idx].item()
        fused_at_max_rel = fused_out.reshape(-1)[max_rel_idx].item()
        msg = (
            f"[split_vs_fused]: rank={rank_id} "
            f"split_us={split_us:>6.2f}, fused_us={fused_us:>6.2f}, "
            f"max_abs={max_abs:.8f}, max_rel={max_rel:.8f}, "
            f"idx={max_rel_idx}, split_val={split_at_max_rel:.8f}, fused_val={fused_at_max_rel:.8f}. "
        )
        print(msg)
        msg=f"split_vs_fused: {shape=} {dtype=} {withGraph=} rank={rank_id}"
        checkAllclose(split_out, fused_out, rtol=0.03, atol=0.03, msg=msg)


def acc_test(
    tp_size, pp_size, shape, dtype, 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)
    rets = []
    cpu_rslt = []
    weight_list = []
    # print(type(shape[0]), shape[1], ref.device)
    m = shape[0]
    n = shape[1]
    eps = 1e-6
    for i in range(tp_size):
        x = torch.randn(shape, dtype=dtype)
        ref += x
        weight = torch.randn((n,), dtype=dtype)
        weight_list.append(weight)
        rets.append(
            pool.apply_async(
                get_acc_value_only,
                args=(tp_size, pp_size, i, x, weight, eps, 1, distributed_init_method),
            )
        )
    pool.close()
    pool.join()

    ar_rslt = []
    for i, ret in enumerate(rets):
        rslt = ret.get()
        ar_rslt.append(rslt)
    for i in ar_rslt:
        checkAllclose(ref, i.to(ref))


def acc_test_cudagraph_on(
    tp_size,
    pp_size,
    shape,
    dtype,
    loop_time=1,
    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)
    rets = []
    cpu_rslt = []
    weight_list = []
    # print(type(shape[0]), shape[1], ref.device)
    m = shape[0]
    n = shape[1]
    eps = 1e-6
    for i in range(tp_size):
        x = torch.randn(shape, dtype=dtype)
        ref += x
        weight = torch.randn((n,), dtype=dtype)
        weight_list.append(weight)
        rets.append(
            pool.apply_async(
                get_acc_value_with_cudagraph,
                args=(
                    tp_size,
                    pp_size,
                    i,
                    x,
                    weight,
                    eps,
                    loop_time,
                    distributed_init_method,
                ),
            )
        )
    pool.close()
    pool.join()

    ar_rslt = []
    for i, ret in enumerate(rets):
        rslt = ret.get()
        ar_rslt.append(rslt)
    for i in ar_rslt:
        checkAllclose(ref, i.to(ref))


# def acc_test(tp_size, pp_size, shape, dtype):
#     os.environ["MASTER_ADDR"] = "127.0.0.1"
#     os.environ["MASTER_PORT"] = "49373"
#     pool = Pool(processes=tp_size)
#     ref = torch.zeros(shape, dtype=dtype)
#     rets = []
#     cpu_rslt = []
#     weight_list = []
#     # print(type(shape[0]), shape[1], ref.device)
#     m = shape[0]
#     n = shape[1]
#     eps = 1e-6
#     for i in range(tp_size):
#         x = torch.randn(shape, dtype=dtype)
#         print(f"device {i}, x[0][0] = {x[0][0]}")
#         ref += x
#         weight = torch.randn((n,), dtype=dtype)
#         weight_list.append(weight)
#         rets.append(
#             pool.apply_async(get_acc_value_only, args=(tp_size, pp_size, i, x, weight, eps))
#         )
#     pool.close()
#     pool.join()
#     for i in range(tp_size):
#         host_rslt = F.rms_norm(
#             input=ref, normalized_shape=(ref.shape[-1],), weight=weight_list[i], eps=eps
#         )
#         cpu_rslt.append(host_rslt)
#
#     ar_rslt = []
#     for i, ret in enumerate(rets):
#         rslt = ret.get()
#         ar_rslt.append(rslt)
#     for i in range(len(ar_rslt)):
#         checkAllclose(cpu_rslt[i], ar_rslt[i].to(ref))

l_dtype = ["bf16"]
l_shape = [(6,7168)]
l_tp = [8]
l_pp = [1]
l_graph = [True, False]

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(
    "-t",
    "--tp",
    type=int,
    nargs="?",
    const=None,
    default=None,
    help="tp num. e.g. -t 8",
)

parser.add_argument(
    "-p",
    "--pp",
    type=int,
    nargs="?",
    const=None,
    default=None,
    help="tp num. e.g. -p 1",
)

parser.add_argument(
    "-g",
    "--graphon",
    type=int,
    nargs="?",
    const=None,
    default=None,
    help="open cudagraph. e.g. -g 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.shape is not None:
        l_shape = [args.shape]
    if args.tp is not None:
        l_tp = [args.tp]
    if args.pp is not None:
        l_pp = [args.pp]
    if args.graphon is not None:
        print(args.graphon)
        l_graph = [args.graphon]

    # Prebuild split-path JIT kernel to avoid workers hanging in JIT lock wait.
    for dtype in l_dtype:
        prebuild_rmsnorm_module(dtype)

    for dtype, shape, tp, pp, graph_on in itertools.product(
        l_dtype, l_shape, l_tp, l_pp, l_graph
    ):
        test_split_ar_rmsnorm(
            tp,
            pp,
            shape,
            dtype,
            withGraph=graph_on,
            distributed_init_method=get_distributed_init_method(
                get_ip(), get_open_port()
            ),
        )
        test_fused_ar_rmsnorm(
            tp,
            pp,
            shape,
            dtype,
            withGraph=graph_on,
            distributed_init_method=get_distributed_init_method(
                get_ip(), get_open_port()
            ),
        )
        
        test_split_vs_fused_ar_rmsnorm(
            tp,
            pp,
            shape,
            dtype,
            withGraph=graph_on,
            distributed_init_method=get_distributed_init_method(
                get_ip(), get_open_port()
            ),
        )
