# File: ChebyInterp.py
# Author: Ian May
# Purpose: Specialize the PolyInterp class to Chebyshev interpolation
#          to demonstrate inheritance
# Notes: This is very much *not* the most efficient way to implement this

import sys
import numpy as np
import matplotlib.pyplot as plt

from PolyInterp import PolyInterp, runge, runge_der

# Chebyshev polynomial interpolation, specializes PolyInterp
# Constructor requires number of points desired
class ChebyInterp(PolyInterp):
    def __init__(self, N):
        # Invoke PolyInterp constructor using Chebyshev nodes
        k = np.arange(1,2*N,2) # Odd nums up to 2*N
        PolyInterp.__init__(self, np.cos(np.pi*k/(2.0*N)))

if __name__ == "__main__":
    # Number of nodes to use (deg+1)
    N = 9
    if len(sys.argv)>1:
        N = int(sys.argv[1])
        
    # Points to plot on
    points = np.linspace(-1,1,300)
    exact = runge(points)
    
    # Construct interpolant using equispaced nodes
    equi = PolyInterp(np.linspace(-1,1,N))
    equi.interp(runge(equi.nodes))
    eq_interp = equi.evalInt(points)
    
    # Construct interpolant using Chebyshev nodes
    cheb = ChebyInterp(N)
    cheb.interp(runge(cheb.nodes))
    cv_interp = cheb.evalInt(points) # function inherited from base class
    
    # Report maximum error
    print('Max error for equispaced: ',np.max(np.abs(exact-eq_interp)))
    print('Max error for Chebyshev: ',np.max(np.abs(exact-cv_interp)))
    
    # Create plot, cap y-limit for visibility
    plt.plot(points,exact,'-k',points,eq_interp,'-r',points,cv_interp,'-b')
    plt.plot(equi.nodes,runge(equi.nodes),'ro',cheb.nodes,runge(cheb.nodes),'bs')
    plt.legend(('Exact','Equispaced','Chebyshev'),loc='best')
    plt.title('Comparison of equispaced and Chebyshev interpolation')
    plt.grid('both')
    plt.ylim([-0.1,1.1])
    plt.show()