#
# Algebraic solver class
#
import casadi
import numpy as np
from scipy import optimize
from scipy.sparse import issparse

import pybamm


class AlgebraicSolver(pybamm.BaseSolver):
    """Solve a discretised model which contains only (time independent) algebraic
    equations using a root finding algorithm.
    Uses scipy.optimize.root.
    Note: this solver could be extended for quasi-static models, or models in
    which the time derivative is manually discretised and results in a (possibly
    nonlinear) algebaric system at each time level.

    Parameters
    ----------
    method : str, optional
        The method to use to solve the system (default is "lm"). If it starts with
        "lsq", least-squares minimization is used. The method for least-squares can be
        specified in the form "lsq_methodname"
    tol : float, optional
        The tolerance for the solver (default is 1e-6).
    extra_options : dict, optional
        Any options to pass to the rootfinder. Vary depending on which method is chosen.
        Please consult `SciPy documentation
        <https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.show_options.html>`_
        for details. Addititional options to pass to the solver, by default:

        .. code-block:: python

            extra_options = {
                # Tolerance for termination by the change of the independent variables.
                "xtol": 1e-12,
                # Tolerance for termination by the norm of the gradient.
                "gtol": 1e-12,
            }
    """

    def __init__(self, method="lm", tol=1e-6, extra_options=None):
        super().__init__(method=method)
        self.tol = tol

        default_extra_options = {
            "xtol": 1e-12,
            "gtol": 1e-12,
        }
        extra_options = extra_options or {}
        self.extra_options = default_extra_options | extra_options

        self.name = f"Algebraic solver ({method})"
        self._algebraic_solver = True
        pybamm.citations.register("Virtanen2020")

    def set_up_root_solver(self, model, inputs_dict, t_eval):
        """Create and return a rootfinder object. Not used for `pybamm.AlgebraicSolver`.

        Parameters
        ----------
        model : :class:`pybamm.BaseModel`
            The model whose solution to calculate.
        inputs_dict : dict
            Dictionary of inputs.
        t_eval : :class:`numpy.array`, size (k,)
            The times at which to compute the solution.
        """
        pass

    @property
    def tol(self):
        return self._tol

    @tol.setter
    def tol(self, value):
        self._tol = value

    def _integrate(self, model, t_eval, inputs_dict=None, t_interp=None):
        """
        Calculate the solution of the algebraic equations through root-finding

        Parameters
        ----------
        model : :class:`pybamm.BaseModel`
            The model whose solution to calculate.
        t_eval : :class:`numpy.array`, size (k,)
            The times at which to compute the solution
        inputs_dict : dict, optional
            Any input parameters to pass to the model when solving
        """
        inputs_dict = inputs_dict or {}
        if model.convert_to_format == "casadi":
            inputs = casadi.vertcat(*[x for x in inputs_dict.values()])
        else:
            inputs = inputs_dict

        y0 = model.y0
        if isinstance(y0, casadi.DM):
            y0 = y0.full()
        y0 = y0.flatten()

        # The casadi algebraic solver can read rhs equations, but leaves them unchanged
        # i.e. the part of the solution vector that corresponds to the differential
        # equations will be equal to the initial condition provided. This allows this
        # solver to be used for initialising the DAE solvers
        # Split y0 into differential and algebraic
        if model.rhs == {}:
            len_rhs = 0
        else:
            len_rhs = model.rhs_eval(t_eval[0], y0, inputs).shape[0]
        y0_diff, y0_alg = np.split(y0, [len_rhs])

        test_result = model.algebraic_eval(0, y0, inputs)

        if isinstance(test_result, casadi.DM):

            def algebraic(t, y):
                result = model.algebraic_eval(t, y, inputs)
                return result.full().flatten()

        else:

            def algebraic(t, y):
                result = model.algebraic_eval(t, y, inputs)
                return result.flatten()

        y_alg = np.empty((len(y0_alg), len(t_eval)))

        timer = pybamm.Timer()
        integration_time = 0
        for idx, t in enumerate(t_eval):

            def root_fun(y_alg, t=t):
                "Evaluates algebraic using y"
                y = np.concatenate([y0_diff, y_alg])
                out = algebraic(t, y)
                pybamm.logger.debug(
                    f"Evaluating algebraic equations at t={t}, L2-norm is {np.linalg.norm(out)}"
                )
                return out

            jac = model.jac_algebraic_eval
            if jac:
                if issparse(jac(t_eval[0], y0, inputs)):

                    def jac_fn(y_alg, jac=jac):
                        """
                        Evaluates Jacobian using y0_diff (fixed) and y_alg (varying)
                        """
                        y = np.concatenate([y0_diff, y_alg])
                        return jac(0, y, inputs)[:, len_rhs:].toarray()

                else:

                    def jac_fn(y_alg, jac=jac):
                        """
                        Evaluates Jacobian using y0_diff (fixed) and y_alg (varying)
                        """
                        y = np.concatenate([y0_diff, y_alg])
                        return jac(0, y, inputs)[:, len_rhs:]

            else:
                jac_fn = None

            itr = 0
            maxiter = 2
            success = False
            while not success:
                # Methods which use least-squares are specified as either "lsq",
                # which uses the default method, or with "lsq__methodname"
                if self.method.startswith("lsq"):
                    if self.method == "lsq":
                        method = "trf"
                    else:
                        method = self.method[5:]
                    if jac_fn is None:
                        jac_fn = "2-point"
                    timer.reset()
                    sol = optimize.least_squares(
                        root_fun,
                        y0_alg,
                        method=method,
                        ftol=self.tol,
                        jac=jac_fn,
                        bounds=model.bounds,
                        **self.extra_options,
                    )
                    integration_time += timer.time()
                    success |= sol.success
                # Methods which use minimize are specified as either "minimize",
                # which uses the default method, or with "minimize__methodname"
                elif self.method.startswith("minimize"):
                    # Adapt the root function for minimize
                    def root_norm(y):
                        return np.sum(root_fun(y) ** 2)

                    if jac_fn is None:
                        jac_norm = None
                    else:

                        def jac_norm(y, jac_fn=jac_fn):
                            return np.sum(2 * root_fun(y) * jac_fn(y), 0)

                    if self.method == "minimize":
                        method = None
                    else:
                        method = self.method[10:]
                    extra_options = self.extra_options
                    if np.any(model.bounds[0] != -np.inf) or np.any(
                        model.bounds[1] != np.inf
                    ):
                        bounds = [
                            (lb, ub)
                            for lb, ub in zip(
                                model.bounds[0], model.bounds[1], strict=True
                            )
                        ]
                    else:
                        bounds = None

                    hess = extra_options.get("hess", None)
                    hessp = extra_options.get("hessp", None)
                    constraints = extra_options.get("constraints", ())
                    callback = extra_options.get("callback", None)

                    timer.reset()
                    sol = optimize.minimize(
                        root_norm,
                        y0_alg,
                        method=method,
                        tol=self.tol,
                        jac=jac_norm,
                        bounds=bounds,
                        hess=hess,
                        hessp=hessp,
                        constraints=constraints,
                        callback=callback,
                        options=extra_options,
                    )
                    integration_time += timer.time()
                    # The solution.success flag is unreliable, so we ignore
                    # it and use the norm of the residual instead
                else:
                    timer.reset()
                    sol = optimize.root(
                        root_fun,
                        y0_alg,
                        method=self.method,
                        tol=self.tol,
                        jac=jac_fn,
                        options=self.extra_options,
                    )
                    integration_time += timer.time()
                    # The solution.success flag is unreliable, so we ignore
                    # it and use the norm of the residual instead

                y0_alg = sol.x
                success |= np.all(abs(sol.fun) < self.tol)
                if success:
                    y_alg[:, idx] = y0_alg
                    break

                if itr > maxiter:
                    raise pybamm.SolverError(
                        "Could not find acceptable solution: solver terminated "
                        "unsuccessfully and maximum solution error "
                        f"({np.max(abs(sol.fun))}) above tolerance ({self.tol})"
                    )
                itr += 1

        # Concatenate differential part
        y_diff = np.r_[[y0_diff] * len(t_eval)].T
        y_sol = np.r_[y_diff, y_alg]
        # Return solution object (no events, so pass None to t_event, y_event)
        sol = pybamm.Solution(
            t_eval, y_sol, model, inputs_dict, termination="final time"
        )
        sol.integration_time = integration_time
        return sol
