import pytest

from ray.tune.search import BasicVariantGenerator, ConcurrencyLimiter, Searcher
from ray.tune.search.repeater import Repeater
from ray.tune.search.search_generator import SearchGenerator
from ray.tune.utils.mock_trainable import MOCK_TRAINABLE_NAME, register_mock_trainable


@pytest.fixture(autouse=True)
def register_test_trainable():
    register_mock_trainable()


def test_nested_suggestion():
    class TestSuggestion(Searcher):
        def suggest(self, trial_id):
            return {"a": {"b": {"c": {"d": 4, "e": 5}}}}

    searcher = TestSuggestion()
    alg = SearchGenerator(searcher)
    alg.add_configurations({"test": {"run": MOCK_TRAINABLE_NAME}})
    trial = alg.next_trial()
    assert "e=5" in trial.experiment_tag
    assert "d=4" in trial.experiment_tag


def _repeat_trials(num_samples: int, repeat: int):
    class TestSuggestion(Searcher):
        index = 0

        def suggest(self, trial_id):
            self.index += 1
            return {"test_variable": 5 + self.index}

        def on_trial_complete(self, *args, **kwargs):
            return

    searcher = TestSuggestion(metric="episode_reward_mean")
    repeat_searcher = Repeater(searcher, repeat=repeat, set_index=False)
    alg = SearchGenerator(repeat_searcher)

    alg.add_configurations(
        {
            "test": {
                "run": MOCK_TRAINABLE_NAME,
                "num_samples": num_samples,
                "stop": {"training_iteration": 1},
            }
        }
    )

    trials = []
    while not alg.is_finished():
        trials.append(alg.next_trial())

    return trials


def test_repeat_1():
    trials = _repeat_trials(num_samples=2, repeat=1)
    assert len(trials) == 2
    parameter_set = {t.evaluated_params["test_variable"] for t in trials}
    assert len(parameter_set) == 2


def test_repeat_4():
    trials = _repeat_trials(num_samples=12, repeat=4)
    assert len(trials) == 12
    parameter_set = {t.evaluated_params["test_variable"] for t in trials}
    assert len(parameter_set) == 3


def test_odd_repeat():
    trials = _repeat_trials(num_samples=11, repeat=5)
    assert len(trials) == 11
    parameter_set = {t.evaluated_params["test_variable"] for t in trials}
    assert len(parameter_set) == 3


def test_set_get_repeater():
    class TestSuggestion(Searcher):
        def __init__(self, index):
            self.index = index
            self.returned_result = []
            super().__init__(metric="result", mode="max")

        def suggest(self, trial_id):
            self.index += 1
            return {"score": self.index}

        def on_trial_complete(self, trial_id, result=None, **kwargs):
            self.returned_result.append(result)

    searcher = TestSuggestion(0)
    repeater1 = Repeater(searcher, repeat=3, set_index=False)
    for i in range(3):
        assert repeater1.suggest(f"test_{i}")["score"] == 1
    for i in range(2):  # An incomplete set of results
        assert repeater1.suggest(f"test_{i}_2")["score"] == 2

    # Restore a new one
    state = repeater1.get_state()
    del repeater1
    new_repeater = Repeater(searcher, repeat=1, set_index=True)
    new_repeater.set_state(state)
    assert new_repeater.repeat == 3
    assert new_repeater.suggest("test_2_2")["score"] == 2
    assert new_repeater.suggest("test_x")["score"] == 3

    # Report results
    for i in range(3):
        new_repeater.on_trial_complete(f"test_{i}", {"result": 2})

    for i in range(3):
        new_repeater.on_trial_complete(f"test_{i}_2", {"result": -i * 10})

    assert len(new_repeater.searcher.returned_result) == 2
    assert new_repeater.searcher.returned_result[-1] == {"result": -10}

    # Finish the rest of the last trial group
    new_repeater.on_trial_complete("test_x", {"result": 3})
    assert new_repeater.suggest("test_y")["score"] == 3
    new_repeater.on_trial_complete("test_y", {"result": 3})
    assert len(new_repeater.searcher.returned_result) == 2
    assert new_repeater.suggest("test_z")["score"] == 3
    new_repeater.on_trial_complete("test_z", {"result": 3})
    assert len(new_repeater.searcher.returned_result) == 3
    assert new_repeater.searcher.returned_result[-1] == {"result": 3}


def test_set_get_limiter():
    class TestSuggestion(Searcher):
        def __init__(self, index):
            self.index = index
            self.returned_result = []
            super().__init__(metric="result", mode="max")

        def suggest(self, trial_id):
            self.index += 1
            return {"score": self.index}

        def on_trial_complete(self, trial_id, result=None, **kwargs):
            self.returned_result.append(result)

    searcher = TestSuggestion(0)
    limiter = ConcurrencyLimiter(searcher, max_concurrent=2)
    assert limiter.suggest("test_1")["score"] == 1
    assert limiter.suggest("test_2")["score"] == 2
    assert limiter.suggest("test_3") is None

    state = limiter.get_state()
    del limiter
    limiter2 = ConcurrencyLimiter(searcher, max_concurrent=3)
    limiter2.set_state(state)
    assert limiter2.suggest("test_4") is None
    assert limiter2.suggest("test_5") is None
    limiter2.on_trial_complete("test_1", {"result": 3})
    limiter2.on_trial_complete("test_2", {"result": 3})
    assert limiter2.suggest("test_3")["score"] == 3


def test_basic_variant_limiter():
    search_alg = BasicVariantGenerator(max_concurrent=2)

    experiment_spec = {
        "run": MOCK_TRAINABLE_NAME,
        "num_samples": 5,
        "stop": {"training_iteration": 1},
    }
    search_alg.add_configurations({"test": experiment_spec})

    trial1 = search_alg.next_trial()
    assert trial1

    trial2 = search_alg.next_trial()
    assert trial2

    # Returns None because of limiting
    trial3 = search_alg.next_trial()
    assert not trial3

    # Finish trial, now trial 3 should be created
    search_alg.on_trial_complete(trial1.trial_id, None, False)
    trial3 = search_alg.next_trial()
    assert trial3

    trial4 = search_alg.next_trial()
    assert not trial4

    search_alg.on_trial_complete(trial2.trial_id, None, False)
    search_alg.on_trial_complete(trial3.trial_id, None, False)

    trial4 = search_alg.next_trial()
    assert trial4

    trial5 = search_alg.next_trial()
    assert trial5

    search_alg.on_trial_complete(trial4.trial_id, None, False)

    # Should also be None because search is finished
    trial6 = search_alg.next_trial()
    assert not trial6


def test_batch_limiter():
    class TestSuggestion(Searcher):
        def __init__(self, index):
            self.index = index
            self.returned_result = []
            super().__init__(metric="result", mode="max")

        def suggest(self, trial_id):
            self.index += 1
            return {"score": self.index}

        def on_trial_complete(self, trial_id, result=None, **kwargs):
            self.returned_result.append(result)

    searcher = TestSuggestion(0)
    limiter = ConcurrencyLimiter(searcher, max_concurrent=2, batch=True)
    assert limiter.suggest("test_1")["score"] == 1
    assert limiter.suggest("test_2")["score"] == 2
    assert limiter.suggest("test_3") is None

    limiter.on_trial_complete("test_1", {"result": 3})
    assert limiter.suggest("test_3") is None
    limiter.on_trial_complete("test_2", {"result": 3})
    assert limiter.suggest("test_3") is not None


def test_batch_limiter_infinite_loop():
    """Check whether an infinite loop when less than max_concurrent trials
    are suggested with batch mode is avoided.
    """

    class TestSuggestion(Searcher):
        def __init__(self, index, max_suggestions=10):
            self.index = index
            self.max_suggestions = max_suggestions
            self.returned_result = []
            super().__init__(metric="result", mode="max")

        def suggest(self, trial_id):
            self.index += 1
            if self.index > self.max_suggestions:
                return None
            return {"score": self.index}

        def on_trial_complete(self, trial_id, result=None, **kwargs):
            self.returned_result.append(result)
            self.index = 0

    searcher = TestSuggestion(0, 2)
    limiter = ConcurrencyLimiter(searcher, max_concurrent=5, batch=True)
    limiter.suggest("test_1")
    limiter.suggest("test_2")
    limiter.suggest("test_3")  # TestSuggestion return None

    limiter.on_trial_complete("test_1", {"result": 3})
    limiter.on_trial_complete("test_2", {"result": 3})
    assert limiter.searcher.returned_result

    searcher = TestSuggestion(0, 10)
    limiter = ConcurrencyLimiter(searcher, max_concurrent=5, batch=True)
    limiter.suggest("test_1")
    limiter.suggest("test_2")
    limiter.suggest("test_3")

    limiter.on_trial_complete("test_1", {"result": 3})
    limiter.on_trial_complete("test_2", {"result": 3})
    assert not limiter.searcher.returned_result


def test_set_max_concurrency():
    """Test whether ``set_max_concurrency`` is called by the
    ``ConcurrencyLimiter`` and works correctly.
    """

    class TestSuggestion(Searcher):
        def __init__(self, index):
            self.index = index
            self.returned_result = []
            self._max_concurrent = 1
            super().__init__(metric="result", mode="max")

        def suggest(self, trial_id):
            self.index += 1
            return {"score": self.index}

        def on_trial_complete(self, trial_id, result=None, **kwargs):
            self.returned_result.append(result)

        def set_max_concurrency(self, max_concurrent: int) -> bool:
            self._max_concurrent = max_concurrent
            return True

    searcher = TestSuggestion(0)
    limiter_max_concurrent = 2
    limiter = ConcurrencyLimiter(
        searcher, max_concurrent=limiter_max_concurrent, batch=True
    )
    assert limiter.searcher._max_concurrent == limiter_max_concurrent
    # Since set_max_concurrency returns True, ConcurrencyLimiter should not
    # be limiting concurrency itself
    assert not limiter._limit_concurrency
    assert limiter.suggest("test_1")["score"] == 1
    assert limiter.suggest("test_2")["score"] == 2
    assert limiter.suggest("test_3")["score"] == 3


if __name__ == "__main__":
    import sys

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