"""
Unit tests for host collective on stream functionality in nvshmem.core
"""
try:
    import torch
    from torch import float32
    _torch_enabled = True
except:
    float32 = None
    _torch_enabled = False

try:
    import cupy
    _cupy_enabled = True
except:
    _cupy_enabled = False


from utils import uid_init, mpi_init
import argparse
import os

import nvshmem.core

from cuda.core.experimental import Device, system

from mpi4py import MPI

all_types_cupy = ["float16", "float32", "float64", "uint8", "int8", "int16", "int32", "int64", "bool"]
all_types_torch = [torch.float16, torch.bfloat16, torch.float32, torch.uint8, torch.int16, torch.int32, torch.int64, torch.bool]
all_types_nvshmem = ["half", "bfloat16", "uint8", "int8", "int16", "int32", "int64", "double", "float"]
all_ops = ["sum", "min", "max"]

###
# Collectives with no data
###

def test_barrier():
    print("Testing barrier")
    if not _cupy_enabled:
        print("WARNING: CuPy not found. Not running CuPy collective test")
        return
    # Get a buffer
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    print("Rank:", local_rank_per_node)
    dev = Device(local_rank_per_node)
    stream = dev.create_stream()
    # Print dst, src before
    print(f"Before barrier from {nvshmem.core.my_pe()}")
    nvshmem.core.barrier(nvshmem.core.Teams.TEAM_NODE, stream=stream)

    # Print dst, src after
    print(f"After barrier from PE {nvshmem.core.my_pe()}")
    print("Done testing barrier")

def test_team_sync():
    print("Testing team sync")
    if not _cupy_enabled:
        print("WARNING: CuPy not found. Not running CuPy collective test")
        return
    # Get a buffer
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    print("Rank:", local_rank_per_node)
    dev = Device(local_rank_per_node)
    stream = dev.create_stream()
    # Print dst, src before
    print(f"Before team sync from {nvshmem.core.my_pe()}")
    nvshmem.core.sync(nvshmem.core.Teams.TEAM_NODE, stream=stream)

    # Print dst, src after
    print(f"After team sync from PE {nvshmem.core.my_pe()}")
    print("Done testing team sync")

###
# CuPy Collectives
###

def test_alltoall_cupy():
    print("Testing CuPy array Alltoall")
    if not _cupy_enabled:
        print("WARNING: CuPy not found. Not running CuPy collective test")
        return
    # Get a buffer
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    print("Rank:", local_rank_per_node)
    dev = Device(local_rank_per_node)
    stream = dev.create_stream()

    for dtype in all_types_cupy:
        arr_src = nvshmem.core.array((2,2), dtype=dtype)
        # We're going to alltoall from src to dst, so it needs to be N times as big
        arr_dst = nvshmem.core.array((nvshmem.core.n_pes(),2,2), dtype=dtype)
        # Zero out the dest and set src to my_pe
        arr_dst[:] = 0
        # as usual, local_rank_per_node + 1 so there's no zeroes
        arr_src[:] = local_rank_per_node + 1

        # Print dst, src before
        print(f"Dest before collective from PE {nvshmem.core.my_pe()}:", arr_dst)
        print(f"Src before collective from PE {nvshmem.core.my_pe()}:", arr_src)

        nvshmem.core.alltoall(nvshmem.core.Teams.TEAM_NODE, arr_dst, arr_src, stream=stream)
        # Explicit sync because we didn't set a CuPy stream
        dev.sync()

        # Print dst, src after
        print(f"Dest after collective from PE {nvshmem.core.my_pe()}:", arr_dst)
        print(f"Src after collective from PE {nvshmem.core.my_pe()}:", arr_src)

        # Free Buffers
        nvshmem.core.free_array(arr_src)
        nvshmem.core.free_array(arr_dst)
        print(f"Done testing CuPy array Alltoall with type={dtype}")

def test_fcollect_cupy():
    print("Testing CuPy array Fcollect")
    if not _cupy_enabled:
        print("WARNING: CuPy not found. Not running CuPy collective test")
        return
    # Get a buffer
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    print("Rank:", local_rank_per_node)
    dev = Device(local_rank_per_node)
    stream = dev.create_stream()

    for dtype in all_types_cupy:
        arr_src = nvshmem.core.array((2,2), dtype=dtype)
        # We're going to Fcollect from src to dst, so it needs to be N times as big
        arr_dst = nvshmem.core.array((nvshmem.core.n_pes(),2,2), dtype=dtype)
        # Zero out the dest and set src to my_pe
        arr_dst[:] = 0
        # as usual, local_rank_per_node + 1 so there's no zeroes
        arr_src[:] = local_rank_per_node + 1

        # Print dst, src before
        print(f"Dest before collective from PE {nvshmem.core.my_pe()}:", arr_dst)
        print(f"Src before collective from PE {nvshmem.core.my_pe()}:", arr_src)

        nvshmem.core.fcollect(nvshmem.core.Teams.TEAM_NODE, arr_dst, arr_src, stream=stream)
        # Explicit sync because we didn't set a CuPy stream
        dev.sync()

        # Print dst, src after
        print(f"Dest after collective from PE {nvshmem.core.my_pe()}:", arr_dst)
        print(f"Src after collective from PE {nvshmem.core.my_pe()}:", arr_src)

        # Free Buffers
        nvshmem.core.free_array(arr_src)
        nvshmem.core.free_array(arr_dst)
        print(f"Done testing CuPy array Fcollect with type={dtype}")

    print("Done testing CuPy array Fcollect")

def test_reduce_cupy():
    print("Testing CuPy array Reduce")
    if not _cupy_enabled:
        print("WARNING: CuPy not found. Not running CuPy collective test")
        return
    # Get a buffer
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    print("Rank:", local_rank_per_node)
    dev = Device(local_rank_per_node)
    stream = dev.create_stream()

    for dtype in all_types_cupy:
        for op in all_ops:
            arr_src = nvshmem.core.array((2,2), dtype=dtype)
            arr_dst = nvshmem.core.array((2,2), dtype=dtype)
            # Zero out the dest and set src to my_pe
            arr_dst[:] = 0
            # as usual, local_rank_per_node + 1 so there's no zeroes
            arr_src[:] = local_rank_per_node + 1

            # Print dst, src before
            print(f"Dest before collective from PE {nvshmem.core.my_pe()}:", arr_dst)
            print(f"Src before collective from PE {nvshmem.core.my_pe()}:", arr_src)

            nvshmem.core.reduce(nvshmem.core.Teams.TEAM_WORLD, arr_dst, arr_src, op, stream=stream)
            # Explicit sync because we didn't set a CuPy stream
            # Important! Make sure you sync the stream that you launched the nvshmem collective on
            # (or the whole device like we do in this example)
            # before you use the results of the collective
            dev.sync()

            # Print dst, src after
            print(f"Dest after collective from PE {nvshmem.core.my_pe()}:", arr_dst)
            print(f"Src after collective from PE {nvshmem.core.my_pe()}:", arr_src)

            # Free Buffers
            nvshmem.core.free_array(arr_src)
            nvshmem.core.free_array(arr_dst)
            print(f"Done testing CuPy Array Reduce with dtype {dtype} op {op}")
    print("Done testing CuPy array Reduce")


def test_reducescatter_cupy():
    print("Testing CuPy array Reducescatter")
    if not _cupy_enabled:
        print("WARNING: CuPy not found. Not running CuPy collective test")
        return
    # Get a buffer
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    print("Rank:", local_rank_per_node)
    dev = Device(local_rank_per_node)
    stream = dev.create_stream()

    for dtype in all_types_cupy:
        for op in all_ops:
            arr_src = nvshmem.core.array((2,2, nvshmem.core.n_pes()), dtype=dtype)
            # We're going to reducescatter from src to dst, so it needs to be N times as big
            arr_dst = nvshmem.core.array((2,2), dtype=dtype)
            # Zero out the dest and set src to my_pe
            arr_dst[:] = 0
            # as usual, local_rank_per_node + 1 so there's no zeroes
            arr_src[:] = local_rank_per_node + 1

            # Print dst, src before
            print(f"Dest before collective from PE {nvshmem.core.my_pe()}:", arr_dst)
            print(f"Src before collective from PE {nvshmem.core.my_pe()}:", arr_src)

            nvshmem.core.reducescatter(nvshmem.core.Teams.TEAM_WORLD, arr_dst, arr_src, op, stream=stream)
            # Explicit sync because we didn't set a CuPy stream
            # Important! Make sure you sync the stream that you launched the nvshmem collective on
            # (or the whole device like we do in this example)
            # before you use the results of the collective
            dev.sync()

            # Print dst, src after
            print(f"Dest after collective from PE {nvshmem.core.my_pe()}:", arr_dst)
            print(f"Src after collective from PE {nvshmem.core.my_pe()}:", arr_src)

            # Free Buffers
            nvshmem.core.free_array(arr_src)
            nvshmem.core.free_array(arr_dst)
            print(f"Done testing CuPy Array Reducescatter with dtype {dtype} op {op}")
    print("Done testing CuPy array Reducescatter")

def test_broadcast_cupy():
    print("Testing CuPy array broadcast")
    if not _cupy_enabled:
        print("WARNING: CuPy not found. Not running CuPy collective test")
        return
    # Get a buffer
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    print("Rank:", local_rank_per_node)
    dev = Device(local_rank_per_node)
    stream = dev.create_stream()

    for dtype in all_types_cupy:
        for op in all_ops:
            arr_src = nvshmem.core.array((2,2), dtype=dtype)
            arr_dst = nvshmem.core.array((2,2), dtype=dtype)
            # Zero out the dest and set src to my_pe
            arr_dst[:] = 0
            # as usual, local_rank_per_node + 1 so there's no zeroes
            arr_src[:] = local_rank_per_node + 1

            # Print dst, src before
            print(f"Dest before collective from PE {nvshmem.core.my_pe()}:", arr_dst)
            print(f"Src before collective from PE {nvshmem.core.my_pe()}:", arr_src)

            nvshmem.core.broadcast(nvshmem.core.Teams.TEAM_WORLD, arr_dst, arr_src, root=0, stream=stream)
            # Explicit sync because we didn't set a CuPy stream
            # Important! Make sure you sync the stream that you launched the nvshmem collective on
            # (or the whole device like we do in this example)
            # before you use the results of the collective
            dev.sync()

            # Print dst, src after
            print(f"Dest after collective from PE {nvshmem.core.my_pe()}:", arr_dst)
            print(f"Src after collective from PE {nvshmem.core.my_pe()}:", arr_src)

            # Free Buffers
            nvshmem.core.free_array(arr_src)
            nvshmem.core.free_array(arr_dst)
            print(f"Done testing CuPy Array broadcast with dtype {dtype} op {op}")
    print("Done testing CuPy array broadcast")


###
# Torch Collectives
###

def test_alltoall_torch():
    print("Testing Torch array Alltoall")
    if not _torch_enabled:
        print("WARNING: Torch not found. Not running Torch collective test")
        return
    # Get a buffer
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    print("Rank:", local_rank_per_node)
    dev = Device(local_rank_per_node)
    stream = dev.create_stream()

    for dtype in all_types_torch:
        arr_src = nvshmem.core.tensor((2,2), dtype=dtype)
        # We're going to alltoall from src to dst, so it needs to be N times as big
        arr_dst = nvshmem.core.tensor((nvshmem.core.n_pes(),2,2), dtype=dtype)
        # Zero out the dest and set src to my_pe
        arr_dst[:] = 0
        # as usual, local_rank_per_node + 1 so there's no zeroes
        arr_src[:] = local_rank_per_node + 1

        # Print dst, src before
        print(f"Dest before collective from PE {nvshmem.core.my_pe()}:", arr_dst)
        print(f"Src before collective from PE {nvshmem.core.my_pe()}:", arr_src)

        nvshmem.core.alltoall(nvshmem.core.Teams.TEAM_NODE, arr_dst, arr_src, stream=stream)
        # Explicit sync because we didn't set a Torch stream
        dev.sync()

        # Print dst, src after
        print(f"Dest after collective from PE {nvshmem.core.my_pe()}:", arr_dst)
        print(f"Src after collective from PE {nvshmem.core.my_pe()}:", arr_src)

        # Free Buffers
        nvshmem.core.free_tensor(arr_src)
        nvshmem.core.free_tensor(arr_dst)
        print(f"Done testing Torch array Alltoall with type={dtype}")

def test_fcollect_torch():
    print("Testing Torch array Fcollect")
    if not _torch_enabled:
        print("WARNING: Torch not found. Not running Torch collective test")
        return
    # Get a buffer
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    print("Rank:", local_rank_per_node)
    dev = Device(local_rank_per_node)
    stream = dev.create_stream()

    for dtype in all_types_torch:
        arr_src = nvshmem.core.tensor((2,2), dtype=dtype)
        # We're going to Fcollect from src to dst, so it needs to be N times as big
        arr_dst = nvshmem.core.tensor((nvshmem.core.n_pes(),2,2), dtype=dtype)
        # Zero out the dest and set src to my_pe
        arr_dst[:] = 0
        # as usual, local_rank_per_node + 1 so there's no zeroes
        arr_src[:] = local_rank_per_node + 1

        # Print dst, src before
        print(f"Dest before collective from PE {nvshmem.core.my_pe()}:", arr_dst)
        print(f"Src before collective from PE {nvshmem.core.my_pe()}:", arr_src)

        nvshmem.core.fcollect(nvshmem.core.Teams.TEAM_NODE, arr_dst, arr_src, stream=stream)
        # Explicit sync because we didn't set a Torch stream
        dev.sync()

        # Print dst, src after
        print(f"Dest after collective from PE {nvshmem.core.my_pe()}:", arr_dst)
        print(f"Src after collective from PE {nvshmem.core.my_pe()}:", arr_src)

        # Free Buffers
        nvshmem.core.free_tensor(arr_src)
        nvshmem.core.free_tensor(arr_dst)
        print(f"Done testing Torch array Fcollect with type={dtype}")

    print("Done testing Torch array Fcollect")

def test_reduce_torch():
    print("Testing torch tensor Reduce")
    if not _torch_enabled:
        print("WARNING: torch not found. Not running torch collective test")
        return
    # Get a buffer
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    print("Rank:", local_rank_per_node)
    dev = Device(local_rank_per_node)
    stream = dev.create_stream()

    for dtype in all_types_torch:
        for op in all_ops:
            arr_src = nvshmem.core.tensor((2,2), dtype=dtype)
            arr_dst = nvshmem.core.tensor((2,2), dtype=dtype)
            # Zero out the dest and set src to my_pe
            arr_dst[:] = 0
            # as usual, local_rank_per_node + 1 so there's no zeroes
            arr_src[:] = local_rank_per_node + 1

            # Print dst, src before
            print(f"Dest before collective from PE {nvshmem.core.my_pe()}:", arr_dst)
            print(f"Src before collective from PE {nvshmem.core.my_pe()}:", arr_src)

            nvshmem.core.reduce(nvshmem.core.Teams.TEAM_WORLD, arr_dst, arr_src, op, stream=stream)
            # Explicit sync because we didn't set a torch stream
            # Important! Make sure you sync the stream that you launched the nvshmem collective on
            # (or the whole device like we do in this example)
            # before you use the results of the collective
            dev.sync()

            # Print dst, src after
            print(f"Dest after collective from PE {nvshmem.core.my_pe()}:", arr_dst)
            print(f"Src after collective from PE {nvshmem.core.my_pe()}:", arr_src)

            # Free Buffers
            nvshmem.core.free_tensor(arr_src)
            nvshmem.core.free_tensor(arr_dst)
            print(f"Done testing torch tensor Reduce with dtype {dtype} op {op}")
    print("Done testing torch tensor Reduce")

def test_reducescatter_torch():
    print("Testing torch tensor Reducescatter")
    if not _torch_enabled:
        print("WARNING: torch not found. Not running torch collective test")
        return
    # Get a buffer
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    print("Rank:", local_rank_per_node)
    dev = Device(local_rank_per_node)
    stream = dev.create_stream()

    for dtype in all_types_torch:
        for op in all_ops:
            arr_src = nvshmem.core.tensor((2,2, nvshmem.core.n_pes()), dtype=dtype)
            # We're going to reducescatter from src to dst, so it needs to be N times as big
            arr_dst = nvshmem.core.tensor((2,2), dtype=dtype)
            # Zero out the dest and set src to my_pe
            arr_dst[:] = 0
            # as usual, local_rank_per_node + 1 so there's no zeroes
            arr_src[:] = local_rank_per_node + 1

            # Print dst, src before
            print(f"Dest before collective from PE {nvshmem.core.my_pe()}:", arr_dst)
            print(f"Src before collective from PE {nvshmem.core.my_pe()}:", arr_src)

            nvshmem.core.reducescatter(nvshmem.core.Teams.TEAM_WORLD, arr_dst, arr_src, op, stream=stream)
            # Explicit sync because we didn't set a torch stream
            # Important! Make sure you sync the stream that you launched the nvshmem collective on
            # (or the whole device like we do in this example)
            # before you use the results of the collective
            dev.sync()

            # Print dst, src after
            print(f"Dest after collective from PE {nvshmem.core.my_pe()}:", arr_dst)
            print(f"Src after collective from PE {nvshmem.core.my_pe()}:", arr_src)

            # Free Buffers
            nvshmem.core.free_tensor(arr_src)
            nvshmem.core.free_tensor(arr_dst)
            print(f"Done testing torch tensor Reducescatter with dtype {dtype} op {op}")
    print("Done testing torch tensor Reducescatter")

def test_broadcast_torch():
    print("Testing torch tensor broadcast")
    if not _torch_enabled:
        print("WARNING: torch not found. Not running torch collective test")
        return
    # Get a buffer
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    print("Rank:", local_rank_per_node)
    dev = Device(local_rank_per_node)
    stream = dev.create_stream()

    for dtype in all_types_torch:
        for op in all_ops:
            arr_src = nvshmem.core.tensor((2,2), dtype=dtype)
            arr_dst = nvshmem.core.tensor((2,2), dtype=dtype)
            # Zero out the dest and set src to my_pe
            arr_dst[:] = 0
            # as usual, local_rank_per_node + 1 so there's no zeroes
            arr_src[:] = local_rank_per_node + 1

            # Print dst, src before
            print(f"Dest before collective from PE {nvshmem.core.my_pe()}:", arr_dst)
            print(f"Src before collective from PE {nvshmem.core.my_pe()}:", arr_src)

            nvshmem.core.broadcast(nvshmem.core.Teams.TEAM_WORLD, arr_dst, arr_src, root=0, stream=stream)
            # Explicit sync because we didn't set a torch stream
            # Important! Make sure you sync the stream that you launched the nvshmem collective on
            # (or the whole device like we do in this example)
            # before you use the results of the collective
            dev.sync()

            # Print dst, src after
            print(f"Dest after collective from PE {nvshmem.core.my_pe()}:", arr_dst)
            print(f"Src after collective from PE {nvshmem.core.my_pe()}:", arr_src)

            # Free Buffers
            nvshmem.core.free_tensor(arr_src)
            nvshmem.core.free_tensor(arr_dst)
            print(f"Done testing torch tensor broadcast with dtype {dtype} op {op}")
    print("Done testing torch tensor broadcast")

###
# Collectives on raw buffers
###

def test_collectives_on_raw_buffers():
    print("Testing collectives on raw buffers")
    # Get a buffer
    local_rank_per_node = nvshmem.core.team_my_pe(nvshmem.core.Teams.TEAM_NODE)
    print("Rank:", local_rank_per_node)
    dev = Device(local_rank_per_node)
    stream = dev.create_stream()
    for coll in ["reduce", "reducescatter", "fcollect", "alltoall", "broadcast"]:
        for dtype in all_types_nvshmem: 
            buf_src = nvshmem.core.buffer(nvshmem.core.utils.get_size((2,2,nvshmem.core.n_pes()), dtype))
            buf_dst = nvshmem.core.buffer(nvshmem.core.n_pes() * nvshmem.core.utils.get_size((2,2,nvshmem.core.n_pes()), dtype))

            # Print dst, src before
            print(f"Dest before collective from PE {nvshmem.core.my_pe()}:", buf_dst)
            print(f"Src before collective from PE {nvshmem.core.my_pe()}:", buf_src)
            if coll in ["reduce", "reducescatter"]:
                for op in all_ops:
                    nvshmem.core.collective_on_buffer(coll, nvshmem.core.Teams.TEAM_WORLD, buf_dst, buf_src, dtype=dtype, op=op, stream=stream)
            else:
                nvshmem.core.collective_on_buffer(coll, nvshmem.core.Teams.TEAM_WORLD, buf_dst, buf_src, dtype=dtype, op=None, stream=stream)
            # Important! Make sure you sync the stream that you launched the nvshmem collective on
            # (or the whole device like we do in this example)
            # before you use the results of the collective
            stream.sync()

            # Print dst, src after
            print(f"Dest after collective from PE {nvshmem.core.my_pe()}:", buf_dst)
            print(f"Src after collective from PE {nvshmem.core.my_pe()}:", buf_src)

            # Free Buffers
            nvshmem.core.free(buf_src)
            nvshmem.core.free(buf_dst)
            print(f"Done testing {coll} with dtype {dtype} on raw buffer")
    print("Done testing on raw buffers")


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--init-type", "-i", type=str, help="Init type to use", choices=["mpi", "uid"], default="uid")
    args = parser.parse_args()
    if args.init_type == "uid":
        uid_init()
    elif args.init_type == "mpi":
        mpi_init()

    # Cupy Interop
    test_fcollect_cupy()
    test_reduce_cupy()
    test_alltoall_cupy()
    test_reducescatter_cupy()
    test_broadcast_cupy()

    # Torch interop
    test_alltoall_torch()
    test_fcollect_torch()
    test_reduce_torch()
    test_reducescatter_torch()
    test_broadcast_torch()

    # Raw Buffers
    test_collectives_on_raw_buffers()

    # No arrays
    test_barrier()
    test_team_sync()

    nvshmem.core.finalize()
