"""Test the allgather API on a distributed Ray cluster."""
import pytest
import ray

import cupy as cp
import torch

from ray.util.collective.tests.util import (
    create_collective_multigpu_workers,
    init_tensors_for_gather_scatter_multigpu,
)


@pytest.mark.parametrize("tensor_backend", ["cupy", "torch"])
@pytest.mark.parametrize(
    "array_size", [2, 2**5, 2**10, 2**15, 2**20, [2, 2], [5, 5, 5]]
)
def test_allgather_different_array_size(
    ray_start_distributed_multigpu_2_nodes_4_gpus, array_size, tensor_backend
):
    world_size = 2
    num_gpu_per_worker = 2
    actual_world_size = world_size * num_gpu_per_worker
    actors, _ = create_collective_multigpu_workers(world_size)
    init_tensors_for_gather_scatter_multigpu(
        actors, array_size=array_size, tensor_backend=tensor_backend
    )
    results = ray.get([a.do_allgather_multigpu.remote() for a in actors])
    for i in range(world_size):
        for j in range(num_gpu_per_worker):
            for k in range(actual_world_size):
                if tensor_backend == "cupy":
                    assert (
                        results[i][j][k] == cp.ones(array_size, dtype=cp.float32)
                    ).all()
                else:
                    assert (
                        results[i][j][k]
                        == torch.ones(array_size, dtype=torch.float32).cuda(j)
                    ).all()


def test_allgather_torch_cupy(ray_start_distributed_multigpu_2_nodes_4_gpus):
    world_size = 2
    num_gpu_per_worker = 2
    actual_world_size = world_size * num_gpu_per_worker
    shape = [10, 10]
    actors, _ = create_collective_multigpu_workers(world_size)

    # tensor is pytorch, list is cupy
    for i, a in enumerate(actors):
        ray.get(
            [a.set_buffer.remote(shape, tensor_type0="torch", tensor_type1="torch")]
        )
        ray.get(
            [a.set_list_buffer.remote(shape, tensor_type0="cupy", tensor_type1="cupy")]
        )
    results = ray.get([a.do_allgather_multigpu.remote() for a in actors])
    for i in range(world_size):
        for j in range(num_gpu_per_worker):
            for k in range(actual_world_size):
                assert (results[i][j][k] == cp.ones(shape, dtype=cp.float32)).all()

    # tensor is cupy, list is pytorch
    for i, a in enumerate(actors):
        ray.get([a.set_buffer.remote(shape, tensor_type0="cupy", tensor_type1="cupy")])
        ray.get(
            [
                a.set_list_buffer.remote(
                    shape, tensor_type0="torch", tensor_type1="torch"
                )
            ]
        )
    results = ray.get([a.do_allgather_multigpu.remote() for a in actors])
    for i in range(world_size):
        for j in range(num_gpu_per_worker):
            for k in range(actual_world_size):
                assert (
                    results[i][j][k] == torch.ones(shape, dtype=torch.float32).cuda(j)
                ).all()


if __name__ == "__main__":
    import pytest
    import sys

    sys.exit(pytest.main(["-v", "-x", __file__]))
