import pickle

# Lecture de la base de données.
#
# Il faudra vraisemblablement adapter le chemin vers le fichier wooldridge.pkl.

fid = open('C:/Users/claire.loupias/Desktop/wooldridge23.pkl', 'rb')
alldatasets = pickle.load(fid)
fid.close()


dataset = alldatasets['401k']
#
# Question 1.
#

prate = dataset.prate
mrate = dataset.mrate

mprate = prate.mean()
mmrate = mrate.mean()

print('Le taux de participation moyen est : '+f'{mprate:.2f}')
print('Contribution moyenne des employeurs : '+f'{mmrate:.2f}')


#
# Question 2.
#

import statsmodels.formula.api as smf

results = smf.ols('prate ~ mrate', data=dataset).fit()

print(results.summary())


#
# Question 4.
#

beta = results.params
pred = beta[0]+3.5*beta[1]

print('Le taux de participation moyen si mrate=3,5 est '+f'{pred:.2f}')

import matplotlib.pyplot as plt

plt.plot(mrate, prate, 'k.', label='data', markersize=.5)
plt.plot(mrate, beta[0]+beta[1]*mrate, 'r-', label="OLS fit", linewidth=2)
plt.legend(loc='best')
plt.show()

# Pour imprimer le graphique dans un fichier tex. Installer au préalable le module 
# tikzplotlib (`pip install tikzplotlib`).
#
#import tikzplotlib
#
#tikzplotlib.save("exercice-2.1.tex")