import matplotlib.pyplot as plt
import numpy as np
import matplotlib.ticker as tckr
from matplotlib.animation import FuncAnimation, PillowWriter

N=25


fig, (ax1, ax2) = plt.subplots(1,2)
#fig.set_tight_layout(True)


def handle_nodes(x, order, ax, prefix):
    y = lambda x : 1/(1+25*x*x);
    p = np.poly1d(np.polyfit(x,y(x),order))
    xd = np.linspace(-1,1,1000)
    ax.cla()
    ax.set_ylim(-5,5)
    ax.set_title(prefix + ' N=' + str(order))
    ax.plot(xd,p(xd))
    ax.plot(xd,y(xd))
    (markerLines) = ax.stem(x,y(x),linefmt=' ',basefmt=' ',markerfmt='ro')
    plt.setp(markerLines, markersize = 3)

def update(order):
    if order >=2:
        x = np.linspace(-1,1,order+1)
        handle_nodes(x, order, ax1, 'Equispaced')
        x = np.cos((2*np.arange(order+1)-1)*np.pi/(2*(order)))
        handle_nodes(x, order, ax2, 'Chebyshev')
    else:
        ax1.set_title('')
        ax2.set_title('')

anim = FuncAnimation(fig, update, frames=np.arange(0, N), interval=1000)
plt.show()
writergif = PillowWriter(fps = 2)
writergif.setup(fig, "animation3.gif", dpi = 400) 
fig.set_size_inches(12.8, 12.8, forward=True)
anim.save('animation4.gif', writer=writergif)