# Newton's divided difference polynomials

import numpy as np
import matplotlib.pyplot as plt
from newtdd import newtdd
from nest import nest

a = -1
b =  1
N =  7

# function to interpolate
f = lambda x: np.cos(np.pi/2*x)

# points for interpolation
Xb = np.linspace(a,b,N)
Yb  = f(Xb)

# Newton's divided differences
c = newtdd(Xb,Yb)

fs = 24 #font size

# plot the interpolation points
plt.figure(1)
plt.plot(Xb,np.zeros_like(Xb),'bx',linewidth=3,markeredgewidth=3,markersize=10)
plt.grid(True)
plt.xlabel('x',fontsize=fs)
plt.ylabel('y',fontsize=fs)

# points for visualization
Xf = np.linspace(a,b,100)

# plot the function
plt.plot(Xf,f(Xf),'b--',linewidth=3,markeredgewidth=3,markersize=10)
plt.show(block=False)
# wait for a key press
_ = input("Press [enter] to continue.")

cm = plt.cm.gist_ncar
colors = [cm(i) for i in np.linspace(0,1,N+1)]


for i in range(N):
    #Newton's polynomial based on points 1:i
    Yf = nest(c[0:i+1],Xf,Xb[0:i+1])
    # plot it
    plt.plot(Xf,Yf,'-',Xb[0:i+1],Yb[0:i+1],'o',linewidth=3,markeredgewidth=3,markersize=10,color=colors[i])
    plt.show(block=False)
    plt.draw()
    # wait for a key press
    if i<N:
        _ = input("Press [enter] to continue.")

plt.show(block=True)