# studying problem (-pu')' + qu = f, u(0)=0 
# user specified functions p,q,f and grid x as input to code

# this particular code considers Legendre equation
# Legendre polynomials are computed through a discrete version of
# the eigenvalue problem, see comments.

import numpy as np
import scipy as sp
import matplotlib.pyplot as plt

def p(x):
    return x**2 - 1

def q(x):
    return 1

def f(x):
    return 0

def U_eigs(x,k):
    n = x.shape[0]
    h = np.diff(x)
    A = np.zeros((n,n))
    M = np.zeros((n,n))
    b = np.zeros((n,1))
    # loop over elements
    for i in range(n-1):
        bary = np.mean(x[i:i+2])
        A[i:i+2][:,i:i+2] = A[i:i+2][:,i:i+2] + p(bary)*np.array([[1,-1],[-1,1]])/h[i]
        M[i:i+2][:,i:i+2] = M[i:i+2][:,i:i+2] + q(bary)*h[i]*np.array([[2,1],[1,2]])/6
        # use midpoint ('barycentric') one node quadrature
        b[i:i+2] = b[i:i+2] + f(bary)*h[i]/2
    # here we could implement boundary conditions as per the other codes.
    # instead, we do something different for fun - implement on one side only:
    A[0] = np.zeros((1,n))
    A[0][0] = 1
    M[0] = np.zeros((1,n))
    M[0][0] = 1
    b[0] = 0
    # now we can solve a discrete version of the eigenvalue problem, i.e.
    # find v such that -Av = lambda.M.v, for different lambdas
    # the exact solutions are lambda = n(n+1), for n=1,3,5,...
    # the corresponding eigenvectors approximate nth degree Legendre polynomials
    [w,v]=sp.linalg.eigh(-A[1:n][:,1:n],M[1:n][:,1:n],eigvals=(0,k-1))
    print(w)
    v = np.vstack([np.zeros((1,k)),v])
    return v

k=5
x = np.linspace(0,1,100)
v=U_eigs(x,k)   
plt.plot(x,v[:,0]/v[-1,0])
plt.plot(x,v[:,1]/v[-1,1])
plt.plot(x,v[:,2]/v[-1,2])