from utils import wrap_untrustworthy, check, final
import numpy as np
import tinytweaks as tt


def make_birect(deepness=23, *, longest=False, pruning=False):
    from birect import birect

    def f(objective, n_trials, n_dim, with_count):
        feval_count = 0

        def _objective(x):
            nonlocal feval_count
            feval_count += 1
            return objective(x)

        xopt, fopt = birect(
            _objective,
            (0,) * n_dim,
            (1,) * n_dim,
            min_diag=2.0 ** -float(deepness),
            max_evals=n_trials,
            by_longest=longest,
            pruning=pruning,
            F=np.float64,
        )
        return (fopt, xopt, feval_count) if with_count else (fopt, xopt)

    name = f"birect{deepness:02}"
    name += "_longest" if longest else ""
    name += "_pruning" if pruning else ""
    f.__name__ = name + "_cube"
    return f


def make_soo(deepness=None, *, K=3):
    if deepness is None:
        deepness = int(31 * np.log(2) / np.log(K) - 1e-8)
    assert K >= 2
    from soo import soo

    def f(objective, n_trials, n_dim, with_count):
        feval_count = 0

        def _objective(x):
            nonlocal feval_count
            feval_count += 1
            return objective(x)

        xopt, history = soo(
            _objective, np.full(n_dim, 0.5), 0.5, n_trials, K=K, h_max=deepness
        )
        fopt = history[-1]
        return (fopt, xopt, feval_count) if with_count else (fopt, xopt)

    name = f"soo{deepness}_k{K}"
    f.__name__ = name + "_cube"
    return f


def make_mercury(
    flags, bounding="clip", *, isigma=tt.IV, popsize=2, irate=1, seed=None
):
    from hg import minimize as hg

    def f(objective, n_trials, n_dim, with_count):
        _objective = wrap_untrustworthy(objective, n_trials, bounding=bounding)

        init = (0.5,) * n_dim
        iterations = (n_trials - 1 + popsize * 2) // (popsize * 2 + 1)
        center, history = hg(
            _objective,
            init,
            iterations,
            sigma=isigma if callable(isigma) else 1 / isigma,
            popsize=popsize,
            rate=irate if callable(irate) else 1 / irate,
            seed=seed,
            _flags=flags,
        )

        # use our remaining evals on something, anything.
        feval_count = _objective(check)
        wasted = max(0, -(feval_count - n_trials))
        if wasted:
            print(f"wasting {wasted} of {n_trials} evaluations")
        for _ in range(wasted):
            prng = np.random.default_rng()
            x = prng.uniform(size=n_dim)
            fx = _objective(x)

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

    name = f"hg{flags:02}"
    name += f"_{bounding}" if bounding != "clip" else ""
    name += f"_ps{popsize:02}" if popsize != 2 else ""
    if isigma != tt.IV:
        name += f"_is{isigma.__name__}" if callable(isigma) else f"_is{isigma:02}"
    if irate != 1:
        name += f"_ir{irate.__name__}" if callable(irate) else f"_ir{irate:02}"
    f.__name__ = name + "_cube"
    return f