"""
Exercise 5.18 from "A primer on..."
Write a function that fits polynomials to 
input data x and y, using the Numpy functions
polyfit and poly1d. Then apply the function
to the data from ex 5.16. Two points worth noting:
- We reuse the function from ex 5.16 by importing it 
- The last argument to the function fit(x, y, deg) is supposed
to be alist of integer values, and the fitting is to be performed for
all these values. This is accomplished with a for-loop over the
elements in deg, repeating the fitting and plotting for each element. 
"""


import numpy as np
import matplotlib.pyplot as plt
from read_density_data import read_density_data


def fit(x, y, deg):
    for d in deg:
        coeff = np.polyfit(x, y, d)
        p = np.poly1d(coeff)
        y_fit = p(x)
        plt.plot(x, y_fit, label = f'd={d}')
    plt.plot(x, y, 'ro', label = 'data')
    plt.legend()
    plt.show()


#fit and plot for both the density datasets from ex. 5.16:
temp, dens = read_density_data('density_air.txt')
fit(temp, dens, [1,2])

temp, dens = read_density_data('density_water.txt')
fit(temp, dens, [1,2])