from scipy.optimize import Bounds
from unittest.mock import patch
from utils import wrap_untrustworthy, final
import numpy as np


def _fix_logging(original_function):
    from functools import wraps

    @wraps(original_function)
    def wrapped_function(*args, **kwargs):
        with patch("logging.FileHandler"):  # do not create a file on import
            import fcmaes.optimizer
        fcmaes.optimizer._logger = None
        return original_function(*args, **kwargs)

    return wrapped_function


@_fix_logging
def make_biteopt(depth=1):
    from fcmaes.optimizer import Bite_cpp

    def f(objective, n_trials, n_dim, with_count):
        _objective = wrap_untrustworthy(objective, n_trials)
        bounds = Bounds([0.0] * n_dim, [1.0] * n_dim)
        optim = Bite_cpp(max_evaluations=n_trials, M=depth)
        _xopt, _fopt, _feval_count = optim.minimize(_objective, bounds)
        fopt, xopt, feval_count = _objective(final)
        return (fopt, xopt, feval_count) if with_count else (fopt, xopt)

    name = f"fcmaes_biteopt_d{depth}"
    f.__name__ = name + "_cube"
    return f


@_fix_logging
def make_csma(isigma=3):
    from fcmaes.optimizer import Csma_cpp

    def f(objective, n_trials, n_dim, with_count):
        _objective = wrap_untrustworthy(objective, n_trials)
        bounds = Bounds([0.0] * n_dim, [1.0] * n_dim)
        optim = Csma_cpp(max_evaluations=n_trials)
        _xopt, _fopt, _feval_count = optim.minimize(
            _objective, bounds, sdevs=1 / isigma
        )
        fopt, xopt, feval_count = _objective(final)
        return (fopt, xopt, feval_count) if with_count else (fopt, xopt)

    name = f"fcmaes_csma_is{isigma:02}"
    f.__name__ = name + "_cube"
    return f


@_fix_logging
def make_fcmaes(popsize=None):
    from fcmaes.optimizer import Cma_cpp

    def f(objective, n_trials, n_dim, with_count):
        _objective = wrap_untrustworthy(objective, n_trials)
        bounds = Bounds([0.0] * n_dim, [1.0] * n_dim)
        ps = int(4 + 3 * np.log(n_dim)) if popsize is None else popsize
        optim = Cma_cpp(
            max_evaluations=n_trials,
            popsize=ps,
            stop_hist=0.0,
            workers=1,
            delayed_update=False,
        )
        _xopt, _fopt, _feval_count = optim.minimize(_objective, bounds, sdevs=1 / 3)
        fopt, xopt, feval_count = _objective(final)
        return (fopt, xopt, feval_count) if with_count else (fopt, xopt)

    name = "fcmaes_cma"
    name += f"_auto" if popsize is None else f"_ps{popsize}"
    f.__name__ = name + "_cube"
    return f


@_fix_logging
def make_crfmnes(popsize=None):
    from fcmaes.optimizer import Crfmnes_cpp

    def f(objective, n_trials, n_dim, with_count):
        _objective = wrap_untrustworthy(objective, n_trials)
        bounds = Bounds([0.0] * n_dim, [1.0] * n_dim)
        ps = int(4 + 3 * np.log(n_dim)) if popsize is None else popsize
        optim = Crfmnes_cpp(max_evaluations=n_trials, popsize=ps)
        _xopt, _fopt, _feval_count = optim.minimize(_objective, bounds, sdevs=1 / 3)
        fopt, xopt, feval_count = _objective(final)
        return (fopt, xopt, feval_count) if with_count else (fopt, xopt)

    name = "fcmaes_crfmnes"
    name += f"_auto" if popsize is None else f"_ps{popsize}"
    f.__name__ = name + "_cube"
    return f


_broken = """
@_fix_logging
def make_lde(popsize=None):
    from fcmaes.optimizer import LDe_cpp

    def f(objective, n_trials, n_dim, with_count):
        _objective = wrap_untrustworthy(objective, n_trials)
        bounds = Bounds([0.0] * n_dim, [1.0] * n_dim)
        optim = LDe_cpp(max_evaluations=n_trials, popsize=popsize)
        _xopt, _fopt, _feval_count = optim.minimize(_objective, bounds)
        fopt, xopt, feval_count = _objective(final)
        return (fopt, xopt, feval_count) if with_count else (fopt, xopt)

    name = "fcmaes_lde"
    name += f"_auto" if popsize is None else f"_ps{popsize}"
    f.__name__ = name + "_cube"
    return f
"""


@_fix_logging
def make_da(local=True):
    from fcmaes.optimizer import Da_cpp

    def f(objective, n_trials, n_dim, with_count):
        _objective = wrap_untrustworthy(objective, n_trials)
        bounds = Bounds([0.0] * n_dim, [1.0] * n_dim)
        optim = Da_cpp(max_evaluations=n_trials, use_local_search=local)
        _xopt, _fopt, _feval_count = optim.minimize(_objective, bounds)
        fopt, xopt, feval_count = _objective(final)
        return (fopt, xopt, feval_count) if with_count else (fopt, xopt)

    name = "fcmaes_da"
    if local:
        name += "_local"
    f.__name__ = name + "_cube"
    return f


@_fix_logging
def make_gclde(popsize=None):
    from fcmaes.optimizer import GCLDE_cpp

    def f(objective, n_trials, n_dim, with_count):
        _objective = wrap_untrustworthy(objective, n_trials)
        bounds = Bounds([0.0] * n_dim, [1.0] * n_dim)
        optim = GCLDE_cpp(max_evaluations=n_trials, popsize=popsize)
        _xopt, _fopt, _feval_count = optim.minimize(_objective, bounds)
        fopt, xopt, feval_count = _objective(final)
        return (fopt, xopt, feval_count) if with_count else (fopt, xopt)

    name = "fcmaes_gclde"
    name += f"_auto" if popsize is None else f"_ps{popsize}"
    f.__name__ = name + "_cube"
    return f


@_fix_logging
def make_lclde(popsize=None):
    from fcmaes.optimizer import LCLDE_cpp

    def f(objective, n_trials, n_dim, with_count):
        _objective = wrap_untrustworthy(objective, n_trials)
        bounds = Bounds([0.0] * n_dim, [1.0] * n_dim)
        optim = LCLDE_cpp(max_evaluations=n_trials, popsize=popsize)
        _xopt, _fopt, _feval_count = optim.minimize(_objective, bounds)
        fopt, xopt, feval_count = _objective(final)
        return (fopt, xopt, feval_count) if with_count else (fopt, xopt)

    name = "fcmaes_lclde"
    name += f"_auto" if popsize is None else f"_ps{popsize}"
    f.__name__ = name + "_cube"
    return f