text: # -*- coding: utf-8 -*-
  """Example usage and tests for :mod:`orion.algo.random`."""
  
  import pytest
  from base import (
      ExploitStub,
      ExploreStub,
      ObjectiveStub,
      TrialStub,
      no_shutil_copytree,
      sample_trials,
      space,
  )
  
  from orion.algo.pbt.pbt import PBT, compute_fidelities
  from orion.core.worker.primary_algo import SpaceTransformAlgoWrapper
  from orion.core.worker.trial import Trial
  from orion.testing.algo import BaseAlgoTests
  
  
  class TestComputeFidelities:
      def test_base_1(self):
          assert compute_fidelities(10, 10, 20, 1) == list(map(float, range(10, 21)))
  
      def test_other_bases(self):
          assert compute_fidelities(9, 2, 2 ** 10, 2) == [2 ** i for i in range(1, 11)]
  
  
  class TestPBTObserve:
      def test_triage_unknown_trial(self, space):
          pbt = SpaceTransformAlgoWrapper(PBT, space).algorithm
          trial = pbt.space.sample(1, seed=1)[0]
          trials_to_verify = pbt._triage([trial])
  
          assert trials_to_verify == []
          assert len(pbt.lineages) == 0
  
      @pytest.mark.parametrize("status", ["new", "reserved", "interrupted"])
      def test_triage_root_not_ready(self, status, space):
          pbt = SpaceTransformAlgoWrapper(PBT, space).algorithm
  
          trial = sample_trials(pbt.space, num=1, status=status)[0]
  
          pbt.register(trial)
  
          trials_to_verify = pbt._triage([trial])
  
          assert trials_to_verify == []
          assert pbt.has_suggested(trial)
          assert not pbt.has_observed(trial)
          assert len(pbt.lineages) == 1
  
      @pytest.mark.parametrize("status", ["broken", "completed"])
      def test_triage_root_ready(self, status, space):
          pbt = SpaceTransformAlgoWrapper(PBT, space).algorithm
  
          trial = sample_trials(pbt.space, num=1, status="new")[0]
  
          pbt.register(trial)
  
          trial.status = status
          trial._results.append(Trial.Result(name="objective", type="objective", value=1))
  
          trials_to_verify = pbt._triage([trial])
  
          assert trials_to_verify == [trial]
  
          assert pbt.has_suggested(trial)
          assert pbt.has_observed(trial)
          assert len(pbt.lineages) == 1
  
      @pytest.mark.parametrize("status", ["broken", "completed"])
      def test_triage_root_observed(self, status, space):
          pbt = SpaceTransformAlgoWrapper(PBT, space).algorithm
  
          trial = sample_trials(pbt.space, num=1, status="completed", objective=1)[0]
  
          pbt.register(trial)
  
          trials_to_verify = pbt._triage([trial])
  
          assert trials_to_verify == []
  
          assert pbt.has_suggested(trial)
          assert pbt.has_observed(trial)
          assert len(pbt.lineages) == 1
  
      @pytest.mark.usefixtures("no_shutil_copytree")
      def test_dont_queue_broken_root_for_promotions(self, space):
          pbt = SpaceTransformAlgoWrapper(PBT, space).algorithm
  
          trial = sample_trials(pbt.space, num=1, status="broken")[0]
          pbt.register(trial)
  
          # Should not queue anything
          pbt._queue_trials_for_promotions([trial])
          assert len(pbt._queue) == 0
  
      @pytest.mark.usefixtures("no_shutil_copytree")
      def test_queue_broken_trials_for_promotions(self, space):
          pbt = SpaceTransformAlgoWrapper(PBT, space).algorithm
          trial = sample_trials(pbt.space, num=1, status="completed", objective=1)[0]
          pbt.register(trial)
  
          new_trial = trial.branch(params={"f": pbt.fidelities[trial.params["f"]]})
          pbt.lineages.fork(trial, new_trial)
  
          new_trial.status = "broken"
          pbt.register(new_trial)
  
          # Should queue the parent of the broken trial
          pbt._queue_trials_for_promotions([new_trial])
          assert len(pbt._queue) == 1
          assert pbt._queue[0].id == trial.id
  
      @pytest.mark.usefixtures("no_shutil_copytree")
      def test_queue_broken_trials_from_jump_for_promotions(self, space):
          pbt = SpaceTransformAlgoWrapper(PBT, space).algorithm
  
          parent_trial = sample_trials(pbt.space, num=1, status="completed", objective=1)[
              0
          ]
          base_trial = sample_trials(
              pbt.space, num=1, seed=2, status="completed", objective=1
          )[0]
  
          pbt.register(parent_trial)
          pbt.register(base_trial)
  
          new_trial = parent_trial.branch(
              params={"f": pbt.fidelities[parent_trial.params["f"]]}
          )
          pbt.lineages.fork(parent_trial, new_trial)
          pbt.lineages.set_jump(base_trial, new_trial)
  
          new_trial.status = "broken"
          pbt.register(new_trial)
  
          # Should queue the parent of the broken trial
          pbt._queue_trials_for_promotions([new_trial])
          assert len(pbt._queue) == 1
          assert pbt._queue[0].id == base_trial.id
  
      @pytest.mark.usefixtures("no_shutil_copytree")
      def test_queue_completed_trials_for_promotions(self, space):
          pbt = SpaceTransformAlgoWrapper(PBT, space).algorithm
  
          trial = sample_trials(pbt.space, num=1, status="completed", objective=1)[0]
          pbt.register(trial)
  
          # Should queue the trial itself
          pbt._queue_trials_for_promotions([trial])
          assert len(pbt._queue) == 1
          assert pbt._queue[0].id == trial.id
  
          new_trial = trial.branch(params={"f": pbt.fidelities[trial.params["f"]]})
          pbt.lineages.fork(trial, new_trial)
  
          new_trial.status = "completed"
          new_trial._results.append(
              Trial.Result(name="objective", type="objective", value=1)
          )
          pbt.register(new_trial)
  
          # Should queue the parent of the broken trial
          pbt._queue_trials_for_promotions([new_trial])
          assert len(pbt._queue) == 2
          assert pbt._queue[1].id == new_trial.id
  
      @pytest.mark.parametrize("status", ["new", "reserved", "interrupted"])
      def test_dont_queue_pending_trials_for_promotions(self, space, status):
          pbt = SpaceTransformAlgoWrapper(PBT, space).algorithm
  
          trial = sample_trials(pbt.space, num=1, status=status)[0]
          pbt.register(trial)
  
          # Should not queue anything
          pbt._queue_trials_for_promotions([trial])
          assert len(pbt._queue) == 0
  
  
  class TestPBTSuggest:
      def test_generate_offspring_unknown_trial(self, space):
  
          pbt = SpaceTransformAlgoWrapper(PBT, space).algorithm
          trial = sample_trials(pbt.space, 1)[0]
          with pytest.raises(RuntimeError, match="Trying to fork a trial that"):
              pbt._generate_offspring(trial)
  
      def test_generate_offspring_exploit_skip(self, space):
  
          pbt = SpaceTransformAlgoWrapper(
              PBT, space, exploit=ExploitStub(skip=True).configuration
          ).algorithm
          trial = sample_trials(pbt.space, 1, status="completed", objective=1)[0]
          pbt.register(trial)
  
          trial_to_branch, new_trial = pbt._generate_offspring(trial)
          assert trial_to_branch is None
          assert new_trial is None
  
      def test_generate_offspring_exploit_promote(self, space):
  
          pbt = SpaceTransformAlgoWrapper(
              PBT,
              space,
              exploit=ExploitStub().configuration,
              explore=ExploreStub(no_call=True).configuration,
          ).algorithm
          trial = sample_trials(pbt.space, 1, status="completed", objective=1)[0]
  
          # Apply the transformation and revert it to have lossy effect (like small precision)
          trial = pbt.space.transform(pbt.space.reverse(pbt.space.transform(trial)))
  
          pbt.register(trial)
  
          new_params_expected = trial.params
          new_params_expected["f"] = 10.9
  
          trial_to_branch, new_trial = pbt._generate_offspring(trial)
          assert trial_to_branch is trial
          assert new_trial.params == new_params_expected
  
      def test_generate_offspring_exploit_branch(self, space):
          pbt = SpaceTransformAlgoWrapper(
              PBT,
              space,
              exploit=ExploitStub(rval="toset").configuration,
              explore=ExploreStub(rval="toset").configuration,
          ).algorithm
  
          trials = sample_trials(pbt.space, 3, status="completed", objective=1)
  
          trial_to_promote = trials[0]
          exploited_trial = trials[1]
          new_params_expected = trials[2].params
  
          pbt.exploit_func.rval = exploited_trial
          pbt.explore_func.rval = new_params_expected
  
          # Make sure they are different
          assert new_params_expected != trial_to_promote.params
          assert new_params_expected != exploited_trial.params
  
          pbt.register(trials[0])
          pbt.register(trials[1])
  
          trial_to_branch, new_trial = pbt._generate_offspring(trial_to_promote)
  
          new_params_expected["f"] = 10.9
  
          assert trial_to_branch is exploited_trial
          assert new_trial.params["f"] == new_params_expected["f"]
          assert new_trial.params == new_params_expected
  
      def test_generate_offspring_timeout(self, space):
  
          pbt = SpaceTransformAlgoWrapper(
              PBT,
              space,
              exploit=ExploitStub(rval=None).configuration,
              explore=ExploreStub(rval="toset").configuration,
              fork_timeout=0.05,
          ).algorithm
          trial = sample_trials(pbt.space, 1, status="completed", objective=1)[0]
          pbt.explore_func.rval = trial.params
  
          pbt.register(trial)
          parent = trial.branch(params={"f": pbt.fidelities[space["f"].low]})
          pbt.register(parent)
  
          with pytest.raises(RuntimeError):
              pbt._generate_offspring(trial)
  
      def test_generate_offspring_retry_using_same_trial(self, space, monkeypatch):
          """Test that when exploit returns another trial, the base one is reused and case of
          duplicate samples
          """
          pbt = SpaceTransformAlgoWrapper(
              PBT,
              space,
              exploit=ExploitStub(rval="toset", should_receive="toset").configuration,
              explore=ExploreStub(rval="toset").configuration,
              fork_timeout=0.0001,
          ).algorithm
  
          trials = sample_trials(pbt.space, 3, status="completed", objective=1)
          parent_trial = trials[0]
          base_trial = trials[1]
          sample_params = trials[2].params
  
          pbt.exploit_func.rval = parent_trial
          pbt.exploit_func.should_receive = base_trial
          pbt.explore_func.rval = sample_params
  
          pbt.register(parent_trial)
          pbt.register(base_trial)
  
          # The trial sampled will already be registered
          sample_params["f"] = pbt.fidelities[space["f"].low]
          child = parent_trial.branch(params=sample_params)
          pbt.register(child)
  
          # Exploit will return parent_trial, but Explore will return params of child, sampling
          # a duplite, since child is already registered. ExploitStub.should_receive will
          # test that base_trial is passed as expected to exploit when attempting more attemps
          # of exploit and explore.
          with pytest.raises(RuntimeError):
              pbt._generate_offspring(base_trial)
  
      def test_fork_lineages_empty_queue(self, space):
          pbt = SpaceTransformAlgoWrapper(PBT, space).algorithm
          assert pbt._fork_lineages(10) == []
  
      def test_fork_lineages_skip_and_requeue_trials(self, space):
          num = 10
          pbt = SpaceTransformAlgoWrapper(
              PBT,
              space,
              exploit=ExploitStub(skip=True).configuration,
          ).algorithm
  
          trials = sample_trials(pbt.space, num, status="completed", objective=1)
  
          for trial in trials:
              pbt.register(trial)
  
          pbt._queue = trials[:]
  
          assert pbt._fork_lineages(num) == []
          assert len(pbt._queue) == num
          assert pbt._queue == trials
  
      @pytest.mark.usefixtures("no_shutil_copytree")
      def test_fork_lineages_promote_trial(self, space):
          num = 10
          pbt = SpaceTransformAlgoWrapper(
              PBT,
              space,
              exploit=ExploitStub(skip=None).configuration,
          ).algorithm
  
          trials = sample_trials(pbt.space, num, status="completed", objective=1)
  
          for trial in trials:
              pbt.register(trial)
  
          pbt._queue = trials[:]
  
          branched_trials = pbt._fork_lineages(num)
          assert len(trials) == num
          assert len(branched_trials) == num
          assert pbt._queue == []
  
          for trial, branched_trial in zip(trials, branched_trials):
              expected_params = trial.params
              expected_params["f"] = 10.9
              assert branched_trial.params == expected_params
  
      @pytest.mark.usefixtures("no_shutil_copytree")
      def test_fork_lineages_branch_trials(self, space):
          num = 10
  
          pbt = SpaceTransformAlgoWrapper(
              PBT,
              space,
              exploit=ExploitStub(rval="toset").configuration,
              fork_timeout=0.05,
          ).algorithm
  
          trials = sample_trials(pbt.space, num + 1, status="completed", objective=1)
          trial_to_branch = trials[-1]
          pbt.exploit_func.rval = trial_to_branch
          for trial in trials:
              pbt.register(trial)
  
          pbt._queue = trials[:-1]
  
          branched_trials = pbt._fork_lineages(num)
  
          assert len(trials) == num + 1
          assert len(branched_trials) == num
          assert pbt._queue == []
  
          for trial, branched_trial in zip(trials, branched_trials):
              # Check if parent is correct
              assert branched_trial.parent == trial_to_branch.id
              # Check in lineage if jump is set from correct base trial
              assert pbt.lineages.get_lineage(branched_trial).base.item.id == trial.id
              # Check if params are not duplicated
              should_not_be_params = trial_to_branch.params
              should_not_be_params["f"] = 10.9
              assert branched_trial.params["f"] == should_not_be_params["f"]
              assert branched_trial.params != should_not_be_params
  
      @pytest.mark.usefixtures("no_shutil_copytree")
      def test_fork_lineages_branch_duplicates(self, space):
          num = 10
          pbt = SpaceTransformAlgoWrapper(
              PBT,
              space,
              exploit=ExploitStub(rval="toset").configuration,
              explore=ExploreStub(rval="toset").configuration,
              fork_timeout=0.05,
          ).algorithm
  
          trials = sample_trials(pbt.space, num + 1, status="completed", objective=1)
          new_params_expected = trials[-1].params
          pbt.exploit_func.rval = trials[-1]
          pbt.explore_func.rval = new_params_expected
          for trial in trials:
              pbt.register(trial)
  
          pbt._queue = trials[:-1]
  
          with pytest.raises(RuntimeError):
              pbt._fork_lineages(num)
  
          # First queue.pop is fine, fails on second queue.pop.
          assert len(pbt._queue) == num - 2
  
      @pytest.mark.usefixtures("no_shutil_copytree")
      def test_fork_lineages_num_larger_than_queue(self, space):
          num = 10
          pbt = SpaceTransformAlgoWrapper(
              PBT,
              space,
              exploit=ExploitStub(rval=None).configuration,
          ).algorithm
  
          trials = sample_trials(pbt.space, num, status="completed", objective=1)
          for trial in trials:
              pbt.register(trial)
  
          pbt._queue = trials[:]
  
          num_fork = 4
          branched_trials = pbt._fork_lineages(num_fork)
  
          assert len(branched_trials) == num_fork
          assert len(pbt._queue) == num - num_fork
  
          trial_ids = [trial.id for trial in trials]
  
          assert [trial.parent for trial in branched_trials] == trial_ids[:num_fork]
          assert [trial.id for trial in pbt._queue] == trial_ids[num_fork:]
  
      @pytest.mark.usefixtures("no_shutil_copytree")
      def test_fork_lineages_num_smaller_than_queue(self, space):
          num = 4
          pbt = SpaceTransformAlgoWrapper(
              PBT,
              space,
              exploit=ExploitStub(rval=None).configuration,
          ).algorithm
  
          trials = sample_trials(pbt.space, num, status="completed", objective=1)
          for trial in trials:
              pbt.register(trial)
  
          pbt._queue = trials[:]
  
          num_fork = 10
          branched_trials = pbt._fork_lineages(num_fork)
  
          assert len(branched_trials) == num
          assert len(pbt._queue) == 0
  
          trial_ids = [trial.id for trial in trials]
  
          assert [trial.parent for trial in branched_trials] == trial_ids
  
      def test_suggest_num_population_size_sample(self, space, mocker):
          population_size = 10
          pbt = SpaceTransformAlgoWrapper(
              PBT, space, population_size=population_size
          ).algorithm
  
          pbt_sample_mock = mocker.spy(pbt, "_sample")
          pbt_fork_mock = mocker.spy(pbt, "_fork_lineages")
  
          num = 6
          assert len(pbt.suggest(num)) == num
  
          pbt_sample_mock.assert_called_with(num)
          pbt_fork_mock.assert_called_with(0)
  
          assert len(pbt.suggest(num)) == 4
  
          pbt_sample_mock.assert_called_with(4)
          pbt_fork_mock.assert_called_with(2)
  
      def test_suggest_num_population_size_sample_broken(self, space, mocker):
          population_size = 10
          pbt = SpaceTransformAlgoWrapper(
              PBT, space, population_size=population_size
          ).algorithm
  
          pbt_sample_mock = mocker.spy(pbt, "_sample")
          pbt_fork_mock = mocker.spy(pbt, "_fork_lineages")
  
          num = 10
          trials = pbt.suggest(num)
          assert len(trials) == num
  
          pbt_sample_mock.assert_called_with(num)
          pbt_fork_mock.assert_called_with(0)
  
          n_broken = 3
          for trial in trials[:n_broken]:
              trial.status = "broken"
  
          pbt.observe(trials)
  
          assert len(pbt.suggest(num)) == n_broken
  
          # 3 trials are broken, need to resample 3 trials, and can try to fork 7 trials
          pbt_sample_mock.assert_called_with(n_broken)
          pbt_fork_mock.assert_called_with(7)
  
      @pytest.mark.usefixtures("no_shutil_copytree")
      def test_suggest_num_population_size_fork_completed(self, space, mocker):
          population_size = 10
          pbt = SpaceTransformAlgoWrapper(
              PBT,
              space,
              population_size=population_size,
              exploit=ExploitStub(rval=None).configuration,
          ).algorithm
  
          pbt_sample_mock = mocker.spy(pbt, "_sample")
          pbt_fork_mock = mocker.spy(pbt, "_fork_lineages")
  
          num = 4
          trials = pbt.suggest(num)
          assert len(trials) == num
  
          pbt_sample_mock.assert_called_with(num)
          pbt_fork_mock.assert_called_with(0)
  
          n_completed = 3
          for trial in trials[:n_completed]:
              trial.exp_working_dir = "/nothing"
              trial.status = "completed"
              trial._results.append(
                  Trial.Result(name="objective", type="objective", value=1)
              )
  
          pbt.observe(trials)
          assert len(pbt._queue) == n_completed
  
          # There are 4 trials sampled, out of which 3 are completed. Still missing 6 trials
          # for base population.
          assert len(pbt.suggest(num)) == num
          pbt_sample_mock.assert_called_with(num)
          pbt_fork_mock.assert_called_with(0)
  
          # There are 8 trials sampled, out of which 3 are completed. Still missing 2 trials
          # for base population.
          assert len(pbt.suggest(num)) == num
          pbt_sample_mock.assert_called_with(2)
          pbt_fork_mock.assert_called_with(2)
  
  
  population_size = 10
  generations = 5
  
  
  @pytest.mark.usefixtures("no_shutil_copytree")
  class TestGenericPBT(BaseAlgoTests):
      algo_name = "pbt"
      max_trials = population_size * generations
      config = {
          "seed": 123456,
          "population_size": population_size,
          "generations": generations,
          "exploit": {
              "of_type": "PipelineExploit",
              "exploit_configs": [
                  {
                      "of_type": "BacktrackExploit",
                      "min_forking_population": population_size / 2,
                      "candidate_pool_ratio": 0.0,
                      "truncation_quantile": 1.0,
                  },
                  {
                      "of_type": "TruncateExploit",
                      "min_forking_population": population_size / 2,
                      "candidate_pool_ratio": 0.3,
                      "truncation_quantile": 0.9,
                  },
              ],
          },
          "explore": {
              "of_type": "PipelineExplore",
              "explore_configs": [
                  {
                      "of_type": "ResampleExplore",
                      "probability": 0.3,
                  },
                  {
                      "of_type": "PerturbExplore",
                      "factor": 1.5,
                      "volatility": 0.005,
                  },
              ],
          },
          "fork_timeout": 5,
      }
      space = {"x": "uniform(0, 1)", "y": "uniform(0, 1)", "f": "fidelity(1, 10, base=1)"}
  
      def test_no_fidelity(self):
          space = self.create_space({"x": "uniform(0, 1)", "y": "uniform(0, 1)"})
  
          with pytest.raises(
              RuntimeError, match="PBT cannot be used if space does not contain"
          ):
              self.create_algo(space=space)
  
      @pytest.mark.skip(
          reason="There are no good reasons to use PBT if search space is so small"
      )
      def test_is_done_cardinality(self):
          pass
  
      @pytest.mark.parametrize("num", [100000, 1])
      def test_is_done_max_trials(self, num):
          space = self.create_space()
  
          MAX_TRIALS = 10
          algo = self.create_algo(space=space)
          algo.algorithm.max_trials = MAX_TRIALS
  
          objective = 0
          while not algo.is_done:
              trials = algo.suggest(num)
              assert trials is not None
              if trials:
                  self.observe_trials(trials, algo, objective)
                  objective += len(trials)
  
          # BPT should ignore max trials.
          assert algo.n_observed > MAX_TRIALS
          # BPT should stop when all trials of last generation are completed.
          assert algo.n_observed == population_size * (generations + 1)
          assert algo.is_done
  
      @pytest.mark.skip(reason="See https://github.com/Epistimio/orion/issues/599")
      def test_optimize_branin(self):
          pass
  
      def assert_callbacks(self, spy, num, algo):
          def check_population_size(gen_population_size, depth, expected):
              assert (
                  gen_population_size == expected
              ), f"population of {gen_population_size} at depth {depth}, should be {expected}"
  
          pbt = algo.algorithm
          remaining_num = num
  
          for depth in range(generations):
              gen_population_size = len(pbt.lineages.get_trials_at_depth(depth))
              if remaining_num > population_size:
                  expected_population_size = population_size
              else:
                  expected_population_size = remaining_num
  
              check_population_size(gen_population_size, depth, expected_population_size)
  
              remaining_num = max(remaining_num - expected_population_size, 0)
  
  
  TestGenericPBT.set_phases(
      [
          ("random", 5, "space.sample"),
          ("generation_2", 2 * population_size, "_generate_offspring"),
          ("generation_3", 3 * population_size, "_generate_offspring"),
      ]
  )
  