from utils import wrap_untrustworthy, final, ExhaustedTrialsError
import numpy as np


def make_evolopy(optimizer_name, popsize=None):
    from evolopy.BAT import BAT
    from evolopy.CS import CS
    from evolopy.DE import DE
    from evolopy.FFA import FFA
    from evolopy.GA import GA
    from evolopy.GWO import GWO
    from evolopy.HHO import HHO
    from evolopy.JAYA import JAYA
    from evolopy.MFO import MFO
    from evolopy.MVO import MVO
    from evolopy.PSO import PSO
    from evolopy.SCA import SCA
    from evolopy.SSA import SSA
    from evolopy.WOA import WOA
    from unittest.mock import patch

    def no_print(*args, **kwargs):
        pass

    optimizers = dict(
        BAT=BAT,
        CS=CS,
        DE=DE,
        FFA=FFA,
        GA=GA,
        GWO=GWO,
        HHO=HHO,
        JAYA=JAYA,
        MFO=MFO,
        MVO=MVO,
        PSO=PSO,
        SCA=SCA,
        SSA=SSA,
        WOA=WOA,
    )
    optimizers = {
        k: patch("builtins.print", no_print)(v) for k, v in optimizers.items()
    }
    optimizer = optimizers[optimizer_name.upper()]

    def f(objective, n_trials, n_dim, with_count):
        # ps = n_dim if popsize is None else popsize
        # ps = max(4, ps) if optimizer_name.upper() == "DE" else ps
        ps = int(4 + 3 * np.log(n_dim)) if popsize is None else popsize
        if optimizer_name.upper() == "GA" and ps & 1:
            ps += 1  # force popsize to be even

        _objective = wrap_untrustworthy(
            objective, n_trials, bounding="sine", raising=True
        )
        try:
            optimizer(_objective, 0.0, 1.0, n_dim, ps, n_trials)
        except ExhaustedTrialsError:
            pass

        fopt, xopt, feval_count = _objective(final)
        return (fopt, xopt, feval_count) if with_count else (fopt, xopt)

    name = f"evolopy_{optimizer.__name__.lower()}"
    name += f"_auto" if popsize is None else f"_ps{popsize:02}"
    f.__name__ = name + "_cube"
    return f