from utils import wrap_untrustworthy, check, final, ExhaustedTrialsError
import numpy as np
import scipy.optimize as scopt


def make_scipy(method, *, jacobian=None, hessian=None):
    # does not include:
    # basinhopping, shgo
    assert method in (
        "Nelder-Mead",
        "Powell",
        "CG",
        "BFGS",
        # "Newton-CG",
        "L-BFGS-B",
        "TNC",
        "COBYLA",
        "SLSQP",
        "trust-constr",
        # "dogleg",
        # "trust-ncg",
        # "trust-exact",
        # "trust-krylov",
    ), method

    def f(objective, n_trials, n_dim, with_count):
        prng = np.random.default_rng()
        _objective = wrap_untrustworthy(
            objective, n_trials, raising=True, bounding="sine"
        )
        x0 = np.full(n_dim, 0.5)
        bounds = scopt.Bounds([0.0] * n_dim, [1.0] * n_dim)
        # jac = "cs" if method == "dogleg" else None  # doesn't work
        jac = "2-point" if jacobian is None else jacobian
        hess = "2-point" if hessian is None else hessian
        # Alternatively, objects implementing the HessianUpdateStrategy interface
        # can be used to approximate the Hessian. Available quasi-Newton methods
        # implementing this interface are: BFGS; SR1.
        tol = None  # 0.0
        # options = dict(maxfun=n_trials) if method == "TNC" else dict(maxiter=n_trials)

        if method in ("BFGS", "CG", "COBYLA", "SLSQP", "trust-constr"):
            options = dict(maxiter=n_trials)
        elif method in ("Nelder-Mead", "Powell"):
            options = dict(maxfev=n_trials, maxiter=n_trials)
        elif method in ("L-BFGS-B", "TNC"):
            options = dict(maxfun=n_trials, maxiter=n_trials)
        else:
            options = dict(maxfun=n_trials, maxfev=n_trials, maxiter=n_trials)

        # silence some warnings:
        if method in ("Nelder-Mead", "Powell"):
            jac = None
            hess = None
        if method in ("COBYLA",):
            jac = None
            hess = None
            bounds = None
        if method in ("BFGS", "CG"):
            hess = None
            bounds = None
        if method in ("L-BFGS-B", "SLSQP", "TNC", "trust-constr"):
            hess = None

        checks = []

        def check_evals():
            evals = _objective(check)
            checks.append(evals)
            return evals < n_trials

        first_try = True
        while check_evals():
            if not first_try:
                x0 = prng.uniform(size=n_dim)

            try:
                res = scopt.minimize(
                    _objective,
                    x0,
                    method=method,
                    jac=jac,
                    hess=hess,
                    bounds=bounds,
                    tol=tol,
                    options=options,
                )
            except ExhaustedTrialsError:
                break
            else:
                # well, this is pointless.
                fopt, xopt, feval_count = res.fun, res.x, res.nfev
                # print("success:", res.success)
                # if not res.success and res.nfev < n_trials // 2:
                shut_up = (
                    method == "SLSQP"
                    and res.message == "Inequality constraints incompatible"
                ) or (
                    res.message
                    == "The maximum number of function evaluations is exceeded."
                )
                if not shut_up and not res.success and res.nfev < n_trials // 3:
                    print("", method, res, "", sep="\n")

            first_try = False

        # TODO: run without this, try to minimize number of attempts (i.e. list length)
        # if len(checks) >= 5: print(method, [b - a for a, b in zip(checks, checks[1:])])

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

    name = f"scipy_{method.replace('-', '').lower()}"
    if jacobian == "2-point":
        name += "_2j"
    elif jacobian == "3-point":
        name += "_3j"
    elif jacobian is not None:
        assert False, jacobian
    if hessian == "2-point":
        name += "_2h"
    elif hessian == "3-point":
        name += "_3h"
    elif hessian is not None:
        assert False, hessian
    f.__name__ = name + "_cube"
    return f


def scipy_basinhopping_cube(objective, n_trials, n_dim, with_count):
    progress = 1e-2  # TODO: make configurable?
    # NOTE: could also callbacks to extract solutions instead of wrapping objective functions?

    def accept_bounded(x_new=None, x_old=None, f_new=None, f_old=None):
        return np.all(x_new >= 0.0) and np.all(x_new <= 1.0)

    def dummy_minimizer(fun, x0, args, **options):
        return scopt.OptimizeResult(x=x0, fun=fun(x0), success=True, nfev=1)

    x0 = np.full(n_dim, 0.5)
    res = scopt.basinhopping(
        objective,
        x0,
        minimizer_kwargs=dict(method=dummy_minimizer),
        accept_test=accept_bounded,
        disp=False,
        niter=n_trials,
        # TODO: try without any progress vars at all.
        T=progress,
        stepsize=progress / 2,
    )
    fopt, xopt, feval_count = res.fun, res.x, res.nfev
    # print("success:", res.success)
    return (fopt, xopt, feval_count) if with_count else (fopt, xopt)


def scipy_direct_cube(objective, n_trials, n_dim, with_count):
    bounds = scopt.Bounds([0.0] * n_dim, [1.0] * n_dim)
    # TODO: try different values of eps. default 0.0001
    res = scopt.direct(
        objective,
        bounds=bounds,
        maxfun=n_trials,
        maxiter=1_000_000,
        vol_tol=0,
    )
    fopt, xopt, feval_count = res.fun, res.x, res.nfev
    # print("success:", res.success)
    return (fopt, xopt, feval_count) if with_count else (fopt, xopt)


def scipy_direct_l_cube(objective, n_trials, n_dim, with_count):
    bounds = scopt.Bounds([0.0] * n_dim, [1.0] * n_dim)
    # TODO: try different values of eps. default 0.0001
    res = scopt.direct(
        objective,
        bounds=bounds,
        maxfun=n_trials,
        maxiter=1_000_000,
        vol_tol=0,
        locally_biased=True,
        len_tol=0.0,  # only for locally_biased=True
    )
    fopt, xopt, feval_count = res.fun, res.x, res.nfev
    # print("success:", res.success)
    return (fopt, xopt, feval_count) if with_count else (fopt, xopt)


scipy_bfgs_2j_cube = make_scipy("BFGS", jacobian="2-point")
scipy_bfgs_3j_cube = make_scipy("BFGS", jacobian="3-point")
scipy_cg_2j_cube = make_scipy("CG", jacobian="2-point")
scipy_cg_3j_cube = make_scipy("CG", jacobian="3-point")
scipy_cobyla_cube = make_scipy("COBYLA")
# scipy_dogleg_cube = make_scipy("dogleg")  # ValueError: Jacobian is required for dogleg minimization
scipy_lbfgsb_2j_cube = make_scipy("L-BFGS-B", jacobian="2-point")
scipy_lbfgsb_3j_cube = make_scipy("L-BFGS-B", jacobian="3-point")
scipy_neldermead_cube = make_scipy("Nelder-Mead")
# scipy_newtoncg_cube = make_scipy("Newton-CG")  # ValueError: Jacobian is required for Newton-CG method
scipy_powell_cube = make_scipy("Powell")
scipy_slsqp_2j_cube = make_scipy("SLSQP", jacobian="2-point")
scipy_slsqp_3j_cube = make_scipy("SLSQP", jacobian="3-point")
scipy_tnc_2j_cube = make_scipy("TNC", jacobian="2-point")
scipy_tnc_3j_cube = make_scipy("TNC", jacobian="3-point")
scipy_trustconstr_2j_cube = make_scipy("trust-constr", jacobian="2-point")
scipy_trustconstr_3j_cube = make_scipy("trust-constr", jacobian="3-point")
# scipy_trustexact_2j_cube = make_scipy("trust-exact", jacobian="2-point")
# scipy_trustexact_3j_cube = make_scipy("trust-exact", jacobian="3-point")
# scipy_trustkrylov_cube = make_scipy("trust-krylov")  # ValueError: ('Jacobian is required for trust region ', 'exact minimization.')
# scipy_trustncg_cube = make_scipy("trust-ncg")  # ValueError: Jacobian is required for Newton-CG trust-region minimization