"""
contique: Numerical continuation of nonlinear equilibrium equations.
"""

import numpy as np
from scipy import sparse

from .helpers import control, one_hot
from .jacobian import jacobian
from .newton import newtonrhapson


def funxt(y, one_hot_vector, ymax, fun, jac=None, jacmode=3, jaceps=None, args=(None,)):
    """Extend the given equilibrium equations.

    Parameters
    ----------
    y : array
        1d-array of unknowns
    one_hot_vector : array
        (pre-evaluated) one-hot vector
    ymax : array
        1d-array with max. allowed values of unknows
    fun : function
        1d-array of equilibrium equations
    jac : function, optional
        jacobian of fun w.r.t. the extended unknows y
    jacmode : int, optional
        forward (2) or central (3) finite-differences approx. of the jacobian
    jaceps : float, optional
        user-specified stepwidth (if None, this defaults to eps^(1/mode))
    args : tuple, optional
        Optional tuple of arguments which are passed to the function. Even if
        only one argument is passed, it has to be encapsulated in a tuple
        (default is (None,)).

    Returns
    -------
    array
        extended 1d-array of equilibrium equations
        with control equation
    """

    # split the unknowns
    x, lpf = y[:-1], y[-1]

    # evaluate the given function
    f = fun(x, lpf, *args)

    if sparse.issparse(f):
        # convert function vector to array
        f = f.toarray()

    # extend the function
    return np.append(f, np.dot(one_hot_vector, (y - ymax)))


def jacxt(y, one_hot_vector, ymax, fun, jac=None, jacmode=3, jaceps=None, args=(None,)):
    """Jacobian of extended equilibrium equations.

    Parameters
    ----------
    y : ndarray
        1d-array of extended unknows
    one_hot_vector : ndarray
        1d-array with pre-evaluated one-hot vector
    ymax : ndarray
        1d-array with max. allowed incremental increase values of y
    fun : function
        function in terms of unknows x and optional args which returns the
        equilibrium equations.
    jac : function, optional
        jacobian of fun w.r.t. the extended unknows y
    jacmode : int, optional
        forward (2) or central (3) finite-differences approx. of the jacobian
    jaceps : float, optional
        user-specified stepwidth (if None, this defaults to eps^(1/mode))
    args : tuple, optional
        Optional tuple of arguments which are passed to the function. Even if
        only one argument is passed, it has to be encapsulated in a tuple
        (default is (None,)).

    Returns
    -------
        ndarray
        jacobian of fun w.r.t. y (contains both derivatives of x and lpf)
        as 2d-array
    """

    # split the unknowns
    x, lpf = y[:-1], y[-1]

    if jac is None:
        # evaluate by finite differences method
        dfundx = jacobian(fun, argnum=0, mode=jacmode, h=jaceps)
        dfundl = jacobian(fun, argnum=1, mode=jacmode, h=jaceps)
    else:
        dfundx, dfundl = jac

    # evaluate the given jacobian
    dfdx = dfundx(x, lpf, *args)
    dfdl = dfundl(x, lpf, *args).reshape(-1, 1)

    # define horizontal and vertical stack operations based on evaluated
    # sparse or dense jacobian
    if sparse.issparse(dfdx):
        hstack = sparse.hstack
        vstack = sparse.vstack
        array = sparse.csr_matrix
    else:
        hstack = np.hstack
        vstack = np.vstack
        array = np.array

    # extend the jacobian
    dfdy = hstack([array(dfdx), array(dfdl)])
    dgdy = vstack([dfdy, array(one_hot_vector)])

    if sparse.issparse(dfdx):
        # convert to compressed sparse row format
        dgdy = dgdy.tocsr()

    return dgdy


def newtonxt(
    fun,
    jac,
    y0,
    control0,
    dymax,
    jacmode=3,
    jaceps=None,
    args=(None,),
    maxiter=20,
    tol=1e-8,
    solve=None,
):
    """Solve equilibrium equations starting from an initial solution
    with a given control component and a max. allowed increase of unknowns.

    Parameters
    ----------
    fun : function
        function in terms of extended unknows and optional args which returns
        the extended equilibrium equations
    jac : function, optional
        jacobian of fun w.r.t. the extended unknows
    y0 : ndarray
        1d-array of initial extended unknows
    control0 : tuple of int, optional
        initial tuple of control component and sign
    dxmax : float, optional
        max. allowed absolute incremental increase of extended unknowns per step
    jacmode : int, optional
        forward (2) or central (3) finite-differences approx. of the jacobian
    jaceps : float, optional
        user-specified stepwidth (if None, this defaults to eps^(1/mode))
    args : tuple, optional
        Optional tuple of arguments which are passed to the function. Eeven if only
        one argument is passed, it has to be encapsulated in a tuple (default is
        (None,)).
    maxiter : int, optional
        max. number of Newton-iterations per cycle
    tol : float, optional
        tolerated residual of the norm of the equilibrium equation (default is 1e-8)
    solve: callable, optional
        A solver.

    Returns
    -------
    res : NewtonResult
        Instance of NewtonResult with res.x being the final extended unknowns
    """

    # init one-hot vector and obtain ymax
    component0, sign0 = control0
    one_hot_vector = one_hot(component0, len(y0))
    ymax = y0 + sign0 * dymax

    # Newton-Rhapson solver
    res = newtonrhapson(
        fun=funxt,
        x0=y0,
        jac=jacxt,
        args=(one_hot_vector, ymax, fun, jac, jacmode, jaceps, args),
        maxiter=maxiter,
        tol=tol,
        solve=solve,
    )

    # normalized dy = dy/dymax
    res.dys = (res.x - y0) / dymax

    # final control component based on normalized dy
    if np.any(np.isnan(res.dys)):
        res.control = control0
    else:
        res.control = control(res.dys)

    return res
