import os
import time
from collections import Counter

import pytest
import sys

import ray
from ray import tune
from ray.air.execution import FixedResourceManager, PlacementGroupResourceManager
from ray.tune import PlacementGroupFactory, TuneError
from ray.tune.execution.tune_controller import TuneController
from ray.tune.experiment import Trial
from ray.tune.schedulers import FIFOScheduler, TrialScheduler
from ray.tune.search import BasicVariantGenerator
from ray.tune.utils.mock import TrialStatusSnapshot, TrialStatusSnapshotTaker


@pytest.fixture(scope="function")
def ray_start_4_cpus_2_gpus_extra():
    address_info = ray.init(num_cpus=4, num_gpus=2, resources={"a": 2})
    yield address_info
    ray.shutdown()


@pytest.mark.parametrize(
    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
)
@pytest.mark.parametrize(
    "bundles",
    [
        [{"CPU": 1}, {"CPU": 3, "GPU": 1}],
        [{"CPU": 1, "a": 2}],
        [{"CPU": 1}, {"a": 2}],
        [{"CPU": 1, "GPU": 1}],
    ],
)
def test_resource_parallelism_single(
    ray_start_4_cpus_2_gpus_extra, resource_manager_cls, bundles
):
    """Test that extra and custom resources are respected for parallelism.

    We schedule two trials with resources according to the bundle. If only
    the head bundle or only CPU/GPU resources were considered, both trials
    could run in parallel.

    However, we assert that the resources in child bundles and extra resources
    are respected and only one trial runs in parallel.

    Legacy test: test_trial_runner.py::TrialRunnerTest::testExtraResources
    Legacy test: test_trial_runner.py::TrialRunnerTest::testCustomResources
    Legacy test: test_trial_runner.py::TrialRunnerTest::testExtraCustomResources
    Legacy test: test_trial_runner.py::TrialRunnerTest::testResourceScheduler
    """
    snapshot = TrialStatusSnapshot()
    runner = TuneController(
        resource_manager_factory=lambda: resource_manager_cls(),
        callbacks=[TrialStatusSnapshotTaker(snapshot)],
    )
    kwargs = {
        "stopping_criterion": {"training_iteration": 1},
        "placement_group_factory": PlacementGroupFactory(bundles),
    }
    trials = [Trial("__fake", **kwargs), Trial("__fake", **kwargs)]
    for t in trials:
        runner.add_trial(t)

    while not runner.is_finished():
        runner.step()

    assert snapshot.max_running_trials() == 1
    assert snapshot.all_trials_are_terminated()


@pytest.mark.parametrize(
    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
)
def test_fractional_gpus(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
    """Test that fractional GPUs lead to more parallelism.

    We schedule four trials with 0.75 GPUs each. Since our cluster has 2 GPUs,
    we should be able to run 2 trials in parallel.

    Legacy test: test_trial_runner.py::TrialRunnerTest::testFractionalGpus
    """
    snapshot = TrialStatusSnapshot()
    runner = TuneController(
        resource_manager_factory=lambda: resource_manager_cls(),
        callbacks=[TrialStatusSnapshotTaker(snapshot)],
    )
    kwargs = {
        "stopping_criterion": {"training_iteration": 1},
        "placement_group_factory": PlacementGroupFactory([{"GPU": 0.75}]),
        "config": {
            "sleep": 1,
        },
    }
    trials = [Trial("__fake", **kwargs) for i in range(4)]
    for t in trials:
        runner.add_trial(t)

    while not runner.is_finished():
        runner.step()

    assert snapshot.max_running_trials() == 2
    assert snapshot.all_trials_are_terminated()


@pytest.mark.parametrize(
    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
)
def test_multi_step(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
    """Test that trials can run for more than one iteration.

    Todo (krfricke): This is not a resource test, so it should be moved.

    Legacy test: test_trial_runner.py::TrialRunnerTest::testMultiStepRun
    Legacy test: test_trial_runner.py::TrialRunnerTest::testMultiStepRun2
    """
    snapshot = TrialStatusSnapshot()
    runner = TuneController(
        resource_manager_factory=lambda: resource_manager_cls(),
        callbacks=[TrialStatusSnapshotTaker(snapshot)],
    )
    kwargs = {
        "stopping_criterion": {"training_iteration": 5},
        "placement_group_factory": PlacementGroupFactory([{"CPU": 1, "GPU": 1}]),
    }
    trials = [Trial("__fake", **kwargs) for i in range(2)]
    for t in trials:
        runner.add_trial(t)

    while not runner.is_finished():
        runner.step()

    # Overstepping should throw error
    # test_trial_runner.py::TrialRunnerTest::testMultiStepRun2
    with pytest.raises(TuneError):
        runner.step()

    assert snapshot.all_trials_are_terminated()
    assert all(t.last_result["training_iteration"] == 5 for t in runner.get_trials())


@pytest.mark.parametrize(
    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
)
def test_resources_changing(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
    """Checks that resource requirements can be changed on fly.

    Legacy test: test_trial_runner.py::TrialRunnerTest::testChangeResources
    """

    class ChangingScheduler(FIFOScheduler):
        def __init__(self):
            self._has_received_one_trial_result = False

        # For figuring out how many runner.step there are.
        def has_received_one_trial_result(self):
            return self._has_received_one_trial_result

        def on_trial_result(self, trial_runner, trial, result):
            if result["training_iteration"] == 1:
                self._has_received_one_trial_result = True
                executor = trial_runner.trial_executor
                executor.pause_trial(trial)
                trial.update_resources(dict(cpu=4, gpu=0))
            return TrialScheduler.NOOP

    scheduler = ChangingScheduler()
    runner = TuneController(
        resource_manager_factory=lambda: resource_manager_cls(), scheduler=scheduler
    )
    kwargs = {
        "stopping_criterion": {"training_iteration": 2},
        "placement_group_factory": PlacementGroupFactory([{"CPU": 2, "GPU": 0}]),
    }
    trials = [Trial("__fake", **kwargs)]
    for t in trials:
        runner.add_trial(t)

    while not trials[0].status == Trial.RUNNING:
        runner.step()

    assert trials[0].status == Trial.RUNNING
    assert runner._actor_manager.get_live_actors_resources().get("CPU") == 2

    with pytest.raises(ValueError):
        trials[0].update_resources(dict(cpu=4, gpu=0))

    while not scheduler.has_received_one_trial_result():
        runner.step()

    assert trials[0].status == Trial.PAUSED

    while not trials[0].status == Trial.RUNNING:
        runner.step()

    assert runner._actor_manager.get_live_actors_resources().get("CPU") == 4

    runner.step()


@pytest.mark.parametrize(
    "resource_manager_cls", [FixedResourceManager, PlacementGroupResourceManager]
)
def test_queue_filling(ray_start_4_cpus_2_gpus_extra, resource_manager_cls):
    """Checks that the trial queue is filled even if only 1 pending trial is allowed.

    Legacy test: test_trial_runner.py::TrialRunnerTest::testQueueFilling
    """
    os.environ["TUNE_MAX_PENDING_TRIALS_PG"] = "1"

    def f1(config):
        for i in range(10):
            yield i
            time.sleep(1)

    tune.register_trainable("f1", f1)

    search_alg = BasicVariantGenerator()
    search_alg.add_configurations(
        {
            "foo": {
                "run": "f1",
                "num_samples": 100,
                "config": {
                    "a": tune.sample_from(lambda spec: 5.0 / 7),
                    "b": tune.sample_from(lambda spec: "long" * 40),
                },
                "resources_per_trial": {"cpu": 2},
            }
        }
    )

    runner = TuneController(
        resource_manager_factory=lambda: resource_manager_cls(), search_alg=search_alg
    )

    while len(runner.get_trials()) < 3:
        runner.step()

    # All trials are enqueued
    assert len(runner.get_trials()) == 3

    status_count = Counter(t.status for t in runner.get_trials())
    while status_count.get(Trial.RUNNING, 0) < 2 and not runner.is_finished():
        runner.step()
        status_count = Counter(t.status for t in runner.get_trials())

    assert len(runner.get_trials()) == 3

    status_count = Counter(t.status for t in runner.get_trials())
    assert status_count.get(Trial.RUNNING) == 2
    assert status_count.get(Trial.PENDING) == 1


if __name__ == "__main__":
    sys.exit(pytest.main(["-v", "--reruns", "3", __file__]))
