"""A unified interface for performing and debugging optimization problems."""

import abc
import warnings

import numpy as np
import scipy.optimize as opt
from scipy.optimize import minimize

from dipy.testing.decorators import warning_for_keywords
from dipy.utils.optpkg import optional_package

cvxpy, have_cvxpy, _ = optional_package("cvxpy", min_version="1.4.1")


class Optimizer:
    @warning_for_keywords()
    def __init__(
        self,
        fun,
        x0,
        args=(),
        *,
        method="L-BFGS-B",
        jac=None,
        hess=None,
        hessp=None,
        bounds=None,
        constraints=(),
        tol=None,
        callback=None,
        options=None,
        evolution=False,
    ):
        """A class for handling minimization of scalar function of one or more
        variables.

        Parameters
        ----------
        fun : callable
            Objective function.

        x0 : ndarray
            Initial guess.

        args : tuple, optional
            Extra arguments passed to the objective function and its
            derivatives (Jacobian, Hessian).

        method : str, optional
            Type of solver.  Should be one of

                - 'Nelder-Mead'
                - 'Powell'
                - 'CG'
                - 'BFGS'
                - 'Newton-CG'
                - 'Anneal'
                - 'L-BFGS-B'
                - 'TNC'
                - 'COBYLA'
                - 'SLSQP'
                - 'dogleg'
                - 'trust-ncg'

        jac : bool or callable, optional
            Jacobian of objective function. Only for CG, BFGS, Newton-CG,
            dogleg, trust-ncg.
            If `jac` is a Boolean and is True, `fun` is assumed to return the
            value of Jacobian along with the objective function. If False, the
            Jacobian will be estimated numerically.
            `jac` can also be a callable returning the Jacobian of the
            objective. In this case, it must accept the same arguments
            as `fun`.

        hess, hessp : callable, optional
            Hessian of objective function or Hessian of objective function
            times an arbitrary vector p.  Only for Newton-CG,
            dogleg, trust-ncg.
            Only one of `hessp` or `hess` needs to be given.  If `hess` is
            provided, then `hessp` will be ignored.  If neither `hess` nor
            `hessp` is provided, then the hessian product will be approximated
            using finite differences on `jac`. `hessp` must compute the Hessian
            times an arbitrary vector.

        bounds : sequence, optional
            Bounds for variables (only for L-BFGS-B, TNC and SLSQP).
            ``(min, max)`` pairs for each element in ``x``, defining
            the bounds on that parameter. Use None for one of ``min`` or
            ``max`` when there is no bound in that direction.

        constraints : dict or sequence of dict, optional
            Constraints definition (only for COBYLA and SLSQP).
            Each constraint is defined in a dictionary with fields:

                type : str
                    Constraint type: 'eq' for equality, 'ineq' for inequality.
                fun : callable
                    The function defining the constraint.
                jac : callable, optional
                    The Jacobian of `fun` (only for SLSQP).
                args : sequence, optional
                    Extra arguments to be passed to the function and Jacobian.

            Equality constraint means that the constraint function result is to
            be zero whereas inequality means that it is to be non-negative.
            Note that COBYLA only supports inequality constraints.

        tol : float, optional
            Tolerance for termination. For detailed control, use
            solver-specific options.

        callback : callable, optional
            Called after each iteration, as ``callback(xk)``, where ``xk`` is
            the current parameter vector. Only available using Scipy >= 0.12.

        options : dict, optional
            A dictionary of solver options. All methods accept the following
            generic options:

                maxiter : int
                    Maximum number of iterations to perform.
                disp : bool
                    Set to True to print convergence messages.

            For method-specific options, see
            `show_options('minimize', method)`.

        evolution : bool, optional
            save history of x for each iteration. Only available using Scipy
            >= 0.12.

        See Also
        --------
        scipy.optimize.minimize

        """
        self.size_of_x = len(x0)
        self._evol_kx = None

        if evolution is True:
            self._evol_kx = []

            def history_of_x(kx):
                self._evol_kx.append(kx)

            res = minimize(
                fun,
                x0,
                args,
                method,
                jac,
                hess,
                hessp,
                bounds,
                constraints,
                tol,
                callback=history_of_x,
                options=options,
            )

        else:
            res = minimize(
                fun,
                x0,
                args,
                method,
                jac,
                hess,
                hessp,
                bounds,
                constraints,
                tol,
                callback,
                options,
            )

        self.res = res

    @property
    def xopt(self):
        return self.res["x"]

    @property
    def fopt(self):
        return self.res["fun"]

    @property
    def nit(self):
        return self.res["nit"]

    @property
    def nfev(self):
        return self.res["nfev"]

    @property
    def message(self):
        return self.res["message"]

    def print_summary(self):
        print(self.res)

    @property
    def evolution(self):
        if self._evol_kx is not None:
            return np.asarray(self._evol_kx)
        else:
            return None


def spdot(A, B):
    """The same as np.dot(A, B), except it works even if A or B or both
    are sparse matrices.

    Parameters
    ----------
    A, B : arrays of shape (m, n), (n, k)

    Returns
    -------
    The matrix product AB. If both A and B are sparse, the result will be a
    sparse matrix. Otherwise, a dense result is returned

    See discussion here:
    http://mail.scipy.org/pipermail/scipy-user/2010-November/027700.html

    """
    return A @ B


@warning_for_keywords()
def sparse_nnls(
    y,
    X,
    *,
    momentum=1,
    step_size=0.01,
    non_neg=True,
    check_error_iter=10,
    max_error_checks=10,
    converge_on_sse=0.99,
):
    """
    Solve y=Xh for h, using gradient descent, with X a sparse matrix.

    Parameters
    ----------
    y : 1-d array of shape (N)
        The data. Needs to be dense.

    X : ndarray. May be either sparse or dense. Shape (N, M)
       The regressors

    momentum : float, optional
        The persistence of the gradient.

    step_size : float, optional
        The increment of parameter update in each iteration

    non_neg : Boolean, optional
        Whether to enforce non-negativity of the solution.

    check_error_iter : int, optional
        How many rounds to run between error evaluation for
        convergence-checking.

    max_error_checks : int, optional
        Don't check errors more than this number of times if no improvement in
        r-squared is seen.

    converge_on_sse : float, optional
      a percentage improvement in SSE that is required each time to say
      that things are still going well.

    Returns
    -------
    h_best : The best estimate of the parameters.

    """
    num_regressors = X.shape[1]
    # Initialize the parameters at the origin:
    h = np.zeros(num_regressors)
    # If nothing good happens, we'll return that:
    h_best = h
    iteration = 1
    ss_residuals_min = np.inf  # This will keep track of the best solution
    sse_best = np.inf  # This will keep track of the best performance so far
    count_bad = 0  # Number of times estimation error has gone up.
    error_checks = 0  # How many error checks have we done so far

    while 1:
        if iteration > 1:
            # The gradient is (Kay 2008 supplemental page 27):
            gradient = spdot(X.T, spdot(X, h) - y)
            gradient += momentum * gradient
            # Normalize to unit-length
            unit_length_gradient = gradient / np.sqrt(np.dot(gradient, gradient))
            # Update the parameters in the direction of the gradient:
            h -= step_size * unit_length_gradient
            if non_neg:
                # Set negative values to 0:
                h[h < 0] = 0

        # Every once in a while check whether it's converged:
        if np.mod(iteration, check_error_iter):
            # This calculates the sum of squared residuals at this point:
            sse = np.sum((y - spdot(X, h)) ** 2)
            # Did we do better this time around?
            if sse < ss_residuals_min:
                # Update your expectations about the minimum error:
                ss_residuals_min = sse
                h_best = h  # This holds the best params we have so far
                # Are we generally (over iterations) converging on
                # sufficient improvement in r-squared?
                if sse < converge_on_sse * sse_best:
                    sse_best = sse
                    count_bad = 0
                else:
                    count_bad += 1
            else:
                count_bad += 1

            if count_bad >= max_error_checks:
                return h_best
            error_checks += 1
        iteration += 1


class SKLearnLinearSolver(metaclass=abc.ABCMeta):
    """
    Provide a sklearn-like uniform interface to algorithms that solve problems
    of the form: $y = Ax$ for $x$

    Sub-classes of SKLearnLinearSolver should provide a 'fit' method that have
    the following signature: `SKLearnLinearSolver.fit(X, y)`, which would set
    an attribute `SKLearnLinearSolver.coef_`, with the shape (X.shape[1],),
    such that an estimate of y can be calculated as:
    `y_hat = np.dot(X, SKLearnLinearSolver.coef_.T)`
    """

    def __init__(self, *args, **kwargs):
        self._args = args
        self._kwargs = kwargs

    @abc.abstractmethod
    def fit(self, X, y):
        """Implement for all derived classes"""

    def predict(self, X):
        """
        Predict using the result of the model

        Parameters
        ----------
        X : array-like (n_samples, n_features)
            Samples.

        Returns
        -------
        C : array, shape = (n_samples,)
            Predicted values.
        """
        X = np.asarray(X)
        return np.dot(X, self.coef_.T)


class NonNegativeLeastSquares(SKLearnLinearSolver):
    """
    A sklearn-like interface to scipy.optimize.nnls

    """

    def fit(self, X, y):
        """
        Fit the NonNegativeLeastSquares linear model to data

        Parameters
        ----------
        X : array-like (n_samples, n_features)
            Samples.
        y : array-like (n_samples,)
            Target values.

        """
        coef, rnorm = opt.nnls(X, y)
        self.coef_ = coef
        return self


class PositiveDefiniteLeastSquares:
    @warning_for_keywords()
    def __init__(self, m, *, A=None, L=None):
        r"""Regularized least squares with linear matrix inequality constraints.

        See :footcite:p:`DelaHaije2020` for further details about the method.

        Generate a CVXPY representation of a regularized least squares
        optimization problem subject to linear matrix inequality constraints.

        Parameters
        ----------
        m : int
            Positive int indicating the number of regressors.
        A : array (t = m + k + 1, p, p), optional
            Constraint matrices $A$.
        L : array (m, m), optional
            Regularization matrix $L$.
            Default: None.

        Notes
        -----
        The basic problem is to solve for $h$ the minimization of

        $c=\|X h - y\|^2 + \|L h\|^2$,

        where $X$ is an (m, m) upper triangular design matrix and $y$ is a set
        of m measurements, subject to the constraint that

        $M=A_0+\sum_{i=0}^{m-1} h_i A_{i+1}+\sum_{j=0}^{k-1} s_j A_{m+j+1}>0$,

        where $s_j$ are slack variables and where the inequality sign denotes
        positive definiteness of the matrix $M$. The sparsity pattern and size
        of $X$ and $y$ are fixed, because every design matrix and set of
        measurements can be reduced to an equivalent (minimal) formulation of
        this type.

        This formulation is used here mainly to enforce polynomial
        sum-of-squares constraints on various models, as described in
        :footcite:p:`DelaHaije2020`.

        References
        ----------
        .. footbibliography::
        """
        # Note: several comments refer to DelaHaije2020 eq 22

        # Input
        self.A = A  # list: [zeros, H(omega), L(alpha)] --- see eq 22
        self.L = L  # this is a regularizer matrix, NOT L(alpha) eq 22

        # Problem size
        t = len(A) if A else 0
        k = t - m - 1  # length of alpha in L(alpha) eq 22

        # Unknowns
        self._X = cvxpy.Parameter((m, m))  # Design matrix
        self._f = cvxpy.Parameter(m)  # Given solution for feasibility check
        self._h = cvxpy.Variable(m)  # Solution to constrained problem
        self._y = cvxpy.Parameter(m)  # Regressand

        # Error output
        self._zeros = np.zeros(m)

        # Objective
        c = self._X @ self._h - self._y
        if L is not None:
            c += L @ self._h

        f_objective = cvxpy.Minimize(0)
        p_objective = cvxpy.Minimize(cvxpy.norm(c))

        # Constraints
        if t:
            M = F = A[0]  # first matrix all zeros (use to initialize)
            if k > 0:
                for i in range(m):  # loop over H(omega) from eq 22
                    F += self._f[i] * A[i + 1]
                    M += self._h[i] * A[i + 1]
                    # A-matrix for 'intercept' (22nd, i=21) should be zeros
                self._s = cvxpy.Variable(k)
                for j in range(k):  # loop over L(alpha) from eq 22
                    F += self._s[j] * A[m + j + 1]
                    M += self._s[j] * A[m + j + 1]
            else:
                for i in range(t - 1):
                    F += self._f[i] * A[i + 1]
                    M += self._h[i] * A[i + 1]
            f_constraints = [F >> 0]
            p_constraints = [M >> 0]
        else:
            f_constraints = p_constraints = []

        # CVXPY problems
        self.problem = cvxpy.Problem(p_objective, p_constraints)
        self.unconstrained_problem = cvxpy.Problem(p_objective)
        self.feasibility_problem = cvxpy.Problem(f_objective, f_constraints)

    @warning_for_keywords()
    def solve(self, design_matrix, measurements, *, check=False, **kwargs):
        r"""Solve CVXPY problem

        Solve a CVXPY problem instance for a given design matrix and a given set
        of observations, and return the optimum.

        Parameters
        ----------
        design_matrix : array (n, m)
            Design matrix.
        measurements : array (n)
            Measurements.
        check : boolean, optional
            If True check whether the unconstrained optimization solution
            already satisfies the constraints, before running the constrained
            optimization. This adds overhead, but can avoid unnecessary
            constrained optimization calls.
        kwargs : keyword arguments
            Arguments passed to the CVXPY solve method.

        Returns
        -------
        h : array (m)
             Estimated optimum for problem variables $h$.
        """

        # Compute and set reduced problem parameters
        try:
            X = np.linalg.cholesky(np.dot(design_matrix.T, design_matrix)).T
        except np.linalg.linalg.LinAlgError:
            msg = "Cholesky decomposition failed, returning zero array. Verify "
            msg += "that the data is sufficient to estimate the model "
            msg += "parameters, and that the design matrix has full rank."
            warnings.warn(msg, stacklevel=2)
            return self._zeros
        self._X.value = X
        self._y.value = np.linalg.multi_dot(
            [X, np.linalg.pinv(design_matrix), measurements]
        )

        try:
            # Check unconstrained solution
            if check:
                # Solve unconstrained problem
                self.unconstrained_problem.solve(**kwargs)

                # Return zeros if optimization failed
                status = self.unconstrained_problem.status
                if status != "optimal":
                    msg = f"Solver failed to produce an optimum: {status}."
                    warnings.warn(msg, stacklevel=2)
                    msg = "Optimization failed, returning zero array."
                    warnings.warn(msg, stacklevel=2)
                    return self._zeros

                # Return unconstrained solution if satisfactory
                self._f.value = self._h.value
                self.feasibility_problem.solve(**kwargs)
                if self.feasibility_problem.status == "optimal":
                    return np.asarray(self._h.value).squeeze()

            # Solve constrained problem
            self.problem.solve(**kwargs)

            # Show warning if solution is not optimal
            status = self.problem.status
            if status != "optimal":
                msg = f"Solver failed to produce an optimum: {status}."
                warnings.warn(msg, stacklevel=2)

            # Return solution
            return np.asarray(self._h.value).squeeze()

        except cvxpy.error.SolverError:
            # Return zeros
            msg = "Optimization failed, returning zero array."
            warnings.warn(msg, stacklevel=2)
            return self._zeros
