"""
This file is for init/fini tests for NVSHMEM4Py
The reason it's separate from the basic sanity test is it needs to be launched in multiples.
Test both the UID and MPI based init/fini
"""

# Usage: mpirun -np <N> python init_test.py -i <init type>
# Must set `-x NVSHMEM_BOOTSTRAP=MPI` if you want to use the MPI init

import nvshmem
import nvshmem.core

import argparse
import sys
import os

import numpy as np
from mpi4py import MPI
from cuda.core.experimental import Device, system

# User should not import this - it's here so we can print stuff
import nvshmem.bindings

def test_mpi_comm_init():
    # Test device init and bootstrap
    local_rank_per_node = MPI.COMM_WORLD.Get_rank() % system.num_devices
    dev = Device(local_rank_per_node)
    dev.set_current()
    nvshmem.core.init(device=dev, uid=None, rank=None, nranks=None,
                      mpi_comm=MPI.COMM_WORLD, initializer_method="mpi")
    nvshmem.core.finalize()
    print("Init/Fini with MPI passed with cuda.core init/fini as well")


def test_multi_init():
    # Test multiple calls to init with overlapping sessions
    local_rank_per_node = MPI.COMM_WORLD.Get_rank() % system.num_devices
    dev = Device(local_rank_per_node)
    dev.set_current()
    nvshmem.core.init(device=dev, uid=None, rank=None, nranks=None,
                      mpi_comm=MPI.COMM_WORLD, initializer_method="mpi")
    print("called init1")
    nvshmem.core.init(device=dev, uid=None, rank=None, nranks=None,
                      mpi_comm=MPI.COMM_WORLD, initializer_method="mpi")
    print("called init2")
    print(f"Hello from PE {nvshmem.bindings.my_pe()} npes {nvshmem.bindings.n_pes()}")
    nvshmem.core.finalize()
    print("called fini1")
    nvshmem.core.finalize()
    print("called fini2")
    print("Init/Fini multi-init with MPI passed")

def test_uid_init():
    # This will use mpi4py to perform a UID based init with bcast.
    comm = MPI.COMM_WORLD
    rank = comm.Get_rank()
    nranks = comm.Get_size()
    local_rank_per_node = rank % system.num_devices
    dev = Device(local_rank_per_node)
    dev.set_current()

    # Create an empty uniqueid for all ranks
    uniqueid = nvshmem.core.get_unique_id(empty=True)
    if rank == 0:
        # Rank 0 gets a real uniqueid
        uniqueid = nvshmem.core.get_unique_id()

    # Broadcast UID to all ranks
    comm.Bcast(uniqueid._data.view(np.int8), root=0)

    nvshmem.core.init(device=dev, uid=uniqueid, rank=rank, nranks=nranks,
                      mpi_comm=None, initializer_method="uid")
    nvshmem.core.finalize()
    print("Init/Fini with UID passed")

def test_emulated_mpi_init():
    # Test device init and bootstrap
    local_rank_per_node = MPI.COMM_WORLD.Get_rank() % system.num_devices
    dev = Device(local_rank_per_node)
    dev.set_current()
    nvshmem.core.init(device=dev, uid=None, rank=None, nranks=None,
                      mpi_comm=MPI.COMM_WORLD, initializer_method="emulated_mpi")
    nvshmem.core.finalize()
    print("Init/Fini with emulated MPI passed with cuda.core init/fini as well")

def test_none_device_init():
    # Test init with None device and allocate with device set
    local_rank_per_node = MPI.COMM_WORLD.Get_rank() % system.num_devices
    
    nvshmem.core.init(device=None, uid=None, rank=None, nranks=None,
                      mpi_comm=MPI.COMM_WORLD, initializer_method="mpi")
    dev = Device(local_rank_per_node)
    dev.set_current()
    buf = nvshmem.core.buffer(1024)
    print(buf)
    nvshmem.core.free(buf)
    nvshmem.core.finalize()
    print("Init/Fini with MPI passed with cuda.core init/fini as well")
    print("2 stage init with None device passed")

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

    if args.init_type == "mpi":
        test_multi_init()
        test_mpi_comm_init()
        test_none_device_init()
    elif args.init_type == "uid":
        test_uid_init()
    elif args.init_type == "emulated_mpi":
        test_emulated_mpi_init()
    else:
        print(f"Unexpected init type: {args.init_type}")
        sys.exit(1)