from __future__ import print_function

old_print = print

from numpy import zeros_like, zeros
from numpy.linalg import solve

import numpy
from numba import guvectorize
from numba import double


def solve(m, sol):

    # seems to segfault on windows

    h,w = m.shape

    for y in range(0,h):
        maxrow = y
        for y2 in range(y+1, h):    # Find max pivot
            if abs(m[y2,y]) > abs(m[maxrow,y]):
                maxrow = y2
        for y2 in range(0,w):
            t = m[y,y2]
            m[y,y2] = m[maxrow,y2]
            m[maxrow,y2] = t

        for y2 in range(y+1, h):    # Eliminate column y
            c = m[y2,y] / m[y,y]
            for x in range(y, w):
                m[y2,x] -= m[y,x] * c

    for y in range(h-1, 0-1, -1): # Backsubstitute
        c  = m[y,y]
        for y2 in range(0,y):
            for x in range(w-1, y-1, -1):
                m[y2,x] -=  m[y,x] * m[y2,y] / c
        m[y,y] /= c
        for x in range(h, w):       # Normalize row y
          m[y,x] /= c

    for y in range(h):
        sol[y] = m[y,w-1]

serial_solve_numba = guvectorize('void(f8[:,:], f8[:])', '(m,n)->(m)')(solve)


from numpy.linalg import solve as linalg_solve


def serial_solve(A, B, diagnose=True):


    if diagnose:

        sol = zeros_like(B)

        for i in range(sol.shape[0]):
            try:
                sol[i,:] = linalg_solve( A[i,:,:], B[i,:])
            except:
                # Should be a special type of exception
                a = Exception("Error solving point {}".format(i))
                a.x = B[i,:]
                a.J = A[i,:,:]
                a.i = i
                raise a

    else:
        M = numpy.concatenate([A,B[:,:,None]],axis=2)
        sol = numpy.zeros_like(B)
        serial_solve_numba(M,sol)

    return sol
#
#     return sol

import time

def newton(f, x, verbose=False, tol=1e-6, maxit=5, jactype='serial'):


    """Solve nonlinear system using safeguarded Newton iterations


    Parameters
    ----------

    Return
    ------
    """
    if verbose:
        print = lambda txt: old_print(txt)
    else:
        print = lambda txt: None

    it = 0
    error = 10
    converged = False
    maxbacksteps = 30

    x0 = x

    if jactype == 'sparse':
        from scipy.sparse.linalg import spsolve as solve
    elif jactype == 'full':
        from numpy.linalg import solve
    else:
        solve = serial_solve

    while it<maxit and not converged:

        [v,dv] = f(x)

        # TODO: rewrite starting here

#        print("Time to evaluate {}".format(ss-tt)0)

        error_0 = abs(v).max()

        if error_0 < tol:

            if verbose:
                print("> System was solved after iteration {}. Residual={}".format(it,error_0))
            converged = True

        else:

            it += 1

            dx = solve(dv, v)

            # norm_dx = abs(dx).max()

            for bck in range(maxbacksteps):
                xx = x - dx*(2**(-bck))
                vm = f(xx)[0]
                err = abs(vm).max()
                if err < error_0:
                    break

            x = xx

            if verbose:
                print("\t> {} | {} | {}".format(it, err, bck))

    if not converged:
        import warnings
        warnings.warn("Did not converge")
    return [x, it]

serial_newton = newton

from numpy import sqrt, finfo, inf

from numpy import isinf, newaxis, diag, zeros

def SerialDifferentiableFunction(f, epsilon=1e-8):

    def df(x):

        v0 = f(x)

        N = v0.shape[0]
        n_v = v0.shape[1]
        assert(x.shape[0] == N)
        n_x = x.shape[1]

        dv = zeros( (N, n_v, n_x) )

        for i in range(n_x):

            xi = x.copy()
            xi[:,i] += epsilon

            vi = f(xi)

            dv[:,:,i] = (vi - v0)/epsilon

        return [v0, dv]

    return df



def test_serial_solve():

    N = 10

    import numpy
    A = numpy.random.random( (N,2,2) )


    B = numpy.random.random( (N,2) )

    print(A)
    print(B)
    out = serial_solve(A,B)
    print("A")
    print(A)
    print("B")
    print(B)
    print("out")
    print(out)

    import numpy.linalg

    print( numpy.linalg.solve(A[0,:,:], B[0,:]) )


if __name__ == "__main__":

    test_serial_solve()
