import logging
import math

import distgfs

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def levi(x, y):
    """
    Levi's function (see https://en.wikipedia.org/wiki/Test_functions_for_optimization).
    Has a global _minimum_ of 0 at x=1, y=1.
    """
    a = math.sin(3.0 * math.pi * x) ** 2
    b = (x - 1) ** 2 * (1 + math.sin(3.0 * math.pi * y) ** 2)
    c = (y - 1) ** 2 * (1 + math.sin(2.0 * math.pi * y) ** 2)
    return a + b + c


def reduce_fun(xs):
    return xs[0]


def obj_fun(pp, pid):
    """Objective function to be _maximized_ by GFS."""
    res = levi(**pp)
    logger.info(f"Iter: {pid}\t x:{pp['x']}, y:{pp['y']}, result:{res}")
    # Since Dlib maximizes, but we want to find the minimum,
    # we negate the result before passing it to the Dlib optimizer.
    return -res


if __name__ == "__main__":
    # For this example, we pretend that we want to keep 'y' fixed at 1.0
    # while optimizing 'x' in the range -4.5 to 4.5
    space = {"x": [-4.5, 4.5]}
    problem_parameters = {"y": 1.0}

    # Create an optimizer parameter set
    distgfs_params = {
        "opt_id": "distgfs_levi",
        "obj_fun_name": "obj_fun",
        "obj_fun_module": "example_distgfs_levi_nprocs",
        "reduce_fun_name": "reduce_fun",
        "reduce_fun_module": "example_distgfs_levi_nprocs",
        "problem_parameters": problem_parameters,
        "space": space,
        "n_iter": 20,
    }

    distgfs.run(distgfs_params, nprocs_per_worker=1, spawn_workers=True, verbose=True)
