import warnings
import numpy as np
from scipy.interpolate import interp1d

from pySDC.projects.DAE.misc.ProblemDAE import ptype_dae


class pendulum_2d(ptype_dae):
    r"""
    Example implementing the well known 2D pendulum as a first order differential-algebraic equation (DAE) of index 3.
    The DAE system is given by the equations

    .. math::
        x' = u,

    .. math::
        \frac{d}{dt} \frac{\partial}{\partial u} L = \frac{\partial L}{\partial x} + f + G^{T} \lambda,

    .. math::
        0 = \phi.

    The pendulum is used in most introductory literature on DAEs, for example on page 8 of [1]_.

    Parameters
    ----------
    nvars : int
        Number of unknowns of the system of DAEs.
    newton_tol : float
        Tolerance for Newton solver.

    Attributes
    ----------
    t_end: float
        The end time at which the reference solution is determined.

    References
    ----------
    .. [1] E. Hairer, C. Lubich, M. Roche. The numerical solution of differential-algebraic systems by Runge-Kutta methods.
        Lect. Notes Math. (1989).
    """

    def __init__(self, nvars, newton_tol):
        """Initialization routine"""
        super().__init__(nvars, newton_tol)
        # load reference solution
        # data file must be generated and stored under misc/data and self.t_end = t[-1]
        # data = np.load(r'pySDC/projects/DAE/misc/data/pendulum.npy')
        # t = data[:, 0]
        # solution = data[:, 1:]
        # self.u_ref = interp1d(t, solution, kind='cubic', axis=0, fill_value='extrapolate')
        self.t_end = 0.0

    def eval_f(self, u, du, t):
        r"""
        Routine to evaluate the implicit representation of the problem, i.e., :math:`F(u, u', t)`.

        Parameters
        ----------
        u : dtype_u
            Current values of the numerical solution at time t.
        du : dtype_u
            Current values of the derivative of the numerical solution at time t.
        t : float
            Current time of the numerical solution.

        Returns
        -------
        f : dtype_f
            Current value of the right-hand side of f (which includes five components).
        """
        g = 9.8
        # The last element of u is a Lagrange multiplier. Not sure if this needs to be time dependent, but must model the
        # weight somehow
        f = self.dtype_f(self.init)
        f[:] = (du[0] - u[2], du[1] - u[3], du[2] + u[4] * u[0], du[3] + u[4] * u[1] + g, u[0] ** 2 + u[1] ** 2 - 1)
        return f

    def u_exact(self, t):
        """
        Approximation of the exact solution generated by spline interpolation of an extremely accurate numerical reference solution.

        Parameters
        ----------
        t : float
            The time of the reference solution.

        Returns
        -------
        me : dtype_u
            The reference solution as mesh object. It contains fixed initial conditions at initial time.
        """
        me = self.dtype_u(self.init)
        if t == 0:
            me[:] = (-1, 0, 0, 0, 0)
        elif t < self.t_end:
            me[:] = self.u_ref(t)
        else:
            warnings.warn("Requested time exceeds domain of the reference solution. Returning zero.")
            me[:] = (0, 0, 0, 0, 0)
        return me


class simple_dae_1(ptype_dae):
    r"""
    Example implementing a smooth linear index-2 differential-algebraic equation (DAE) with known analytical solution.
    The DAE system is given by

    .. math::
        \frac{d u_1 (t)}{dt} = (\alpha - \frac{1}{2 - t}) u_1 (t) + (2-t) \alpha z (t) + \frac{3 - t}{2 - t},

    .. math::
        \frac{d u_2 (t)}{dt} = \frac{1 - \alpha}{t - 2} u_1 (t) - u_2 (t) + (\alpha - 1) z (t) + 2 e^{t},

    .. math::
        0 = (t + 2) u_1 (t) + (t^{2} - 4) u_2 (t) - (t^{2} + t - 2) e^{t}.

    The exact solution of this system is

    .. math::
        u_1 (t) = u_2 (t) = e^{t},

    .. math::
        z (t) = -\frac{e^{t}}{2 - t}.

    This example is commonly used to test that numerical implementations are functioning correctly. See, for example,
    page 267 of [1]_.

    Parameters
    ----------
    nvars : int
        Number of unknowns of the system of DAEs.
    newton_tol : float
        Tolerance for Newton solver.

    References
    ----------
    .. [1] U. Ascher, L. R. Petzold. Computer method for ordinary differential equations and differential-algebraic
        equations. Society for Industrial and Applied Mathematics (1998).
    """

    def eval_f(self, u, du, t):
        r"""
        Routine to evaluate the implicit representation of the problem, i.e., :math:`F(u, u', t)`.

        Parameters
        ----------
        u : dtype_u
            Current values of the numerical solution at time t.
        du : dtype_u
            Current values of the derivative of the numerical solution at time t.
        t : float
            Current time of the numerical solution.

        Returns
        -------
        f : dtype_f
            Current value of the right-hand side of f (which includes three components).
        """
        # Smooth index-2 DAE pg. 267 Ascher and Petzold (also the first example in KDC Minion paper)
        a = 10.0
        f = self.dtype_f(self.init)
        f[:] = (
            -du[0] + (a - 1 / (2 - t)) * u[0] + (2 - t) * a * u[2] + np.exp(t) * (3 - t) / (2 - t),
            -du[1] + (1 - a) / (t - 2) * u[0] - u[1] + (a - 1) * u[2] + 2 * np.exp(t),
            (t + 2) * u[0] + (t**2 - 4) * u[1] - (t**2 + t - 2) * np.exp(t),
        )
        return f

    def u_exact(self, t):
        """
        Routine for the exact solution.

        Parameters
        ----------
        t : float
            The time of the reference solution.

        Returns
        -------
        me : dtype_u
            The reference solution as mesh object containing three components.
        """
        me = self.dtype_u(self.init)
        me[:] = (np.exp(t), np.exp(t), -np.exp(t) / (2 - t))
        return me


class problematic_f(ptype_dae):
    r"""
    Standard example of a very simple fully implicit index-2 differential algebraic equation (DAE) that is not
    numerically solvable for certain choices of the parameter :math:`\eta`. The DAE system is given by

    .. math::
        y (t) + \eta t z (t) = f(t),

    .. math::
        \frac{d y(t)}{dt} + \eta t \frac{d z(t)}{dt} + (1 + \eta) z (t) = g (t).

    See, for example, page 264 of [1]_.

    Parameters
    ----------
    nvars : int
        Number of unknowns of the system of DAEs.
    newton_tol : float
        Tolerance for Newton solver.

    Attributes
    ----------
    eta: float
        Specific parameter of the problem.

    References
    ----------
    .. [1] U. Ascher, L. R. Petzold. Computer method for ordinary differential equations and differential-algebraic
        equations. Society for Industrial and Applied Mathematics (1998).
    """

    def __init__(self, nvars, newton_tol, eta=1):
        """Initialization routine"""
        super().__init__(nvars, newton_tol)
        self._makeAttributeAndRegister('eta', localVars=locals())

    def eval_f(self, u, du, t):
        r"""
        Routine to evaluate the implicit representation of the problem, i.e., :math:`F(u, u', t)`.

        Parameters
        ----------
        u : dtype_u
            Current values of the numerical solution at time t.
        du : dtype_u
            Current values of the derivative of the numerical solution at time t.
        t : float
            Current time of the numerical solution.

        Returns
        -------
        f : dtype_f
            Current value of the right-hand side of f (which includes two components).
        """
        f = self.dtype_f(self.init)
        f[:] = (
            u[0] + self.eta * t * u[1] - np.sin(t),
            du[0] + self.eta * t * du[1] + (1 + self.eta) * u[1] - np.cos(t),
        )
        return f

    def u_exact(self, t):
        """
        Routine for the exact solution.

        Parameters
        ----------
        t : float
            The time of the reference solution.

        Returns
        -------
        me : dtype_u
            The reference solution as mesh object containing two components.
        """
        me = self.dtype_u(self.init)
        me[:] = (np.sin(t), 0)
        return me
