import numpy as np
import pytest
from numpy.testing import assert_almost_equal, assert_array_equal

from skopt import gp_minimize
from skopt.benchmarks import bench1, bench2, bench3, bench4, branin
from skopt.space.space import Categorical, Real, Space
from skopt.utils import cook_estimator


def check_minimize(
    func,
    y_opt,
    bounds,
    acq_optimizer,
    acq_func,
    margin,
    n_calls,
    n_initial_points=10,
    init_gen="random",
):
    r = gp_minimize(
        func,
        bounds,
        acq_optimizer=acq_optimizer,
        acq_func=acq_func,
        n_initial_points=n_initial_points,
        n_calls=n_calls,
        random_state=1,
        initial_point_generator=init_gen,
        noise=1e-10,
    )
    assert r.fun < y_opt + margin


SEARCH = ["sampling", "lbfgs"]
ACQUISITION = ["LCB", "EI"]
INITGEN = ["random", "lhs", "halton", "hammersly", "sobol"]


@pytest.mark.slow_test
@pytest.mark.parametrize("search", SEARCH)
@pytest.mark.parametrize("acq", ACQUISITION)
def test_gp_minimize_bench1(search, acq):
    check_minimize(bench1, 0.0, [(-2.0, 2.0)], search, acq, 0.05, 20)


@pytest.mark.slow_test
@pytest.mark.parametrize("search", ["sampling"])
@pytest.mark.parametrize("acq", ["LCB", "MES"])
@pytest.mark.parametrize("initgen", INITGEN)
def test_gp_minimize_bench1_initgen(search, acq, initgen):
    check_minimize(bench1, 0.0, [(-2.0, 2.0)], search, acq, 0.05, 20, init_gen=initgen)


@pytest.mark.slow_test
@pytest.mark.parametrize("search", SEARCH)
@pytest.mark.parametrize("acq", ACQUISITION)
def test_gp_minimize_bench2(search, acq):
    check_minimize(bench2, -5, [(-6.0, 6.0)], search, acq, 0.05, 20)


@pytest.mark.slow_test
@pytest.mark.parametrize("search", SEARCH)
@pytest.mark.parametrize("acq", ACQUISITION)
def test_gp_minimize_bench3(search, acq):
    check_minimize(bench3, -0.9, [(-2.0, 2.0)], search, acq, 0.05, 20)


@pytest.mark.fast_test
@pytest.mark.parametrize("search", ["sampling"])
@pytest.mark.parametrize("acq", ["LCB", "EI", "MES"])
def test_gp_minimize_bench4(search, acq):
    # this particular random_state picks "2" twice so we can make an extra
    # call to the objective without repeating options
    check_minimize(bench4, 0, [["-2", "-1", "0", "1", "2"]], search, acq, 1.05, 20)


@pytest.mark.fast_test
def test_n_jobs():
    r_single = gp_minimize(
        bench3,
        [(-2.0, 2.0)],
        acq_optimizer="lbfgs",
        acq_func="EI",
        n_calls=4,
        n_initial_points=2,
        random_state=1,
        noise=1e-10,
    )
    r_double = gp_minimize(
        bench3,
        [(-2.0, 2.0)],
        acq_optimizer="lbfgs",
        acq_func="EI",
        n_calls=4,
        n_initial_points=2,
        random_state=1,
        noise=1e-10,
        n_jobs=2,
    )
    assert_array_equal(r_single.x_iters, r_double.x_iters)


@pytest.mark.fast_test
def test_gpr_default():
    """Smoke test that gp_minimize does not fail for default values."""
    gp_minimize(branin, ((-5.0, 10.0), (0.0, 15.0)), n_initial_points=2, n_calls=2)


@pytest.mark.fast_test
def test_use_given_estimator():
    """Test that gp_minimize does not use default estimator if one is passed in
    explicitly."""
    domain = [(1.0, 2.0), (3.0, 4.0)]
    noise_correct = 1e5
    noise_fake = 1e-10
    estimator = cook_estimator("GP", domain, noise=noise_correct)
    res = gp_minimize(
        branin,
        domain,
        n_calls=4,
        n_initial_points=2,
        base_estimator=estimator,
        noise=noise_fake,
    )

    assert res['models'][-1].noise == noise_correct


@pytest.mark.fast_test
def test_use_given_estimator_with_max_model_size():
    """Test that gp_minimize does not use default estimator if one is passed in
    explicitly."""
    domain = [(1.0, 2.0), (3.0, 4.0)]
    noise_correct = 1e5
    noise_fake = 1e-10
    estimator = cook_estimator("GP", domain, noise=noise_correct)
    res = gp_minimize(
        branin,
        domain,
        n_calls=4,
        n_initial_points=2,
        base_estimator=estimator,
        noise=noise_fake,
        model_queue_size=1,
    )
    assert len(res['models']) == 1
    assert res['models'][-1].noise == noise_correct


@pytest.mark.fast_test
def test_categorical_integer():
    def f(params):
        return np.random.uniform()

    dims = [[1]]
    res = gp_minimize(f, dims, n_calls=2, n_initial_points=2, random_state=1)
    assert res.x_iters[0][0] == dims[0][0]


@pytest.mark.parametrize("initgen", INITGEN)
def test_mixed_categoricals(initgen):

    space = Space(
        [
            Categorical(name="x", categories=["1", "2", "3"]),
            Categorical(name="y", categories=[4, 5, 6]),
            Real(name="z", low=1.0, high=5.0),
        ]
    )

    def objective(param_list):
        x = param_list[0]
        y = param_list[1]
        z = param_list[2]
        loss = int(x) + y * z
        return loss

    res = gp_minimize(
        objective, space, n_calls=20, random_state=1, initial_point_generator=initgen
    )
    assert res["x"][:2] in [['1', 4], ['2', 4]]
    assert_almost_equal(res["x"][2], 1.0, decimal=3)


@pytest.mark.parametrize("initgen", INITGEN)
def test_mixed_categoricals2(initgen):
    space = Space(
        [
            Categorical(name="x", categories=["1", "2", "3"]),
            Categorical(name="y", categories=[4, 5, 6]),
        ]
    )

    def objective(param_list):
        x = param_list[0]
        y = param_list[1]
        loss = int(x) + y
        return loss

    res = gp_minimize(
        objective, space, n_calls=12, random_state=1, initial_point_generator=initgen
    )
    assert res["x"] == ['1', 4]
