"""
MATH3476 Example 1.1
Lagrange interpolation through 3 nodes
"""
import numpy as np
import matplotlib.pyplot as plt

n = 2
xn = np.array([-1., 0., 1.])   # n+1 nodes
uu = np.array([2.,  1., 3.])   # function values

# set up the x-axis
n_pt = 500
x = np.linspace(np.amin(xn),np.amax(xn),n_pt)

# construct the polynomials in (1.2)
L = np.ones((n+1,n_pt))   # initialise the array L
for i in range(n+1):
    for j in range(n+1):
        if j != i:
           L[i,:] *= (x - xn[j])/(xn[i] - xn[j])
            
# compute the interpolation
p = np.dot(uu,L)   # use dot() instead of a for-loop

# plotting
plt.figure(1); plt.clf()
plt.plot(x,L[0,:],'r--',label='$\ell_0(x)$')
plt.plot(x,L[1,:],'g--',label='$\ell_1(x)$')
plt.plot(x,L[2,:],'b--',label='$\ell_2(x)$')
plt.plot(xn,uu,'ko',label='date points')
plt.plot(x,p,'k',label='interpolation $p_2(x)$')

plt.xlabel('$x$')
plt.legend()
plt.grid(True)

plt.savefig('lagrange.pdf',dpi=600,bbox_inches='tight')