import pandas as pd
%matplotlib inline
import numpy as np
from sklearn.linear_model import LogisticRegression
import matplotlib.pyplot as plt 

"""
Explanation: Apply logistic regression to categorize whether a county had high mortality rate due to contamination
1. Import the necessary packages to read in the data, plot, and create a logistic regression model
End of explanation
"""


df = pd.read_csv('../data/hanford.csv')

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


df.describe()

"""
Explanation: <img src="../../images/hanford_variables.png"></img>
3. Calculate the basic descriptive statistics on the data
End of explanation
"""


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

df['Mortality'].mean() + 3 * df['Mortality'].std()

df['Exposure'].mean() + 3 * df['Exposure'].std()

high_exposure = df['Exposure'].median()

high_mortality = df['Mortality'].median()

df['Mort_high'] = df['Mortality'].apply(lambda x:1 if x >= high_mortality else 0)
df['Expo_high'] = df['Mortality'].apply(lambda x:1 if x >= high_mortality else 0)

df

"""
Explanation: 4. Find a reasonable threshold to say exposure is high and recode the data
End of explanation
"""


lm = LogisticRegression()
data = np.asarray(df[['Mort_high', 'Expo_high']]) #data is a list of lists
x = data[:,1:] # the second colon is in case you have more than one input in your prediction model
y = data[:,0]
lm.fit(x,y)
lm.score(x,y)  #This gives you r^2, or coefficient of determination
slope = lm.coef_[0] #the coefficients are a list (here, we’re just using one) of the slopes for each x value
intercept = lm.intercept_

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

"""
Explanation: 5. Create a logistic regression model
End of explanation
"""


lm.predict(50)

"""
Explanation: 6. Predict whether the mortality rate (Cancer per 100,000 man years) will be high at an exposure level of 50
End of explanation
"""


