# Copyright © 2022 Intel Corporation
#
# SPDX-License-Identifier: MIT
import copy

import pytest
from mock import Mock

from sigopt.xgboost.constants import DEFAULT_CLASSIFICATION_METRIC, DEFAULT_EVALS_NAME, DEFAULT_REGRESSION_METRIC
from sigopt.xgboost.experiment import XGBExperiment


EXPERIMENT_CONFIG_BASE = dict(
  name="Single metric optimization",
  type="offline",
  parameters=[
    dict(name="eta", type="double", bounds={"min": 0.1, "max": 0.5}),
    dict(name="max_depth", type="int", bounds={"min": 2, "max": 6}),
    dict(name="num_boost_round", type="int", bounds={"min": 2, "max": 6}),
  ],
  metrics=[dict(name="accuracy", strategy="optimize", objective="maximize")],
  parallel_bandwidth=1,
  budget=2,
)
PARAMS_BASE = {
  "num_class": 3,
  "lambda": 1,
}


def verify_experiment_config_integrity(experiment_config):
  assert isinstance(experiment_config, dict)
  assert "type" in experiment_config
  assert "parameters" in experiment_config
  assert "metrics" in experiment_config
  assert "budget" in experiment_config

  parameters = experiment_config["parameters"]
  for parameter in parameters:
    assert "name" in parameter
    assert "type" in parameter
    if parameter["type"] in ["int", "double"]:
      assert "bounds" in parameter
    else:
      assert "categorical_values" in parameter

  metrics = experiment_config["metrics"]
  for metric in metrics:
    assert "name" in metric
    assert "strategy" in metric
    assert "objective" in metric


def parse_and_create_aiexperiment_config(experiment_config, params):
  num_boost_round = None
  run_options = None
  early_stopping_rounds = 10
  d_train = Mock()
  evals = Mock()
  xgb_experiment = XGBExperiment(
    experiment_config,
    d_train,
    evals,
    params,
    num_boost_round,
    run_options,
    early_stopping_rounds,
  )
  xgb_experiment.parse_and_create_metrics()
  xgb_experiment.parse_and_create_parameters()
  return xgb_experiment


class TestExperimentConfig:
  def verify_integrity(self, experiment_config, params):
    xgb_experiment = parse_and_create_aiexperiment_config(experiment_config, params)
    verify_experiment_config_integrity(xgb_experiment.experiment_config_parsed)

  def test_base(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    params = copy.deepcopy(PARAMS_BASE)
    self.verify_integrity(experiment_config, params)

  def test_config_no_search_space(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    params = copy.deepcopy(PARAMS_BASE)
    del experiment_config["parameters"]
    self.verify_integrity(experiment_config, params)

  def test_config_search_space_name_only(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    params = copy.deepcopy(PARAMS_BASE)
    for parameter in experiment_config["parameters"]:
      del parameter["type"]
      del parameter["bounds"]
    self.verify_integrity(experiment_config, params)

  def test_config_detect_log_transformation(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    params = copy.deepcopy(PARAMS_BASE)
    experiment_config["parameters"] = [dict(name="eta")]
    xgb_experiment = parse_and_create_aiexperiment_config(experiment_config, params)
    assert xgb_experiment.experiment_config_parsed["parameters"][0]["transformation"] == "log"

  def test_config_search_space_mixed(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    params = copy.deepcopy(PARAMS_BASE)
    del experiment_config["parameters"][2]["type"]
    del experiment_config["parameters"][2]["bounds"]
    self.verify_integrity(experiment_config, params)

  def test_config_search_space_wrong_type(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    params = copy.deepcopy(PARAMS_BASE)
    experiment_config["parameters"][0]["type"] = "int"
    del experiment_config["parameters"][0]["bounds"]
    with pytest.raises(ValueError):
      self.verify_integrity(experiment_config, params)

  def test_config_search_space_no_type(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    params = copy.deepcopy(PARAMS_BASE)
    del experiment_config["parameters"][0]["type"]
    del experiment_config["parameters"][1]["type"]
    del experiment_config["parameters"][2]["type"]
    self.verify_integrity(experiment_config, params)

  def test_config_search_space_categories_no_type(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    params = copy.deepcopy(PARAMS_BASE)
    experiment_config["parameters"].append(
      dict(
        name="tree_method",
        categorical_values=["auto", "exact", "hist", "gpu_hist"],
      )
    )
    self.verify_integrity(experiment_config, params)

  def test_config_search_space_no_categorical_values(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    params = copy.deepcopy(PARAMS_BASE)
    experiment_config["parameters"].append(
      dict(
        name="tree_method",
        type="categorical",
      )
    )
    with pytest.raises(ValueError):
      self.verify_integrity(experiment_config, params)

  def test_config_search_space_wrong_categories(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    params = copy.deepcopy(PARAMS_BASE)
    experiment_config["parameters"].append(
      dict(
        name="tree_method",
        type="categorical",
        categorical_values=["auto", "exact", "hist", "gpu_hist", "WrongCategory"],
      )
    )
    with pytest.raises(ValueError):
      self.verify_integrity(experiment_config, params)

  def test_config_search_space_fake_categories(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    params = copy.deepcopy(PARAMS_BASE)
    experiment_config["parameters"].append(
      dict(
        name="foo",
        type="categorical",
        categorical_values=["auto", "exact", "hist", "gpu_hist", "WrongCategory"],
      )
    )
    self.verify_integrity(experiment_config, params)

  def test_config_no_supported_bounds(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    experiment_config["parameters"].append(dict(name="max_leaves"))
    params = copy.deepcopy(PARAMS_BASE)
    with pytest.raises(ValueError):
      self.verify_integrity(experiment_config, params)

  def test_autodetect_metric_from_objective(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    del experiment_config["metrics"]
    params = copy.deepcopy(PARAMS_BASE)

    params["objective"] = "binary:logistic"
    xgb_experiment = parse_and_create_aiexperiment_config(experiment_config, params)
    assert xgb_experiment.experiment_config_parsed["metrics"][0]["name"] == "-".join(
      (DEFAULT_EVALS_NAME, DEFAULT_CLASSIFICATION_METRIC)
    )

    params["objective"] = "multi:softmax"
    xgb_experiment = parse_and_create_aiexperiment_config(experiment_config, params)
    assert xgb_experiment.experiment_config_parsed["metrics"][0]["name"] == "-".join(
      (DEFAULT_EVALS_NAME, DEFAULT_CLASSIFICATION_METRIC)
    )

    params["objective"] = "reg:squarederror"
    xgb_experiment = parse_and_create_aiexperiment_config(experiment_config, params)
    assert xgb_experiment.experiment_config_parsed["metrics"][0]["name"] == "-".join(
      (DEFAULT_EVALS_NAME, DEFAULT_REGRESSION_METRIC)
    )

  def test_config_metric_string_only(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    params = copy.deepcopy(PARAMS_BASE)
    experiment_config["metrics"] = "accuracy"
    self.verify_integrity(experiment_config, params)

  def test_config_metric_list(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    experiment_config["metrics"].append(dict(name="f1", strategy="store", objective="maximize"))
    params = copy.deepcopy(PARAMS_BASE)
    self.verify_integrity(experiment_config, params)

  def test_config_param_defined_twice(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    params = copy.deepcopy(PARAMS_BASE)
    params["eta"] = 0.1
    with pytest.raises(ValueError):
      self.verify_integrity(experiment_config, params)

  def test_config_num_boost_round_defined_twice(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    params = copy.deepcopy(PARAMS_BASE)
    params["num_boost_round"] = 10
    with pytest.raises(ValueError):
      self.verify_integrity(experiment_config, params)

  def test_config_wrong_metric_string(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    experiment_config["metrics"] = "NOT_A_METRIC_SUPPORTED"
    params = copy.deepcopy(PARAMS_BASE)
    with pytest.raises(ValueError):
      self.verify_integrity(experiment_config, params)

  def test_config_wrong_metric_list(self):
    experiment_config = copy.deepcopy(EXPERIMENT_CONFIG_BASE)
    experiment_config["metrics"][0]["name"] = "NOT_A_METRIC_SUPPORTED"
    params = copy.deepcopy(PARAMS_BASE)
    with pytest.raises(ValueError):
      self.verify_integrity(experiment_config, params)
