# lagSymPop.py
# Andrew Davison, ad@coe.psu.ac.th, June 2026
'''
Code for the US population question at the end of
the Function Approximations chapter.

The Lagrange goes off the rails very quickly once it's
used for estimation. For example, for 2024, the 
approximation is too large by 12%.
'''


import math

import sympy as sp
from sympy.polys.specialpolys import interpolating_poly

import matplotlib.pyplot as plt
import faUtils

x = sp.symbols('x')

# Interpolation nodes
xs = [1950, 1960, 1970, 1980, 1990, 2000, 2010, 2020]
ys = [ 151325798, 179323175, 203211926, 226545805, 
       248709873, 281421906, 308745538, 331449281]

# Lagrange polynomial
ipoly = interpolating_poly(len(xs), x, X=xs, Y=ys)
p = sp.expand(ipoly)   # multiply out all the terms
print("Expanded:")
print(p); print()

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

print("Estimates Test")
print("Year    Estimate    Approx       RelError")
for est in [ (2021, 332031554), (2022, 333287557), 
             (2023, 336806231), (2024, 340110988) ]:
  approx = pFunc(est[0])
  relErr = abs(est[1] - approx)/est[1]
  print(f"{est[0]}:   {est[1]}   {int(approx)}    {relErr:6f}")


# Generate plot data
xVals = [y for y in range(1950, 2041)]
polyVals = [pFunc(v) for v in xVals]

plt.figure(figsize=(8, 5))
plt.plot(xVals, polyVals, label="Lagrange")
plt.plot(xs, ys, "o", label="Data")

ax = plt.gca()
ax.set_xlim([1940, 2040])

plt.xlabel("year")
plt.ylabel("pop")
plt.grid(True)
plt.legend()
plt.title("US Population by Year")
plt.show()