import pandas as pd
import matplotlib.pyplot as plt
import statsmodels.formula.api as smf
%matplotlib inline
plt.style.use('fivethirtyeight')

"""
Explanation: 1. Import the necessary packages to read in the data, plot, and create a linear regression model
End of explanation
"""


df = pd.read_csv("~/Documents/LEDE/algorithms/class6/data/hanford.csv")

df

"""
Explanation: 2. Read in the hanford.csv file
End of explanation
"""


df.describe()

"""
Explanation: 3. Calculate the basic descriptive statistics on the data
End of explanation
"""


r = df.corr()
r**2

df.plot(kind='scatter', x='Exposure', y='Mortality')

"""
Explanation: 4. Calculate the coefficient of correlation (r) and generate the scatter plot. Does there seem to be a correlation worthy of investigation?
End of explanation
"""


lm = smf.ols(formula='Mortality~Exposure',data=df).fit()

lm.params

intercept, slope = lm.params

"""
Explanation: 5. Create a linear regression model based on the available data to predict the mortality rate given a level of exposure
End of explanation
"""


df.plot(kind='scatter',x='Exposure',y='Mortality')
plt.plot(df['Exposure'],slope*df['Exposure']+intercept,"-")

"""
Explanation: 6. Plot the linear regression line on the scatter plot of values. Calculate the r^2 (coefficient of determination)
End of explanation
"""


def mortality_predictor(exposure):
    return slope*float(exposure)+intercept

mortality_predictor(10)

"""
Explanation: 7. Predict the mortality rate (Cancer per 100,000 man years) given an index of exposure = 10
End of explanation
"""


