# File: polyInterp.py
# Author: Ian May
# Purpose: Implement a polynomial interpolation class as a way to demonstrate OOP
# Notes: This is very much *not* the most efficient way to implement this

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

# Base polynomial interpolation class
# Constructor takes in arbitrary distribution of nodes
class PolyInterp:
    def __init__(self, nodes):
        # Copy in node positions
        self.nodes = np.array(nodes)
        # Set degree, number of nodes, and make space for coeffs
        self.N = nodes.size
        self.deg = self.N - 1
        self.coeff = np.zeros(nodes.size)
        # Build vandermonde matrix
        self.V = np.zeros((self.N,self.N))
        for j in range(0,self.N):
            self.V[:,j] = self.nodes**j
            
    # Generate interpolant given data set
    def interp(self, data):
        self.coeff = np.linalg.solve(self.V,data)
        
    # Evaluate interpolant at a set of point(s)
    def evalInt(self, points):
        interp = np.zeros(points.shape)
        for j in range(0,self.N):
            interp = interp + self.coeff[j]*points**j
        return interp

# Runge's function to test our interpolants on
def runge(x):
    return 1.0/(1.0 + 25.0*x**2)

# Derivative of Runge's function to test our interpolants on
def runge_der(x):
    return -50.0*x/((1.0 + 25.0*x**2)**2)

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)

    # Report maximum error
    print('Max error for equispaced: ',np.max(np.abs(exact-eq_interp)))

    # Create plot, cap y-limit for visibility
    plt.plot(points,exact,'-k',points,eq_interp,'-r')
    plt.plot(equi.nodes,runge(equi.nodes),'ro')
    plt.legend(('Exact','Equispaced'),loc='best')
    plt.title("Equispaced interpolation of Runge's function")
    plt.grid('both')
    plt.ylim([-0.1,1.1])
    plt.show()