import pytest
import torch

import ray
from ray.train import RunConfig, ScalingConfig
from ray.train.constants import TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S
from ray.train.examples.pytorch.torch_linear_example import (
    train_func as linear_train_func,
)
from ray.train.torch import TorchConfig, TorchTrainer
from ray.train.torch.config import _is_backend_nccl
from ray.train.v2._internal.constants import HEALTH_CHECK_INTERVAL_S_ENV_VAR
from ray.train.v2.api.config import FailureConfig
from ray.train.v2.api.exceptions import WorkerGroupError
from ray.train.v2.torch.torchft_config import TorchftConfig


@pytest.fixture(scope="module")
def ray_start_4_cpus():
    ray.init(num_cpus=4)
    yield
    ray.shutdown()


@pytest.fixture(autouse=True)
def reduce_health_check_interval(monkeypatch):
    monkeypatch.setenv(HEALTH_CHECK_INTERVAL_S_ENV_VAR, "0.2")
    yield


def test_minimal(ray_start_4_cpus):
    def train_func():
        pass

    trainer = TorchTrainer(train_func)
    trainer.fit()


@pytest.mark.parametrize("num_workers", [1, 2])
def test_torch_linear(ray_start_4_cpus, num_workers):
    def train_func(config):
        result = linear_train_func(config)
        assert len(result) == epochs
        assert result[-1]["loss"] < result[0]["loss"]

    num_workers = num_workers
    epochs = 3
    scaling_config = ScalingConfig(num_workers=num_workers)
    config = {"lr": 1e-2, "hidden_size": 1, "batch_size": 4, "epochs": epochs}
    trainer = TorchTrainer(
        train_loop_per_worker=train_func,
        train_loop_config=config,
        scaling_config=scaling_config,
    )
    trainer.fit()


@pytest.mark.parametrize(
    "torch_config,expected_world_size",
    [
        (TorchConfig(backend="gloo", init_method="env"), 2),
        (TorchConfig(backend="gloo", init_method="tcp"), 2),
        (
            TorchftConfig(
                backend="gloo", init_method="env", lighthouse_kwargs={"min_replicas": 1}
            ),
            1,
        ),
        (
            TorchftConfig(
                backend="gloo", init_method="tcp", lighthouse_kwargs={"min_replicas": 1}
            ),
            1,
        ),
    ],
)
def test_torch_start_shutdown(ray_start_4_cpus, torch_config, expected_world_size):
    def check_process_group():
        assert (
            torch.distributed.is_initialized()
            and torch.distributed.get_world_size() == expected_world_size
        )

    trainer = TorchTrainer(
        train_loop_per_worker=check_process_group,
        scaling_config=ScalingConfig(num_workers=2),
        torch_config=torch_config,
    )
    trainer.fit()


@pytest.mark.parametrize("timeout_s", [5, 0])
def test_torch_process_group_shutdown_timeout(ray_start_4_cpus, monkeypatch, timeout_s):
    """Tests that we don't more than a predefined timeout
    on Torch process group shutdown."""

    monkeypatch.setenv(TORCH_PROCESS_GROUP_SHUTDOWN_TIMEOUT_S, str(timeout_s))

    trainer = TorchTrainer(
        train_loop_per_worker=lambda: None,
        scaling_config=ScalingConfig(num_workers=2),
        torch_config=TorchConfig(backend="gloo"),
    )
    # Even if shutdown times out (timeout_s=0),
    # the training should complete successfully.
    trainer.fit()


def test_torchft_linear(ray_start_4_cpus):
    """Test torchft linear training: loss goes down and models are equal across workers."""

    from ray.train.v2.examples.pytorch.torchft_linear_example import (
        train_func as torchft_linear_train_func,
    )

    @ray.remote
    class WeightCollector:
        def __init__(self):
            self.weights = {}

        def report(self, rank, weight, bias):
            self.weights[rank] = {"weight": weight, "bias": bias}

        def get_weights(self):
            return self.weights

    collector = WeightCollector.remote()

    def train_func(config):
        result = torchft_linear_train_func(config)
        assert result[-1]["loss"] < result[0]["loss"]
        world_rank = ray.train.get_context().get_world_rank()
        ray.get(
            config["collector"].report.remote(
                world_rank, result[-1]["weight"], result[-1]["bias"]
            )
        )

    trainer = TorchTrainer(
        train_loop_per_worker=train_func,
        train_loop_config={"collector": collector},
        scaling_config=ScalingConfig(num_workers=2),
        torch_config=TorchftConfig(
            backend="gloo", lighthouse_kwargs={"min_replicas": 2}
        ),
    )
    result = trainer.fit()
    assert result.error is None

    # Check that models converged across workers.
    weights = ray.get(collector.get_weights.remote())
    assert len(weights) == 2
    assert weights[0]["weight"] == pytest.approx(weights[1]["weight"], abs=1e-4)
    assert weights[0]["bias"] == pytest.approx(weights[1]["bias"], abs=1e-4)


# TODO(tseah): Add test for lighthouse failures (e.g. lighthouse unreachable).


@pytest.mark.parametrize(
    "min_replicas,max_failures,expect_error,expected_train_fn_calls,expected_reports",
    [
        # TODO(tseah): enable this after we support training with 1/2 workers without
        # trying to restart the replica group.
        # (1, 0, False, 2),
        # This continues training with 1 replica. It does not replay step 10 and its report.
        # TODO(tseah): expected_reports should be 1 when we support training with 1/2 workers.
        (1, 1, False, 3, 0),
        # This errors immediately.
        (2, 0, True, 2, 0),
        # This stops training with 1 replica. It replays step 10 and its report.
        (2, 1, False, 3, 2),
    ],
)
def test_torchft_linear_replica_failure(
    ray_start_4_cpus,
    min_replicas,
    max_failures,
    expect_error,
    expected_train_fn_calls,
    expected_reports,
):
    """Test torchft linear training behavior when a replica fails mid-training."""

    from ray.train.v2.examples.pytorch.torchft_linear_example import (
        train_func as torchft_linear_train_func,
    )

    num_workers = 2
    # TODO(tseah): remove this check once we support training with 1/2 workers.
    training_requires_all_workers = min_replicas == num_workers

    @ray.remote
    class Counter:
        def __init__(self):
            self.count = 0

        def increment(self):
            self.count += 1
            return self.count

        def get_count(self):
            return self.count

    counter = Counter.remote()

    def train_fn(config):
        counter.increment.remote()
        return torchft_linear_train_func(config)

    trainer = TorchTrainer(
        train_loop_per_worker=train_fn,
        train_loop_config={
            "num_steps": 20,
            "error_step": 10,
            "error_rank": 0,
            "num_replicas": min_replicas,
            "training_requires_all_workers": training_requires_all_workers,
        },
        scaling_config=ScalingConfig(num_workers=num_workers),
        torch_config=TorchftConfig(
            backend="gloo", lighthouse_kwargs={"min_replicas": min_replicas}
        ),
        run_config=RunConfig(failure_config=FailureConfig(max_failures=max_failures)),
    )
    if expect_error:
        with pytest.raises(WorkerGroupError):
            trainer.fit()
    else:
        result = trainer.fit()
        assert len(result.best_checkpoints) == expected_reports
        assert result.error is None
    # Fewer train_fn calls indicate partial worker group restarts.
    assert ray.get(counter.get_count.remote()) == expected_train_fn_calls


def test_is_backend_nccl():
    assert _is_backend_nccl("nccl")
    assert _is_backend_nccl("cuda:nccl")
    assert _is_backend_nccl("cpu:gloo,cuda:nccl")
    assert not _is_backend_nccl("gloo")
    assert not _is_backend_nccl("cpu:nccl")
    assert not _is_backend_nccl("cuda:gloo")
    assert not _is_backend_nccl("cpu:gloo,cuda:gloo")


if __name__ == "__main__":
    import sys

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