from .math import scalar_softplus
from .utils_np import do_bounding

check = object()  # secret "key" to pass to wrap_untrustworthy to extract feval_count
final = object()  # secret "key" to pass to wrap_untrustworthy to extract results


class ExhaustedTrialsError(Exception):
    pass


def wrap_untrustworthy(
    objective, budget, *, raising=False, bounding=None, softplus=False, eps=0.0
):
    # also handles bounding now, so it may be used for other purposes as well. whoops.
    feval_count = 0
    best_so_far = None

    def _objective(x):
        nonlocal feval_count, best_so_far
        if x is check:
            return feval_count
        if x is final:
            assert best_so_far is not None
            fopt, xopt = best_so_far
            return fopt, xopt, feval_count
        if raising and feval_count >= budget:
            raise ExhaustedTrialsError()
        if bounding is not None:
            x = do_bounding(x, bounding)
        fx = objective(x)
        feval_count += 1
        if budget is None or feval_count <= budget:
            if best_so_far is None or fx < best_so_far[0]:
                best_so_far = (fx, x.copy())
        return scalar_softplus(fx) + eps if softplus else fx

    return _objective