import numpy as np
import matplotlib.pyplot as plt

from numpy.polynomial.legendre import legval, leggauss

#
# first, Legendre polynomial approximation
#

# Gauss–Legendre nodes and weights
x, w = leggauss(200)

# Function to approximate
y = np.sin(0.5 * np.pi * x)

# First plot
fig = plt.figure()
plt.plot(x, y)
plt.title('y = sin(π*x/2)')
plt.show()

# initialise approximating polynomial to zero
p = 0*x

# Loop over odd k = 1,3,5,7
for k in range(1, 8, 2):

    # Legendre polynomial P_k evaluated at x
    coeffs = np.zeros(k+1)
    coeffs[k] = 1.0
    Pk = legval(x, coeffs)

    # L2 projection coefficient:
    # <f,Pk> / <Pk,Pk>
    a = np.sum(w * y * Pk) / np.sum(w * Pk**2)

    p = p + a * Pk

    fig = plt.figure()
    plt.plot(x, y-p)
    plt.title(f'error of degree {k} Legendre polynomial approximation')
    plt.show()

# second, Chebyshev polynomial approximation

th = np.linspace(0, np.pi, 200)
th = 0.5 * (th[1:] + th[:-1])

x = np.cos(th)
y = np.sin(0.5 * np.pi * x)

# First figure
fig = plt.figure()
plt.plot(x, y)
plt.title('y = sin(π*x/2)')
plt.show()

# initialise approximating polynomial to zero
p = 0*x

# Loop over odd k = 1,3,5,7
for k in range(1, 8, 2):
    a = np.sum(np.cos(k * th) * y) / np.sum(np.cos(k * th)**2)
    p = p + a * np.cos(k * th)

    fig = plt.figure()
    plt.plot(x, y-p)
    plt.title(f'error of degree {k} Chebyshev polynomial approximation')
    plt.show()