# chebPolys.py
# Andrew Davison, ad@coe.psu.ac.th, June 2026
'''
  Use Sympy's chebyshevt(n, x) to generate the first 6 T_n(x)
  Chebyshev polynomials. Each one is ploted between x = [-1,1].
  https://docs.sympy.org/latest/modules/functions/special.html#sympy.functions.special.polynomials.chebyshevt

  chebyshevIters() is another way to generate T_n(x), but
  by using sympy.expand() on its recursive definition.
'''

import sympy as sp
import matplotlib.pyplot as plt
import faUtils

def chebyshevIters(n):
  # the nth Chebyshev polynomial in x, T_n(x)
  if n == 0:
    return 1
  if n == 1:
    return x

  # implement T_n(x) = 2*x*T_{n-1}(x) - T_{n-2}(x)
  t0 = 1
  t1 = x
  for _ in range(2, n+1):
    t2 = sp.expand(2*x*t1 - t0)  
      # expand the polynomial expression
    t0 = t1
    t1 = t2
  return t1


# ----------------------------------------------------------
xs = faUtils.linspace(-1, 1, 1000)

x = sp.symbols('x')

print("Chebyshev Iterated Polys:")
for n in range(6):
  print(f"T{n}(x) = {chebyshevIters(n)}")
print()

print("Chebyshev Built-in Polys:")
for n in range(6):
  # the nth Chebyshev polynomial in x, T_n(x)
  cPoly = sp.chebyshevt(n, x)
  print(f"T{n}(x) = {cPoly}")

  # Convert to a Python function
  pFunc = sp.lambdify(x, cPoly, "math")

  ys = [pFunc(v) for v in xs]
     # evaluate sympy function returning a float
  plt.plot(xs, ys, label=f"T{n}(x)")

plt.axhline(0, color="black")
plt.axvline(0, color="black")
plt.title("Chebyshev Polynomials")
plt.xlabel("x")
plt.ylabel("Tn(x)")
plt.grid()
plt.legend()
plt.show()