import numpy as np
import scipy
from . import options
from . import constants
from .utilities import convergence_check


def mdn_solver(x, mna, circ, T, MAXIT, nv, locked_nodes, time=None,
               print_steps=False, vector_norm=None, debug=True):
    """
    Solves a problem like F(x) = 0 using the Newton Algorithm with a variable
    damping.

    Where:

    .. math::

        F(x) = mna*x + T + T(x)

    * :math:`mna` is the Modified Network Analysis matrix of the circuit
    * :math:`T(x)` is the contribute of nonlinear elements to KCL
    * :math:`T` contains the contributions of the independent sources, time
    * invariant and linear

    If :math:`x(0)` is the initial guess, every :math:`x(n+1)` is given by:

    .. math::
        x(n+1) = x(n) + td \\cdot dx

    Where :math:`td` is a damping coefficient to avoid overflow in non-linear
    components and excessive oscillation in the very first iteration. Afterwards
    :math:`td=1` To calculate :math:`td`, an array of locked nodes is needed.

    The convergence check is done this way:

    **Parameters:**

    x : ndarray
        The initial guess. If set to ``None``, it will be initialized to all
        zeros. Specifying a initial guess may improve the convergence time of
        the algorithm and determine which solution (if any) is found if there
        are more than one.
    mna : ndarray
        The Modified Network Analysis matrix of the circuit, reduced, see above.
    circ : circuit instance
        The circuit instance.
    T : ndarray,
        The :math:`T` vector described above.
    MAXIT : int
        The maximum iterations that the method may perform.
    nv : int
        Number of nodes in the circuit (counting the ref, 0)
    locked_nodes : list of tuples
        A list of ports driving non-linear elements, generated by
        ``circ.get_locked_nodes()``
    time : float or None, optional
        The value of time to be passed to non_linear _and_ time variant
        elements.
    print_steps : boolean, optional
        Show a progress indicator, very verbose. Defaults to ``False``.
    vector_norm : function, optional
        An R^N -> R^1 function returning the norm of a vector, for convergence
        checking. Defaults to the maximum norm, ie :math:`f(x) = max(|x|)`,
    debug : int, optional
        Debug flag that will result in an array being returned containing
        node-by-node convergence information.

    **Returns:**

    sol : ndarray
        The solution.
    err : ndarray
        The remaining error.
    converged : boolean
        A boolean that is set to ``True`` whenever the method exits because of a
        successful convergence check. ``False`` whenever convergence problems
        where found.
    N : int
        The number of NR iterations performed.
    convergence_by_node : list
        If ``debug`` was set to ``True``, this list has the same size of the MNA
        matrix and contains the information regarding which nodes fail to
        converge in the circuit. Ie. ``if convergence_by_node[j] == False``,
        node ``j`` has a convergence problem. This may significantly help
        debugging non-convergent circuits.

    """
    if vector_norm is None:
        vector_norm = lambda v: max(abs(v)),
    mna_size = mna.shape[0]
    nonlinear_circuit = circ.is_nonlinear()
    if x is None:
        # if no guess was specified, its all zeros
        x = np.zeros((mna_size, 1))
    else:
        if x.shape[0] != mna_size:
            raise ValueError("x0s size is different from expected: got "
                             "%d-elements x0 with an MNA of size %d" %
                             (x.shape[0], mna_size))
    if T is None:
        #printing.print_warning(
        #    "dc_analysis.mdn_solver called with T==None, setting T=0. BUG or no sources in circuit?")
        T = np.zeros((mna_size, 1))

    sparse = mna_size > options.dense_matrix_limit
    # We allocate the matrices once and then reuse them
    if sparse:
        mna = scipy.sparse.coo_matrix(mna)
        J = scipy.sparse.lil_matrix((mna_size, mna_size))
    else:
        J = np.zeros((mna_size, mna_size))
    Tx = np.zeros((mna_size, 1))
    converged = False
    iteration = 0
    while iteration < MAXIT:  # newton iteration counter
        iteration += 1
        if nonlinear_circuit:
            # build dT(x)/dx (stored in J) and Tx(x)
            J[:, :] = 0.0
            Tx[:, 0] = 0.0
            for elem in circ:
                if elem.is_nonlinear:
                    _update_J_and_Tx(J, Tx, x, elem, time)
        residuo = mna.dot(x) + T + nonlinear_circuit*Tx

        if sparse:
            lu = scipy.sparse.linalg.splu(scipy.sparse.csc_matrix(mna + nonlinear_circuit*J))
            dx = lu.solve(-residuo)
        else:
            dx = np.linalg.solve(mna + nonlinear_circuit*J, -residuo)
        x = x + get_td(dx, locked_nodes, n=iteration) * dx
        if not nonlinear_circuit:
            converged = True
            break
        elif convergence_check(x, dx, residuo, nv - 1)[0]:
            converged = True
            break
        # if vector_norm(dx) == np.nan: #Overflow
        #   raise OverflowError
    if debug and not converged:
        # re-run the convergence check, only this time get the results
        # by node, so we can show to the users which nodes are misbehaving.
        converged, convergence_by_node = convergence_check(x, dx, residuo, 
                                                           nv - 1, debug=True)
    else:
        convergence_by_node = []
    return (x, residuo, converged, iteration, convergence_by_node)

def _update_J_and_Tx(J, Tx, x, elem, time):
    out_ports = elem.get_output_ports()
    for index in range(len(out_ports)):
        n1, n2 = out_ports[index]
        n1m1, n2m1 = n1 - 1, n2 - 1
        dports = elem.get_drive_ports(index)
        v_dports = []
        for port in dports:
            v = 0.  # build v: remember we removed the 0 row and 0 col of mna -> -1
            if port[0]:
                v = v + x[port[0] - 1, 0]
            if port[1]:
                v = v - x[port[1] - 1, 0]
            v_dports.append(v)
        if hasattr(elem, 'gstamp') and hasattr(elem, 'istamp'):
            iis, gs = elem.gstamp(v_dports, time)
            J[iis] += gs.reshape(-1)
            iis, i = elem.istamp(v_dports, time)
            Tx[iis] += i.reshape(-1)
            continue
        if n1 or n2:
            iel = elem.i(index, v_dports, time)
        if n1:
            Tx[n1m1, 0] = Tx[n1m1, 0] + iel
        if n2:
            Tx[n2m1, 0] = Tx[n2m1, 0] - iel
        for iindex in range(len(dports)):
            if n1 or n2:
                g = elem.g(index, v_dports, iindex, time)
            if n1:
                if dports[iindex][0]:
                    J[n1m1, dports[iindex][0] - 1] += g
                if dports[iindex][1]:
                    J[n1m1, dports[iindex][1] - 1] -= g
            if n2:
                if dports[iindex][0]:
                    J[n2m1, dports[iindex][0] - 1] -= g
                if dports[iindex][1]:
                    J[n2m1, dports[iindex][1] - 1] += g


def get_td(dx, locked_nodes, n=-1):
    """Calculates the damping coefficient for the Newthon method.

    The damping coefficient is choosen as the lowest between:

    - the damping required for the first NR iterations, a parameter which is set
      through the integer ``options.nr_damp_first_iters``.
    - If ``options.nl_voltages_lock`` evaluates to ``True``, the biggest damping
      factor that keeps the change in voltage across the locked nodes pairs less
      than the maximum variation allowed, set by:
      ``(options.nl_voltages_lock_factor * Vth)``
    - Unity.

    **Parameters:**

    dx : ndarray
        The undamped increment returned by the NR solver.
    locked_nodes : list
        A vector of tuples of (internal) nodes that are a port of a non-linear
        component.
    n : int, optional
        The NR iteration counter

    .. note::

        If ``n`` is set to ``-1`` (or any negative value), ``td`` is independent
        from the iteration number and ``options.nr_damp_first_iters`` is ignored.

    **Returns:**

    td : float
        The damping coefficient.

    """
    if not options.nr_damp_first_iters or n < 0:
        td = 1
    else:
        if n < 10:
            td = 1e-2
        elif n < 20:
            td = 0.1
        else:
            td = 1
    td_new = 1
    if options.nl_voltages_lock:
        for (n1, n2) in locked_nodes:
            if n1 != 0:
                if n2 != 0:
                    if abs(dx[n1 - 1, 0] - dx[n2 - 1, 0]) > options.nl_voltages_lock_factor * constants.Vth():
                        td_new = (options.nl_voltages_lock_factor * constants.Vth()) / abs(
                            dx[n1 - 1, 0] - dx[n2 - 1, 0])
                else:
                    if abs(dx[n1 - 1, 0]) > options.nl_voltages_lock_factor * constants.Vth():
                        td_new = (options.nl_voltages_lock_factor * constants.Vth()) / abs(
                            dx[n1 - 1, 0])
            else:
                if abs(dx[n2 - 1, 0]) > options.nl_voltages_lock_factor * constants.Vth():
                    td_new = (options.nl_voltages_lock_factor * constants.Vth()) / abs(
                        dx[n2 - 1, 0])
            if td_new < td:
                td = td_new
    return td
