# Copyright 2023 Google LLC.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

"""Tests for hpob."""

from absl import app
from absl import logging

from vizier._src.benchmarks.experimenters import hpob_experimenter
from vizier._src.benchmarks.experimenters.hpob import handler as handler_lib

from absl.testing import absltest
from absl.testing import parameterized

# Global variable to be filled in with the class (not an instance) that shall
# be executed as the actual test. The class will be generated by
# generate_test_class().
class_to_test = None


def generate_test_class():
  """Generates all test cases. Must be called after InitGoogle()."""

  handler = handler_lib.HPOBHandler(
      root_dir=hpob_experimenter.ROOT_DIR,
      mode=hpob_experimenter.DEFAULT_TEST_MODE,
      surrogates_dir=hpob_experimenter.SURROGATES_DIR)

  class HpobTest(parameterized.TestCase):

    @absltest.skip("Files must be installed manually.")
    @parameterized.parameters(
        dict(normalize_y=True), dict(normalize_y=False),
        dict(normalize_y=True, use_surrogate_values=True),
        dict(normalize_y=False, use_surrogate_values=True),
        *tuple(dict(na_policy=x) for x in hpob_experimenter.NaPolicy))
    def test_all(self, **kwargs):
      container = hpob_experimenter.HPOBContainer(handler, **kwargs)
      for search_space_id, dataset_id in container.dataset_keys(
          hpob_experimenter.TEST):
        problem_and_trials = container.get_problem_and_trials(
            search_space_id, dataset_id)
        experimenter = container.get_experimenter(search_space_id, dataset_id)
        _ = experimenter.problem_statement()
        # TODO: Add more validations on the problem statement.
        objective = experimenter.EvaluateContinuous(
            problem_and_trials.trials[0])

        experimenter.evaluate([problem_and_trials.trials[0]])

        logging.info('First five trial metrics: %s', [
            t.final_measurement.metrics[hpob_experimenter.METRIC_NAME].value
            for t in problem_and_trials.trials[:5]
        ])
        logging.info('objective: %f', objective)
        logging.info('evaluated: %s', problem_and_trials.trials[0])

        self.assertAlmostEqual(
            problem_and_trials.trials[0].final_measurement.metrics[
                hpob_experimenter.METRIC_NAME].value,
            objective,
            places=1)

        break

  global class_to_test
  class_to_test = HpobTest


if __name__ == '__main__':
  app.call_after_init(generate_test_class)
  absltest.main()
