import numpy
import scipy
import scipy.linalg
import matplotlib.pyplot as plt

def schroedinger_matrix(V, h):
    N = len(V)
    A = numpy.zeros((N,N))

    for i in range(1,N):
        A[i,i-1] = -1.0/(h*h)
    for i in range(0,N):
        A[i,i] = +2.0/(h*h) + V[i]
    for i in range(0, N-1):
        A[i,i+1] = -1.0/(h*h)

    return A

def solve_schroedinger_scipy(V, h, n):
    A = schroedinger_matrix(V, h)
    N = len(V)

    print 'Computing eigenvalues...'
    E, psi = scipy.linalg.eigh(A, eigvals=(0,n-1))

    return E, psi

def qr_eigenvalues(A, tolerance):
    n = A.shape[0]
    I = numpy.identity(n)
    Ak = A.copy()
    error = tolerance+1
    iteration = 0
    while error > tolerance:
        shift = Ak[-1,-1]
        Q, R = scipy.linalg.qr(Ak - shift*I)
        Ak = numpy.dot(R, Q) + shift*I
        error = max([abs(Ak[i,k]) \
                         for i in range(n) \
                         for k in range(i)])
        iteration += 1
        print "  iteration %d: error = %lf" % (iteration, error)
    return numpy.array([Ak[i,i] for i in range(n)])

def inverse_iteration(A, l, tolerance):
    n = A.shape[0]
    x = numpy.ones(n)
    Ashift = A - l*numpy.identity(n)
    error = tolerance + 1
    while error > tolerance:
        x = scipy.linalg.solve(Ashift, x)
        x = x / scipy.linalg.norm(x)
        error = scipy.linalg.norm(l*x - numpy.dot(A, x))
    return x

def solve_schroedinger_qr(V, h, n, tolerance=1e-1):
    A = schroedinger_matrix(V, h)
    N = len(V)
    A = numpy.zeros((N,N))

    for i in range(1,N):
        A[i,i-1] = -1.0/(h*h)
    for i in range(0,N):
        A[i,i] = +2.0/(h*h) + V[i]
    for i in range(0, N-1):
        A[i,i+1] = -1.0/(h*h)

    print 'Computing eigenvalues...'
    E = qr_eigenvalues(A, tolerance)[:n]
    psi = numpy.empty((N, n))

    print 'Computing inverse iteration...'
    for i in range(n):
        psi[:, i] = inverse_iteration(A, E[i], tolerance)

    return E, psi

N = 50
n = 5
L = 10.0

h = L/N

x = numpy.linspace(0.0, L, N)
V = numpy.zeros(N)

# # finite well
V[:N/5] = 2.0
V[-N/5:] = 2.0

# harmonic well
#V = 0.1*(x-L/2)**2

# double well
#V[N*0.4:N*0.6] = 2.0

# add "infinite" walls
#V[0] = 1000.0
#V[-1] = 1000.0

# plot potential
p1 = plt.subplot(221)
p1.plot(x[1:-1], V[1:-1], 'o-')
p1.set_title('psi')

p2 = plt.subplot(222)
p2.plot(x[1:-1], V[1:-1], 'o-')
p2.set_title('psi**2')

p3 = plt.subplot(223)
p3.plot(x[1:-1], V[1:-1], 'o-')
p3.set_title('psi')

p4 = plt.subplot(224)
p4.plot(x[1:-1], V[1:-1], 'o-')
p4.set_title('psi**2')

# V
legends = ['V']

print "SciPy Solving..."
Es, psis = solve_schroedinger_scipy(V, h, n)

for i in range(n):
    E = Es[i]
    psi = psis[:,i]
    p1.plot(x, E+psi, 'o-')
    p2.plot(x, E+psi**2, 'o-')
    legends.append('E = %5.2f' % E)
p2.legend(legends)

# V
legends = ['V']

print "QR Solving..."
Es, psis = solve_schroedinger_qr(V, h, n, tolerance=1e-1)

for i in range(n):
    E = Es[i]
    psi = psis[:,i]
    p3.plot(x, E+psi, 'o-')
    p4.plot(x, E+psi**2, 'o-')
    legends.append('E = %5.2f' % E)
p4.legend(legends)

E = [(i*numpy.pi/L)**2 for i in range(n)] 
print "E = ", E

plt.show()