import logging
import os
import warnings
from functools import wraps

import numpy as np

# Optimize imported here and below as the jax-optimized one is jax or passthrough, but this is required regardless
import scipy.optimize
from pymbar.utils import ensure_type, check_w_normalized, ParameterError

logger = logging.getLogger(__name__)

use_jit = False


# Temporary until we can make a proper setting to enable/disable by choice at runtime
def _setup_jax_acceleration():
    return os.environ.get("PYMBAR_DISABLE_JAX", "").lower() in ("true", "yes", "1")


# Setup if we should use jax or not
force_no_jax = _setup_jax_acceleration()

try:
    #### JAX related imports
    if force_no_jax:
        # Capture user-disabled JAX instead "JAX not found"
        raise ImportError("Jax disabled by force_no_jax in mbar_solvers.py")
    try:
        # This works with pre- and post- JAX 0.3.25
        # If you use try...except around a from jax.config import config then
        # config gets mangled and the x64_enabled flag doesn't exist.
        from jax import config

        if not config.x64_enabled:
            # Warn that we're going to be setting 64 bit jax
            logger.warning(
                "\n"
                "****** PyMBAR will use 64-bit JAX! *******\n"
                "* JAX is currently set to 32-bit bitsize *\n"
                "* which is its default.                  *\n"
                "*                                        *\n"
                "* PyMBAR requires 64-bit mode and WILL   *\n"
                "* enable JAX's 64-bit mode when called.  *\n"
                "*                                        *\n"
                "* This MAY cause problems with other     *\n"
                "* Uses of JAX in the same code.          *\n"
                "******************************************\n"
            )

        from jax.numpy import exp, sum, newaxis, diag, dot, s_
        from jax.numpy import pad as npad
        from jax.numpy.linalg import lstsq
        import jax.scipy.optimize as optimize_maybe_jax
        from jax.scipy.special import logsumexp

        from jax import jit as jit_or_passthrough

        use_jit = True
    except ImportError:
        # Catch no JAX and throw a warning
        logger.warning(
            "\n"
            "********* JAX NOT FOUND *********\n"
            " PyMBAR can run faster with JAX  \n"
            " But will work fine without it   \n"
            "Either install with pip or conda:\n"
            "      pip install pymbar[jax]     \n"
            "               OR                \n"
            "      conda install pymbar       \n"
            "*********************************"
        )
        raise  # Continue with the raised Import Error

except ImportError:
    # No JAX found, overlap imports
    # These imports MUST align exactly
    from numpy import exp, sum, newaxis, diag, dot, s_
    from numpy import pad as npad
    from numpy.linalg import lstsq
    import scipy.optimize as optimize_maybe_jax  # pylint: disable=reimported
    from scipy.special import logsumexp

    # No jit, so make a passthrough decorator
    def jit_or_passthrough(fn):
        return fn


# Note on "pylint: disable=invalid-unary-operand-type"
# Known issue with astroid<2.12 and numpy array returns, but 2.12 doesn't fix it due to returns being jax.
# Can be mostly ignored

if use_jit is False:
    logger.info("JAX was either not detected or disabled, using standard NumPy and SciPy")
else:
    logger.info("JAX detected. Using JAX acceleration.")

# Below are the recommended default protocols (ordered sequence of minimization algorithms / NLE solvers) for solving
# the MBAR equations.
# Note: we use tuples instead of lists to avoid accidental mutability.
JAX_SOLVER_PROTOCOL = (
    dict(method="BFGS", continuation=True),
    dict(method="adaptive", options=dict(min_sc_iter=0)),
)

DEFAULT_SOLVER_PROTOCOL = (
    dict(method="hybr", continuation=True),
    dict(method="adaptive", options=dict(min_sc_iter=0)),
)

ROBUST_SOLVER_PROTOCOL = (
    dict(method="adaptive", options=dict(maxiter=1000)),
    dict(method="L-BFGS-B", options=dict(maxiter=1000)),
)

BOOTSTRAP_SOLVER_PROTOCOL = (dict(method="adaptive", options=dict(min_sc_iter=0)),)

# Allows all of the gradient based methods, but not the non-gradient methods ["Nelder-Mead", "Powell", "COBYLA"]",
scipy_minimize_options = [
    "L-BFGS-B",
    "dogleg",
    "CG",
    "BFGS",
    "Newton-CG",
    "TNC",
    "trust-ncg",
    "trust-krylov",
    "trust-exact",
    "SLSQP",
]
scipy_nohess_options = [
    "L-BFGS-B",
    "BFGS",
    "CG",
    "TNC",
    "SLSQP",
]  # don't pass a hessian to these to avoid warnings to these.
scipy_root_options = ["hybr", "lm"]  # only use root options with the hessian included


def jit_or_pass_after_bitsize(jitable_fn):
    """
    Attempt to set JAX precision if present. This does nothing if JAX is not present

    Parameters
    ----------
    jitable_fn: function
        A function which can be jit'd
    """

    @wraps(
        jitable_fn
    )  # Helper to ensure the decorated function still registers for docs and inspection
    def staggered_jit(*args, **kwargs):
        # This will only trigger if JAX is set
        if use_jit and not config.x64_enabled:
            # Warn that JAX 64-bit will being turned on
            logger.warning(
                "\n"
                "******* JAX 64-bit mode is now on! *******\n"
                "*     JAX is now set to 64-bit mode!     *\n"
                "*   This MAY cause problems with other   *\n"
                "*      uses of JAX in the same code.     *\n"
                "******************************************\n"
            )
            config.update("jax_enable_x64", True)
        jited_fn = jit_or_passthrough(jitable_fn)
        return jited_fn(*args, **kwargs)

    return staggered_jit


def validate_inputs(u_kn, N_k, f_k):
    """Check types and return inputs for MBAR calculations.

    Parameters
    ----------
    u_kn or q_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
        The reduced potential energies or unnormalized probabilities
    N_k : np.ndarray, shape=(n_states), dtype='int'
        The number of samples in each state
    f_k : np.ndarray, shape=(n_states), dtype='float'
        The reduced free energies of each state

    Returns
    -------
    u_kn or q_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
        The reduced potential energies or unnormalized probabilities
    N_k : np.ndarray, shape=(n_states), dtype='float'
        The number of samples in each state.  Converted to float because this cast is required when log is calculated.
    f_k : np.ndarray, shape=(n_states), dtype='float'
        The reduced free energies of each state
    """
    n_states, n_samples = u_kn.shape

    u_kn = ensure_type(u_kn, "float", 2, "u_kn or Q_kn", shape=(n_states, n_samples))
    N_k = ensure_type(
        N_k, "float", 1, "N_k", shape=(n_states,), warn_on_cast=False
    )  # Autocast to float because will be eventually used in float calculations.
    f_k = ensure_type(f_k, "float", 1, "f_k", shape=(n_states,))

    return u_kn, N_k, f_k


def self_consistent_update(u_kn, N_k, f_k, states_with_samples=None):
    """Return an improved guess for the dimensionless free energies

    Parameters
    ----------
    u_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
        The reduced potential energies, i.e. -log unnormalized probabilities
    N_k : np.ndarray, shape=(n_states), dtype='int'
        The number of samples in each state
    f_k : np.ndarray, shape=(n_states), dtype='float'
        The reduced free energies of each state

    Returns
    -------
    f_k : np.ndarray, shape=(n_states), dtype='float'
        Updated estimate of f_k

    Notes
    -----
    Equation C3 in MBAR JCP paper.
    """

    return jax_self_consistent_update(u_kn, N_k, f_k, states_with_samples=states_with_samples)


@jit_or_pass_after_bitsize
def _jit_self_consistent_update(u_kn, N_k, f_k):
    """JAX version of self_consistent update.  For parameters, see self_consistent_update.
    N_k must be float (should be cast at a higher level)

    """
    # Asteroid
    log_denominator_n = logsumexp(f_k - u_kn.T, b=N_k, axis=1)
    # All states can contribute to the numerator term. Check transpose
    return -1.0 * logsumexp(
        -log_denominator_n - u_kn, axis=1
    )  # pylint: disable=invalid-unary-operand-type


def jax_self_consistent_update(u_kn, N_k, f_k, states_with_samples=None):
    """JAX version of self_consistent update.  For parameters, see self_consistent_update.
    N_k must be float (should be cast at a higher level)

    """
    # Only the states with samples can contribute to the denominator term.
    # Precondition before feeding the op to the JIT'd function
    # In theory, this can be computed with jax.lax.cond, but trying to reuse code for non-jax paths
    states_with_samples = s_[:] if states_with_samples is None else states_with_samples
    # Feed to the JIT'd function. Can't pass slice types, so slice here
    return _jit_self_consistent_update(
        u_kn[states_with_samples], N_k[states_with_samples], f_k[states_with_samples]
    )


def mbar_gradient(u_kn, N_k, f_k):
    """Gradient of MBAR objective function.

    Parameters
    ----------
    u_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
        The reduced potential energies, i.e. -log unnormalized probabilities
    N_k : np.ndarray, shape=(n_states), dtype='int'
        The number of samples in each state
    f_k : np.ndarray, shape=(n_states), dtype='float'
        The reduced free energies of each state

    Returns
    -------
    grad : np.ndarray, dtype=float, shape=(n_states)
        Gradient of mbar_objective

    Notes
    -----
    This is equation C6 in the JCP MBAR paper.
    """
    return jax_mbar_gradient(u_kn, N_k, f_k)


@jit_or_pass_after_bitsize
def jax_mbar_gradient(u_kn, N_k, f_k):
    """JAX version of MBAR gradient function. See documentation of mbar_gradient.
    N_k must be float (should be cast at a higher level)
    """

    log_denominator_n = logsumexp(f_k - u_kn.T, b=N_k, axis=1)
    log_numerator_k = logsumexp(-log_denominator_n - u_kn, axis=1)
    return -1 * N_k * (1.0 - exp(f_k + log_numerator_k))


def mbar_objective(u_kn, N_k, f_k):
    """Calculates objective function for MBAR.

    Parameters
    ----------
    u_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
        The reduced potential energies, i.e. -log unnormalized probabilities
    N_k : np.ndarray, shape=(n_states), dtype='int'
        The number of samples in each state
    f_k : np.ndarray, shape=(n_states), dtype='float'
        The reduced free energies of each state


    Returns
    -------
    obj : float
        Objective function

    Notes
    -----
    This objective function is essentially a doubly-summed partition function and is
    quite sensitive to precision loss from both overflow and underflow. For optimal
    results, u_kn can be preconditioned by subtracting out a `n` dependent
    vector.

    More optimal precision, the objective function uses math.fsum for the
    outermost sum and logsumexp for the inner sum.
    """

    return jax_mbar_objective(u_kn, N_k, f_k)


@jit_or_pass_after_bitsize
def jax_mbar_objective(u_kn, N_k, f_k):
    """JAX version of mbar_objective.
    For parameters, mbar_objective_and_Gradient
    N_k must be float (should be cast at a higher level)

    """

    log_denominator_n = logsumexp(f_k - u_kn.T, b=N_k, axis=1)
    obj = sum(log_denominator_n) - dot(N_k, f_k)

    return obj


@jit_or_pass_after_bitsize
def jax_mbar_objective_and_gradient(u_kn, N_k, f_k):
    """JAX version of mbar_objective_and_gradient.
    For parameters, mbar_objective_and_Gradient
    N_k must be float (should be cast at a higher level)

    """

    log_denominator_n = logsumexp(f_k - u_kn.T, b=N_k, axis=1)
    log_numerator_k = logsumexp(-log_denominator_n - u_kn, axis=1)
    grad = -1 * N_k * (1.0 - exp(f_k + log_numerator_k))

    obj = sum(log_denominator_n) - dot(N_k, f_k)

    return obj, grad


def mbar_objective_and_gradient(u_kn, N_k, f_k):
    """Calculates both objective function and gradient for MBAR.

    Parameters
    ----------
    u_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
        The reduced potential energies, i.e. -log unnormalized probabilities
    N_k : np.ndarray, shape=(n_states), dtype='int'
        The number of samples in each state
    f_k : np.ndarray, shape=(n_states), dtype='float'
        The reduced free energies of each state


    Returns
    -------
    obj : float
        Objective function
    grad : np.ndarray, dtype=float, shape=(n_states)
        Gradient of objective function

    Notes
    -----
    This objective function is essentially a doubly-summed partition function and is
    quite sensitive to precision loss from both overflow and underflow. For optimal
    results, u_kn can be preconditioned by subtracting out a `n` dependent
    vector.

    More optimal precision, the objective function uses math.fsum for the
    outermost sum and logsumexp for the inner sum.

    The gradient is equation C6 in the JCP MBAR paper; the objective
    function is its integral.
    """

    return jax_mbar_objective_and_gradient(u_kn, N_k, f_k)


@jit_or_pass_after_bitsize
def jax_mbar_hessian(u_kn, N_k, f_k):
    """JAX version of mbar_hessian.
    For parameters, see mbar_hessian
    N_k must be float (should be cast at a higher level)

    """

    log_denominator_n = logsumexp(f_k - u_kn.T, b=N_k, axis=1)
    logW = f_k - u_kn.T - log_denominator_n[:, newaxis]
    W = exp(logW)

    H = dot(W.T, W)
    H *= N_k
    H *= N_k[:, newaxis]
    H -= diag(W.sum(0) * N_k)
    return -1.0 * H


def mbar_hessian(u_kn, N_k, f_k):
    """Hessian of MBAR objective function.

    Parameters
    ----------
    u_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
        The reduced potential energies, i.e. -log unnormalized probabilities
    N_k : np.ndarray, shape=(n_states), dtype='int'
        The number of samples in each state
    f_k : np.ndarray, shape=(n_states), dtype='float'
        The reduced free energies of each state

    Returns
    -------
    H : np.ndarray, dtype=float, shape=(n_states, n_states)
        Hessian of mbar objective function.

    Notes
    -----
    Equation (C9) in JCP MBAR paper.
    """

    return jax_mbar_hessian(u_kn, N_k, f_k)


@jit_or_pass_after_bitsize
def jax_mbar_log_W_nk(u_kn, N_k, f_k):
    """JAX version of mbar_log_W_nk.
    For parameters, see mbar_log_W_nk
    N_k must be float (should be cast at a higher level)

    """

    log_denominator_n = logsumexp(f_k - u_kn.T, b=N_k, axis=1)
    logW = f_k - u_kn.T - log_denominator_n[:, newaxis]
    return logW


def mbar_log_W_nk(u_kn, N_k, f_k):
    """Calculate the log weight matrix.

    Parameters
    ----------
    u_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
        The reduced potential energies, i.e. -log unnormalized probabilities
    N_k : np.ndarray, shape=(n_states), dtype='int'
        The number of samples in each state
    f_k : np.ndarray, shape=(n_states), dtype='float'
        The reduced free energies of each state

    Returns
    -------
    logW_nk : np.ndarray, dtype='float', shape=(n_samples, n_states)
        The normalized log weights.

    Notes
    -----
    Equation (9) in JCP MBAR paper.
    """
    return jax_mbar_log_W_nk(u_kn, N_k, f_k)


@jit_or_pass_after_bitsize
def jax_mbar_W_nk(u_kn, N_k, f_k):
    """JAX version of mbar_W_nk.
    For parameters, see mbar_W_nk
    N_k must be float (should be cast at a higher level)

    """
    return exp(jax_mbar_log_W_nk(u_kn, N_k, f_k))


def mbar_W_nk(u_kn, N_k, f_k):
    """Calculate the weight matrix.

    Parameters
    ----------
    u_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
        The reduced potential energies, i.e. -log unnormalized probabilities
    N_k : np.ndarray, shape=(n_states), dtype='int'
        The number of samples in each state
    f_k : np.ndarray, shape=(n_states), dtype='float'
        The reduced free energies of each state

    Returns
    -------
    W_nk : np.ndarray, dtype='float', shape=(n_samples, n_states)
        The normalized weights.

    Notes
    -----
    Equation (9) in JCP MBAR paper.
    """
    return jax_mbar_W_nk(u_kn, N_k, f_k)


def adaptive(u_kn, N_k, f_k, tol=1.0e-8, options=None):
    """
    Determine dimensionless free energies by a combination of Newton-Raphson iteration and self-consistent iteration.
    Picks whichever method gives the lowest gradient.
    Is slower than NR since it calculates the log norms twice each iteration.

    OPTIONAL ARGUMENTS
    tol (float between 0 and 1) - relative tolerance for convergence (default 1.0e-12)

    options : dictionary of options
        gamma (float between 0 and 1) - incrementor for NR iterations (default 1.0).  Usually not changed now, since adaptively switch.
        maxiter (int) - maximum number of Newton-Raphson iterations (default 10000: either NR converges or doesn't, pretty quickly)
        verbose (boolean) - verbosity level for debug output

    NOTES

    This method determines the dimensionless free energies by
    minimizing a convex function whose solution is the desired
    estimator.  The original idea came from the construction of a
    likelihood function that independently reproduced the work of
    Geyer (see [1] and Section 6 of [2]).  This can alternatively be
    formulated as a root-finding algorithm for the Z-estimator.  More
    details of this procedure will follow in a subsequent paper.  Only
    those states with nonzero counts are include in the estimation
    procedure.

    REFERENCES
    See Appendix C.2 of [1].

    """
    # put the defaults here in case we get passed an 'options' dictionary that is only partial
    options.setdefault("verbose", False)
    options.setdefault("maxiter", 10000)
    options.setdefault("print_warning", False)
    options.setdefault("gamma", 1.0)
    options.setdefault("min_sc_iter", 2)  # set a minimum number of self-consistent iterations

    gamma = options["gamma"]

    doneIterating = False
    if options["verbose"] == True:
        logger.info(
            "Determining dimensionless free energies by Newton-Raphson / self-consistent iteration."
        )

    if tol < 4.0 * np.finfo(float).eps:
        logger.info("Tolerance may be too close to machine precision to converge.")

    success = False  # fail unless solution is found.
    # keep track of Newton-Raphson and self-consistent iterations
    nr_iter = 0
    sci_iter = 0

    f_sci = np.zeros(len(f_k), dtype=np.float64)
    f_nr = np.zeros(len(f_k), dtype=np.float64)

    # Perform Newton-Raphson iterations (with sci computed on the way)

    # usually calculated at the end of the loop and saved, but we need
    # to calculate the first time.
    g = mbar_gradient(u_kn, N_k, f_k)  # Objective function gradient.

    maxiter = options["maxiter"]
    min_sc_iter = options["min_sc_iter"]
    warn = "Did not converge."
    for iteration in range(0, maxiter):
        if use_jit:
            (f_sci, g_sci, gnorm_sci, f_nr, g_nr, gnorm_nr) = jax_core_adaptive(
                u_kn, N_k, f_k, options["gamma"]
            )
        else:
            H = mbar_hessian(u_kn, N_k, f_k)  # Objective function hessian
            Hinvg = np.linalg.lstsq(H, g, rcond=-1)[0]
            Hinvg -= Hinvg[0]
            f_nr = f_k - gamma * Hinvg

            # self-consistent iteration gradient norm and saved log sums.
            f_sci = self_consistent_update(u_kn, N_k, f_k)
            f_sci = f_sci - f_sci[0]  # zero out the minimum
            g_sci = mbar_gradient(u_kn, N_k, f_sci)
            gnorm_sci = dot(g_sci, g_sci)

            # newton raphson gradient norm and saved log sums.
            g_nr = mbar_gradient(u_kn, N_k, f_nr)
            gnorm_nr = dot(g_nr, g_nr)

        # we could save the gradient, for the next round, but it's not too expensive to
        # compute since we are doing the Hessian anyway.

        if options["verbose"]:
            logger.info(
                "self consistent iteration gradient norm is %10.5g, Newton-Raphson gradient norm is %10.5g"
                % (np.sqrt(gnorm_sci), np.sqrt(gnorm_nr))
            )
        # decide which directon to go depending on size of gradient norm
        f_old = f_k

        if gnorm_sci < gnorm_nr or sci_iter < min_sc_iter:
            f_k = f_sci
            g = g_sci
            sci_iter += 1
            if options["verbose"]:
                if sci_iter < min_sc_iter:
                    logger.info(
                        f"Choosing self-consistent iteration on iteration {iteration:d} because min_sci_iter={min_sc_iter:d}"
                    )
                else:
                    logger.info(
                        f"Choosing self-consistent iteration for lower gradient on iteration {iteration:d}"
                    )
        else:
            f_k = f_nr
            g = g_nr
            nr_iter += 1
            if options["verbose"]:
                logger.info(f"Newton-Raphson used on iteration {iteration:}")

        div = np.abs(f_k[1:])  # what we will divide by to get relative difference
        zeroed = np.abs(f_k[1:]) < np.min(
            [10**-8, tol]
        )  # check which values are near enough to zero, hard coded max for now.
        div[zeroed] = 1.0  # for these values, use absolute values.
        max_delta = np.max(np.abs(f_k[1:] - f_old[1:]) / div)
        max_diff = np.max(np.abs(f_sci[1:] - f_nr[1:]) / div)
        # add this just to make sure they are not too different.
        # if we start with bad states, the f_k - f_k_old might be far off.
        if np.isnan(max_delta) or ((max_delta < tol) and max_diff < np.sqrt(tol)):
            doneIterating = True
            success = True
            warn = "Convergence achieved by change in f with respect to previous guess."
            break

    if doneIterating:
        if options["verbose"]:
            logger.info(f"Converged to tolerance of {max_delta:e} in {iteration+1:d} iterations.")
            logger.info(
                f"Of {iteration+1:d} iterations, {nr_iter:d} were Newton-Raphson iterations and {sci_iter:d} were self-consistent iterations"
            )
            if np.all(f_k == 0.0):
                logger.info("WARNING: All f_k appear to be zero.")
    else:
        logger.warning("WARNING: Did not converge to within specified tolerance.")

        if maxiter <= 0:
            logger.warning(
                f"No iterations ran be cause maximum_iterations was <= 0 ({maxiter:s})!"
            )
        else:
            logger.warning(
                f"max_delta = {max_delta:e}, tol = {tol:e}, maximum_iterations = {maxiter:d}, iterations completed = {iteration:d}"
            )

    results = dict()
    results["success"] = success
    results["message"] = warn
    results["x"] = f_k

    return results


@jit_or_pass_after_bitsize
def jax_core_adaptive(u_kn, N_k, f_k, gamma):
    """JAX version of adaptive inner loop.
    N_k must be float (should be cast at a higher level)

    """

    # Perform Newton-Raphson iterations (with sci computed on the way)
    g = mbar_gradient(u_kn, N_k, f_k)  # Objective function gradient
    H = mbar_hessian(u_kn, N_k, f_k)  # Objective function hessian
    Hinvg = lstsq(H, g, rcond=-1)[0]
    Hinvg -= Hinvg[0]
    f_nr = f_k - gamma * Hinvg

    # self-consistent iteration gradient norm and saved log sums.
    f_sci = self_consistent_update(u_kn, N_k, f_k)
    f_sci = f_sci - f_sci[0]  # zero out the minimum
    g_sci = mbar_gradient(u_kn, N_k, f_sci)
    gnorm_sci = dot(g_sci, g_sci)

    # newton raphson gradient norm and saved log sums.
    g_nr = mbar_gradient(u_kn, N_k, f_nr)
    gnorm_nr = dot(g_nr, g_nr)

    return f_sci, g_sci, gnorm_sci, f_nr, g_nr, gnorm_nr


@jit_or_pass_after_bitsize
def jax_precondition_u_kn(u_kn, N_k, f_k):
    """JAX version of precondition_u_kn
    for parameters, see precondition_u_kn
    N_k must be float (should be cast at a higher level)

    """

    u_kn = u_kn - u_kn.min(0)
    u_kn += (logsumexp(f_k - u_kn.T, b=N_k, axis=1)) - dot(N_k, f_k) / N_k.sum()
    return u_kn


def precondition_u_kn(u_kn, N_k, f_k):
    """Subtract a sample-dependent constant from u_kn to improve precision

    Parameters
    ----------
    u_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
        The reduced potential energies, i.e. -log unnormalized probabilities
    N_k : np.ndarray, shape=(n_states), dtype='int'
        The number of samples in each state
    f_k : np.ndarray, shape=(n_states), dtype='float'
        The reduced free energies of each state

    Returns
    -------
    u_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
        The reduced potential energies, i.e. -log unnormalized probabilities

    Notes
    -----
    Returns u_kn - x_n, where x_n is based on the current estimate of f_k.
    Upon subtraction of x_n, the MBAR objective function changes by an
    additive constant, but its derivatives remain unchanged.  We choose
    x_n such that the current objective function value is zero, which
    should give maximum precision in the objective function.
    """
    return jax_precondition_u_kn(u_kn, N_k, f_k)


def solve_mbar_once(
    u_kn_nonzero,
    N_k_nonzero,
    f_k_nonzero,
    method="adaptive",
    tol=1e-12,
    continuation=None,
    options=None,
):
    """Solve MBAR self-consistent equations using some form of equation solver.

    Parameters
    ----------
    u_kn_nonzero : np.ndarray, shape=(n_states, n_samples), dtype='float'
        The reduced potential energies, i.e. -log unnormalized probabilities
        for the nonempty states
    N_k_nonzero : np.ndarray, shape=(n_states), dtype='int'
        The number of samples in each state for the nonempty states
    f_k_nonzero : np.ndarray, shape=(n_states), dtype='float'
        The reduced free energies for the nonempty states
    method : str, optional, default="hybr"
        The optimization routine to use.  This can be any of the methods
        available via scipy.optimize.minimize() or scipy.optimize.root().
    tol : float, optional, default=1E-14
        The convergance tolerance for minimize() or root()
    verbose: bool
        Whether to print information about the solution method.
    options: dict, optional, default=None
        Optional dictionary of algorithm-specific parameters.  See
        scipy.optimize.root or scipy.optimize.minimize for details.

    Returns
    -------
    f_k : np.ndarray
        The converged reduced free energies.
    results : dict
        Dictionary containing entire results of optimization routine, may
        be useful when debugging convergence.

    Notes
    -----
    This function requires that N_k_nonzero > 0--that is, you should have
    already dropped all the states for which you have no samples.
    Internally, this function works in a reduced coordinate system defined
    by subtracting off the first component of f_k and fixing that component
    to be zero.

    For fast but precise convergence, we recommend calling this function
    multiple times to polish the result.  `solve_mbar()` facilitates this.
    """

    # we only validate at the outside of the call
    u_kn_nonzero, N_k_nonzeo, f_k_nonzero = validate_inputs(u_kn_nonzero, N_k_nonzero, f_k_nonzero)
    f_k_nonzero = f_k_nonzero - f_k_nonzero[0]  # Work with reduced dimensions with f_k[0] := 0
    N_k_nonzero = 1.0 * N_k_nonzero  # convert to float for acceleration.
    u_kn_nonzero = precondition_u_kn(u_kn_nonzero, N_k_nonzero, f_k_nonzero)

    pad = lambda x: np.pad(
        x, (1, 0), mode="constant"
    )  # Helper function inserts zero before first element
    unpad_second_arg = lambda obj, grad: (
        obj,
        grad[1:],
    )  # Helper function drops first element of gradient

    # Create objective functions / nonlinear equations to send to scipy.optimize, fixing f_0 = 0
    grad = lambda x: mbar_gradient(u_kn_nonzero, N_k_nonzero, pad(x))[
        1:
    ]  # Objective function gradient

    grad_and_obj = lambda x: unpad_second_arg(
        *mbar_objective_and_gradient(u_kn_nonzero, N_k_nonzero, pad(x))
    )  # Objective function gradient and objective function

    de_jax_grad_and_obj = lambda x: (
        *map(np.array, grad_and_obj(x)),  # (...,) Casts to tuple instead of <map> object
    )  # Force any jax-based array output to normal numpy for scipy.optimize.minimize. np.asarray does not work.

    hess = lambda x: mbar_hessian(u_kn_nonzero, N_k_nonzero, pad(x))[1:][
        :, 1:
    ]  # Hessian of objective function
    with warnings.catch_warnings(record=True) as w:
        if use_jit and method == "BFGS":
            fpad = lambda x: npad(x, (1, 0))
            obj = lambda x: mbar_objective(u_kn_nonzero, N_k_nonzero, fpad(x))
            # objective function to be minimized (for derivative free methods, mostly jit)
            jax_results = optimize_maybe_jax.minimize(
                obj,
                f_k_nonzero[1:],
                method=method,
                tol=tol,
                options=dict(maxiter=options["maxiter"]),
            )
            results = dict()  # there should be a way to copy this.
            results["x"] = jax_results[0]
            f_k_nonzero = pad(results["x"])
            results["success"] = jax_results[1]
        elif method in scipy_minimize_options:
            if method in scipy_nohess_options:
                hess = None  # To suppress warning from passing a hessian function.
            results = scipy.optimize.minimize(
                de_jax_grad_and_obj,
                f_k_nonzero[1:],
                jac=True,
                hess=hess,
                method=method,
                tol=tol,
                options=options,
            )
            f_k_nonzero = pad(results["x"])
        elif method == "adaptive":
            results = adaptive(u_kn_nonzero, N_k_nonzero, f_k_nonzero, tol=tol, options=options)
            f_k_nonzero = results["x"]
        elif method in scipy_root_options:
            # find the root in the gradient.
            results = scipy.optimize.root(
                grad, f_k_nonzero[1:], jac=hess, method=method, tol=tol, options=options
            )
            f_k_nonzero = pad(results["x"])
        else:
            raise ParameterError(f"Method {method} for solution of free energies not recognized")

    # If there were runtime warnings, show the messages
    if len(w) > 0:
        can_ignore = True
        for warn_msg in w:
            if "Unknown solver options" in str(warn_msg.message):
                continue
            warnings.showwarning(
                warn_msg.message,
                warn_msg.category,
                warn_msg.filename,
                warn_msg.lineno,
                warn_msg.file,
                "",
            )
            can_ignore = False  # If any warning is not just unknown options, can not skip check
        if not can_ignore:
            # Ensure MBAR solved correctly
            w_nk_check = mbar_W_nk(u_kn_nonzero, N_k_nonzero, f_k_nonzero)
            check_w_normalized(w_nk_check, N_k_nonzero)
            logger.warning(
                "MBAR weights converged within tolerance, despite the SciPy Warnings. Please validate your results."
            )

    return f_k_nonzero, results


def solve_mbar(u_kn_nonzero, N_k_nonzero, f_k_nonzero, solver_protocol=None):
    """Solve MBAR self-consistent equations using some sequence of equation solvers.

    Parameters
    ----------
    u_kn_nonzero : np.ndarray, shape=(n_states, n_samples), dtype='float'
        The reduced potential energies, i.e. -log unnormalized probabilities
        for the nonempty states
    N_k_nonzero : np.ndarray, shape=(n_states), dtype='int'
        The number of samples in each state for the nonempty states
    f_k_nonzero : np.ndarray, shape=(n_states), dtype='float'
        The reduced free energies for the nonempty states
    solver_protocol : tuple(dict()), optional, default=None
        Optional list of dictionaries of steps in solver protocol.
        If None, a default protocol will be used.

    Returns
    -------
    f_k : np.ndarray
        The converged reduced free energies.
    all_results : list(dict())
        List of results from each step of solver_protocol.  Each element in
        list contains the results dictionary from solve_mbar_once()
        for the corresponding step.

    Notes
    -----
    This function requires that N_k_nonzero > 0--that is, you should have
    already dropped all the states for which you have no samples.
    Internally, this function works in a reduced coordinate system defined
    by subtracting off the first component of f_k and fixing that component
    to be zero.

    This function calls `solve_mbar_once()` multiple times to achieve
    converged results.  Generally, a single call to solve_mbar_once()
    will not give fully converged answers because of limited numerical precision.
    Each call to `solve_mbar_once()` re-conditions the nonlinear
    equations using the current guess.
    """

    if solver_protocol is None:
        solver_protocol = DEFAULT_SOLVER_PROTOCOL

    all_fks = []
    all_gnorms = []
    all_results = []

    for solver in solver_protocol:
        f_k_nonzero_result, results = solve_mbar_once(
            u_kn_nonzero, N_k_nonzero, f_k_nonzero, **solver
        )
        all_fks.append(f_k_nonzero_result)
        all_gnorms.append(
            np.linalg.norm(mbar_gradient(u_kn_nonzero, N_k_nonzero, f_k_nonzero_result))
        )
        all_results.append(results)

        if results["success"]:
            success = True
            best_gnorm = all_gnorms[-1]
            logger.info(f"Reached a solution to within tolerance with {solver['method']}")
            break
        else:
            logger.warning(
                f"Failed to reach a solution to within tolerance with {solver['method']}: trying next method"
            )
        logger.info(f"Ending gnorm of method {solver['method']} = {all_gnorms[-1]:e}")
        if solver["continuation"]:
            f_k_nonzero = f_k_nonzero_result
            logger.info("Will continue with results from previous method")

    if results["success"]:
        logger.info("Solution found within tolerance!")
    else:
        i_best_gnorm = np.argmin(all_gnorms)
        logger.warning("No solution found to within tolerance.")
        best_method = solver_protocol[i_best_gnorm]["method"]
        best_gnorm = all_gnorms[i_best_gnorm]
        logger.warning(
            f"The solution with the smallest gradient {best_gnorm:e} norm is {best_method}"
        )
        f_k_nonzero_result = all_fks[i_best_gnorm]
        logger.warning(
            "Please exercise caution with this solution and consider alternative methods or a different tolerance."
        )

    logger.info(f"Final gradient norm: {best_gnorm:.3g}")

    return f_k_nonzero_result, all_results


def solve_mbar_for_all_states(u_kn, N_k, f_k, states_with_samples, solver_protocol):
    """Solve for free energies of states with samples, then calculate for
    empty states.

    Parameters
    ----------
    u_kn : np.ndarray, shape=(n_states, n_samples), dtype='float'
        The reduced potential energies, i.e. -log unnormalized probabilities
    N_k : np.ndarray, shape=(n_states), dtype='int'
        The number of samples in each state
    f_k : np.ndarray, shape=(n_states), dtype='float'
        The reduced free energies of each state
    solver_protocol : tuple(dict()), optional, default=None
        Sequence of dictionaries of steps in solver protocol for final
        stage of refinement.

    Returns
    -------
    f_k : np.ndarray, shape=(n_states), dtype='float'
        The free energies of states
    """

    if len(states_with_samples) == 1:
        f_k_nonzero = np.array([0.0])
    else:
        f_k_nonzero, all_results = solve_mbar(
            u_kn[states_with_samples],
            N_k[states_with_samples],
            f_k[states_with_samples],
            solver_protocol=solver_protocol,
        )

    f_k[states_with_samples] = np.array(f_k_nonzero)

    # Update all free energies because those from states with zero samples are not correctly computed by solvers.
    f_k = self_consistent_update(u_kn, N_k, f_k)
    # This is necessary because state 0 might have had zero samples,
    # but we still want that state to be the reference with free energy 0.
    f_k -= f_k[0]

    return f_k
