# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# +
import numpy as np
from geopy import *
import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.model_selection import train_test_split
from sklearn.cluster import *
import matplotlib.pyplot as plt
import seaborn as sns

import plotly.graph_objects as go
import plotly.express as px
from plotly.subplots import make_subplots
from scipy.stats import pearsonr
import re

cases = pd.read_csv('data/time_series_covid19_confirmed_US.csv') # https://github.com/CSSEGISandData/COVID-19/blob/master/csse_covid_19_data/csse_covid_19_time_series/time_series_covid19_confirmed_US.csv
vaccinations = pd.read_csv('data/people_vaccinated_us_timeline.csv') # https://raw.githubusercontent.com/govex/COVID-19/master/data_tables/vaccine_data/us_data/time_series/people_vaccinated_us_timeline.csv
counties = pd.read_csv('data/co-est2020.csv', encoding='latin-1') # https://www2.census.gov/programs-surveys/popest/datasets/2010-2020/counties/totals/co-est2020.csv
mask_use = pd.read_csv('data/mask-use-by-county.csv') # https://github.com/nytimes/covid-19-data/blob/master/mask-use/mask-use-by-county.csv
# -

cases = cases.transform(lambda x: x.fillna('') if x.dtype == 'object' else x.fillna(0))
vaccinations = vaccinations.transform(lambda x: x.fillna('') if x.dtype == 'object' else x.fillna(0))
mask_use = mask_use.transform(lambda x: x.fillna('') if x.dtype == 'object' else x.fillna(0))
counties = counties.transform(lambda x: x.fillna('') if x.dtype == 'object' else x.fillna(0))

counties['FIPS'] = counties['STATE']*1000 + counties['COUNTY']
counties['FIPS'] = counties['FIPS'].apply(lambda x: str(int(x)).zfill(5))
counties

cases['FIPS_1'] = cases['FIPS'].apply(lambda x: str(int(x)).zfill(5))
mask_use['FIPS_2'] = mask_use['COUNTYFP'].apply(lambda x: str(x).zfill(5))
county_data = counties.merge(cases,right_on = 'FIPS_1', left_on = 'FIPS').merge(mask_use,left_on = 'FIPS_1', right_on = 'FIPS_2').drop(columns = ['FIPS_x','FIPS_2','FIPS_y']).rename(columns = {'FIPS_1':'FIPS'})
mask_use = mask_use.drop(columns = ['FIPS_2'])
cases = cases.drop(columns = ['FIPS_1'])
county_data

# +
data_5a = county_data.loc[:,'9/12/21':]
data_5a.drop(columns=['FIPS', 'COUNTYFP'], inplace=True)
sns.heatmap(data_5a.corr(), cmap="YlGnBu", annot=True)
plt.title("pairwise correlation heatmap")

plt.show()

# +
from sklearn.metrics import mean_squared_error
X_q5b = county_data[['NEVER', 'RARELY', 'SOMETIMES', 'FREQUENTLY', 'ALWAYS']]
y_q5b = county_data['9/12/21']

X_q5b_train, X_q5b_test, y_q5b_train, y_q5b_test = train_test_split(X_q5b, y_q5b, test_size=0.33, random_state=42)

reg = LinearRegression()
reg.fit(X_q5b_train, y_q5b_train)

y_pred_test = reg.predict(X_q5b_test)
y_pred_train = reg.predict(X_q5b_train)

mse_test = mean_squared_error(y_pred_test, y_q5b_test)
mse_train = mean_squared_error(y_pred_train, y_q5b_train)

test_rmse_cases = np.sqrt(mse_test)
train_rmse_cases = np.sqrt(mse_train)

train_rmse_cases, test_rmse_cases

# +
X_q5d = county_data[['NEVER', 'RARELY', 'SOMETIMES', 'FREQUENTLY', 'ALWAYS']]
y_q5d = county_data['9/12/21'] / county_data['POPESTIMATE2020']

X_q5d_train, X_q5d_test, y_q5d_train, y_q5d_test = train_test_split(X_q5d, y_q5d, test_size=0.33, random_state=42)

reg = LinearRegression()
reg.fit(X_q5b_train, y_q5d_train)

y_pred_test = reg.predict(X_q5d_test)
y_pred_train = reg.predict(X_q5d_train)

mse_test = mean_squared_error(y_pred_test, y_q5d_test)
mse_train = mean_squared_error(y_pred_train, y_q5d_train)

test_rmse_cpc = np.sqrt(mse_test)
train_rmse_cpc = np.sqrt(mse_train)

train_rmse_cpc, test_rmse_cpc

# +
plt.figure(figsize = (10,7))

sns.regplot(x = y_pred_test, y = y_q5d_test, scatter_kws={'s':5})
plt.title("Predictions versus observations for Q5d")
plt.xlabel("Predictions")
plt.ylabel("Observations");
plt.show()

# +
models = []

q5f_Xy = X_q5d_train.copy()
q5f_Xy["cases"] = y_q5d_train

for _ in range(1000):
    temp = q5f_Xy.sample(frac=1.0, replace = True) #with replacement
    tempX = temp.loc[:,:"ALWAYS"]
    tempY = temp["cases"]
    temp_reg = LinearRegression()
    temp_reg.fit(tempX, tempY)
    models.append(temp_reg)

# +
y_predictions = []

for model in models:
    y_pred = model.predict(pd.DataFrame(X_q5d_test.iloc[100]).transpose())
    y_predictions.append(y_pred)
    
variance = np.var(y_predictions)

mse = mean_squared_error(y_predictions, np.array([y_q5d_test.iloc[100]]*1000))

prop_var = variance / mse

prop_var

# +
vote = pd.read_csv('presidential_elections.csv')[['State','2020']].head(51)
vote = vote.replace('D.C.','District of Columbia')

daily_diff = vaccinations[['People_Fully_Vaccinated','People_Partially_Vaccinated']].diff()
daily_vacc = vaccinations[['Province_State','Date']].join(daily_diff) 
daily_vacc = daily_vacc.merge(county_data[['STNAME','POPESTIMATE2020']].groupby('STNAME', as_index = False).agg(sum), 
                 left_on = 'Province_State', right_on= 'STNAME', how = 'left').drop(columns = 'STNAME')
daily_vacc['People_Fully_Vaccinated'] = daily_vacc['People_Fully_Vaccinated']/daily_vacc['POPESTIMATE2020']
daily_vacc['People_Partially_Vaccinated'] = daily_vacc['People_Partially_Vaccinated']/daily_vacc['POPESTIMATE2020']
daily_vacc = daily_vacc.merge(vote, left_on = 'Province_State', right_on = 'State', how = 'right')
daily_vacc['2020'] = daily_vacc['2020'].replace('R',1).replace('D',0)
daily_vacc = daily_vacc[daily_vacc['Date'] > '2020-12-10'].rename(columns = {'2020':'party'}).drop(columns = 'State')
daily_vacc['Date'] = pd.to_datetime(daily_vacc['Date'], format = '%Y-%m-%d')
daily_vacc

# +
#df contains daily cases increase
df = cases[['Province_State']].join(cases.iloc[:, 11:].diff(axis=1, periods=1)).groupby('Province_State', as_index = False).agg(sum)
df= df.transpose()
df.columns = df.iloc[0]
df = df.drop(df.index[0]).reset_index()
df = df.melt(id_vars = 'index', value_vars = vote['State'].values)
df = df.rename(columns = {'index':'Date', 'value': 'daily_case'})
df['Date'] = pd.to_datetime(df['Date'], format = '%m/%d/%y')

#merge df (cases data) with vacc data & time (275 days)
all_features = daily_vacc.merge(df, on = ['Province_State','Date'])
all_features['Date'] = pd.to_datetime(all_features['Date'], format = '%Y-%m-%d')
all_features['daily_case'] = all_features['daily_case']/all_features['POPESTIMATE2020']
# -

#train-test split for time series, this creates a 90/10 train-validation split 
X = all_features[['Date','People_Fully_Vaccinated','People_Partially_Vaccinated','daily_case']]
X['Date'] = pd.to_datetime(X['Date'], format = '%Y-%m-%d')
Y = all_features[['Date','party']]
Y['Date'] = pd.to_datetime(Y['Date'], format = '%Y-%m-%d')
X_train = X[X['Date']<= '2021-08-15'].set_index('Date')
X_test = X[X['Date']> '2021-08-15'].set_index('Date')
Y_train = Y[Y['Date']<= '2021-08-15'].set_index('Date')
Y_test = Y[Y['Date']> '2021-08-15'].set_index('Date')

# +
from sklearn.linear_model import LogisticRegression

model = LogisticRegression(fit_intercept = True, solver = 'lbfgs')
model.fit(X_train, Y_train)

training_accuracy = model.score(X_train, Y_train)
print("Training Accuracy: ", training_accuracy)

# +

#imported daily deaths data and computed the daily difference of deaths 
deaths = pd.read_csv('data/deaths.csv')
deaths = deaths.sort_values(by = ['state','date']).reset_index()[['date','state','deaths']]
deaths['date'] = pd.to_datetime(deaths['date'], format = '%Y-%m-%d')
deaths['deaths_diff'] = deaths['deaths'].diff()

#dictionary to change abbreviations
us_state_to_abbrev = { "Alabama": "AL", "Alaska": "AK", "Arizona": "AZ", "Arkansas": "AR", "California": "CA", "Colorado": "CO", "Connecticut": "CT", "Delaware": "DE", "Florida": "FL", "Georgia": "GA", "Hawaii": "HI", "Idaho": "ID", "Illinois": "IL", "Indiana": "IN", "Iowa": "IA", "Kansas": "KS", "Kentucky": "KY", "Louisiana": "LA", "Maine": "ME", "Maryland": "MD", "Massachusetts": "MA", "Michigan": "MI", "Minnesota": "MN", "Mississippi": "MS", "Missouri": "MO", "Montana": "MT", "Nebraska": "NE", "Nevada": "NV", "New Hampshire": "NH", "New Jersey": "NJ", "New Mexico": "NM", "New York": "NY", "North Carolina": "NC", "North Dakota": "ND", "Ohio": "OH", "Oklahoma": "OK", "Oregon": "OR", "Pennsylvania": "PA", "Rhode Island": "RI", "South Carolina": "SC", "South Dakota": "SD", "Tennessee": "TN", "Texas": "TX", "Utah": "UT", "Vermont": "VT", "Virginia": "VA", "Washington": "WA", "West Virginia": "WV", "Wisconsin": "WI", "Wyoming": "WY", "District of Columbia": "DC", "American Samoa": "AS", "Guam": "GU", "Northern Mariana Islands": "MP", "Puerto Rico": "PR", "United States Minor Outlying Islands": "UM", "U.S. Virgin Islands": "VI", }
us_state_to_abbrev = {y:x for x,y in us_state_to_abbrev.items()}

#imported daily test data and computed the daily difference of tests
tests = pd.read_csv('data/time_series_covid19_US.csv')
tests['state'] = tests[['state']].replace({'state':us_state_to_abbrev})
tests['Date'] = pd.to_datetime(tests['date'], format = '%m/%d/%Y')
tests = tests[['Date','state','tests_combined_total']]
tests = tests.sort_values(by = ['state','Date'])
tests['tests_diff'] = tests['tests_combined_total'].diff()

#merged new features: deaths & tests data to the existing dataframe with original three features
all_features_6b = all_features.merge(deaths, left_on = ['Date','Province_State'], right_on = ['date','state'], how = 'left')
all_features_6b['Date'] = pd.to_datetime(all_features_6b['Date'], format = '%Y-%m-%d')
all_features_6b['deaths_diff_per_capita'] = all_features_6b['deaths_diff']/ all_features_6b['POPESTIMATE2020']
# all_features.drop(columns = ['date','state','deaths'])
# all_features = all_features[['Province_State','Date','People_Fully_Vaccinated','People_Partially_Vaccinated','POPESTIMATE2020','daily_case','deaths_diff_per_capita','party']]
all_features_6b = all_features_6b.merge(tests, left_on = ['Date','Province_State'], right_on = ['Date','state'], how = 'left')
all_features_6b['tests_diff_per_capita'] = all_features_6b['tests_diff']/ all_features_6b['POPESTIMATE2020']
all_features_6b = all_features_6b[['Province_State','Date','People_Fully_Vaccinated','People_Partially_Vaccinated','POPESTIMATE2020','daily_case','deaths_diff_per_capita','tests_diff_per_capita','party']]


# +
import plotly.express as px
plot_6b = all_features_6b.copy().groupby(['Date','party'], as_index = False).agg(np.mean)
plot_6b = plot_6b.merge(all_features_6b[['Date','party','daily_case']].groupby(['Date','party'], as_index = False).agg(np.mean))
fig = px.line(plot_6b, x='Date', y='People_Fully_Vaccinated', color = 'party')
fig.update_layout(
    width=700,
    title = {'text': "Comparison of Parties over Time : Average of Fully Vaccinated People Per Capita"},
    height=700)
fig.show()

fig = px.line(plot_6b, x='Date', y='People_Partially_Vaccinated', color = 'party')
fig.update_layout(
    width=700,
    title = {'text': "Comparison of Parties over Time : Average of Partially Vaccinated People Per Capita"},
    height=700)
fig.show()

fig = px.line(plot_6b, x='Date', y='daily_case', color = 'party')
fig.update_layout(
    width=700,
    title = {'text': "Comparison of Parties over Time : Average of Daily New Cases Per Capita"},
    height=700)
fig.show()

fig = px.line(plot_6b, x='Date', y='deaths_diff_per_capita', color = 'party')
fig.update_layout(
    width=700,
    title = {'text': "Comparison of Parties over Time : Average of Daily New Deaths Per Capita"},
    height=700)
fig.show()

fig = px.line(plot_6b, x='Date', y='tests_diff_per_capita', color = 'party')
fig.update_layout(
    width=700,
    title = {'text': "Comparison of Parties over Time : Average of Daily New Tests Per Capita"},
    height=700)
fig.show()
# -

#train-test split for time series, this creates a 70/30 train-validation split 
X_6b3 = all_features_6b[['Date','People_Fully_Vaccinated','People_Partially_Vaccinated','daily_case','deaths_diff_per_capita','tests_diff_per_capita']]
Y_6b3 = all_features_6b[['Date','party']]
X_train_6b3 = X_6b3[X_6b3['Date']>= '2021-03-04'].set_index('Date')
X_test_6b3 = X_6b3[X_6b3['Date'] <'2021-03-04'].set_index('Date')
Y_train_6b3 = Y_6b3[Y_6b3['Date']>= '2021-03-04'].set_index('Date')
Y_test_6b3 = Y_6b3[Y_6b3['Date']< '2021-03-04'].set_index('Date')

# +
from sklearn.linear_model import LogisticRegression
model_6b3 = LogisticRegression(penalty = 'none', fit_intercept = True, solver = 'lbfgs')
model_6b3.fit(X_train_6b3, Y_train_6b3)

training_accuracy = model_6b3.score(X_train_6b3, Y_train_6b3)
print("Training Accuracy: ", training_accuracy)
testing_accuracy = model_6b3.score(X_test_6b3, Y_test_6b3)
print("Testing Accuracy: ", testing_accuracy)
# -

# 5-fold cross validation 
from sklearn.linear_model import LogisticRegressionCV
model_kfold = LogisticRegressionCV(cv=5,random_state = 0)
model_kfold.fit(X_6b3.set_index('Date'), Y_6b3.set_index('Date'))
acc = model_kfold.score(X_6b3.set_index('Date'), Y_6b3.set_index('Date'))
print("Model Accuracy: ", acc)

y_pred_train = model_kfold.predict(X_6b3.set_index('Date'))
tp = np.sum((y_pred_train == 1) & (Y_6b3['party'] == 1))
tn = np.sum((y_pred_train == 0) & (Y_6b3['party'] == 0))
fp = np.sum((y_pred_train == 1) & (Y_6b3['party'] == 0))
fn = np.sum((y_pred_train == 0) & (Y_6b3['party'] == 1))
precision = tp / (tp + fp)
recall = tp / (tp + fn)
print(precision, recall)

# +
from sklearn.metrics import roc_curve
fpr, tpr, threshold = roc_curve(Y_6b3.set_index('Date'), model_kfold.predict_proba(X_6b3.set_index('Date'))[:, 1])
fig = px.line(x=fpr, y = tpr, hover_name=threshold)
fig.update_xaxes(title="False Positive Rate")
fig.update_yaxes(title="True Positive Rate")
fig.update_layout(
    width=500,
    title={
    'text': "ROC Curve Analysis of the Improved Model",
    'y':0.95,
    'x':0.5},
    height=500)
fig.show()

fpr, tpr, threshold = roc_curve(Y_train, model.predict_proba(X_train)[:, 1])
fig = px.line(x=fpr, y = tpr, hover_name=threshold)
fig.update_xaxes(title="False Positive Rate")
fig.update_yaxes(title="True Positive Rate")
fig.update_layout(
    width=500,
    title={
    'text': "ROC Curve Analysis of the Baseline Model",
    'y':0.95,
    'x':0.5},
    height=500)
fig.show()

# +
from datetime import timedelta
import datetime

end_date = pd.to_datetime("2021-04-07", format = '%Y-%m-%d')

# function tries different k values
def search_k(k, model_type, x_set, y_set): # k: 2
    for d in range(0,k): # 2 iterations where last iteration predicts the day after tmrw 
        next_date = end_date + timedelta(days = d) # increment end_date by 1 day
        x_train = x_set[x_set['Date'] < next_date].set_index('Date')
        y_train = y_set[y_set['Date'] < next_date].set_index('Date')
        tmrw = x_set[x_set['Date'] == next_date].set_index('Date')
        tmrw_true = y_set[y_set['Date'] == next_date][['party']].values
        model_type.fit(x_train, y_train)
        y_pred_array = model_type.predict(tmrw)
        y_pred = pd.DataFrame(data=y_pred_array, index = [next_date]*51, columns=["party"])
        y_train = y_train.append(y_pred)
    return y_train, x_set[x_set['Date'] <= next_date].set_index('Date')


# +
from sklearn.metrics import roc_curve, roc_auc_score

result_table_6d = pd.DataFrame(columns=['model_type','k_values','fpr','tpr','auc'])
k_values = [2,14]
# Train the models and record the results
for k in k_values:
        y, x = search_k(k,model_kfold, X_6b3, Y_6b3)
        y_prob = model_kfold.predict_proba(x)[:,1]
        auc = roc_auc_score(y, y_prob)
        fpr, tpr, threshold = roc_curve(y, model_kfold.predict_proba(x)[:, 1])
        result_table_6d = result_table_6d.append({'model_type':'improved_model',
                                            'k_values' : k,
                                            'fpr' : fpr,
                                            'tpr' : tpr,
                                            'auc':auc}, ignore_index=True)
for k in k_values:
        y, x = search_k(k,model, X, Y)
        y_prob = model.predict_proba(x)[:,1]
        auc = roc_auc_score(y, y_prob)
        fpr, tpr, threshold = roc_curve(y, model.predict_proba(x)[:, 1])
        result_table_6d = result_table_6d.append({'model_type':'baseline_model',
                                            'k_values' : k,
                                            'fpr' : fpr,
                                            'tpr' : tpr,
                                            'auc':auc}, ignore_index=True)    

# +
fig = plt.figure(figsize=(8,6))

for i in result_table_6d.index:
    plt.plot(result_table_6d.loc[i]['fpr'], 
             result_table_6d.loc[i]['tpr'], 
             label= result_table_6d.loc[i]['model_type'][:8] + ", " + str(result_table_6d.loc[i]['k_values']) + " days, " + "AUC={:.3f}".format(result_table_6d.loc[i]['auc']))

plt.xticks(np.arange(0.0, 1.1, step=0.1))
plt.xlabel("False Positive Rate", fontsize=15)
plt.yticks(np.arange(0.0, 1.1, step=0.1))
plt.ylabel("True Positive Rate", fontsize=15)

plt.title('ROC Curve Analysis of Improved vs Baseline Model', fontsize=15)
plt.legend(prop={'size':13}, loc='lower right')

plt.show()

# +
from sklearn.metrics import roc_curve, roc_auc_score

result_table = pd.DataFrame(columns=['model_type','k_values','auc'])

# Train the models and record the results
for k in range(1,51,5):
        y, x = search_k(k,model_kfold, X_6b3, Y_6b3)
        y_prob = model_kfold.predict_proba(x)[:,1]
        auc = roc_auc_score(y, y_prob)
        result_table = result_table.append({'model_type':'model_kfold',
                                            'k_values' : k,
                                            'auc':auc}, ignore_index=True)
for k in range(1,51,5):
        y, x = search_k(k,model, X, Y)
        y_prob = model.predict_proba(x)[:,1]
        auc = roc_auc_score(y, y_prob)
        result_table = result_table.append({'model_type':'model',
                                            'k_values' : k,
                                            'auc':auc}, ignore_index=True)  

# +
result_table = result_table.replace('model_kfold', 'improved_model').replace('model', 'baseline_model')
fig = px.line(result_table, x='k_values', y='auc', color='model_type', markers=True)

fig.update_xaxes(title="Experimented K Values in Days")
fig.update_yaxes(title="AUC Score")
fig.update_layout(
    width=700,
    title={
    'text': "Improved Model vs Baseline Model Predictions for 10 Different K Values",
    'y':0.95,
    'x':0.5},
    height=600)

fig.show()


# +
def search_k_6f(m, model_type, x_set, y_set): # k: 2
    for i in range(0,10): # 2 iterations where last iteration predicts the day after tmrw 
        x_train = x_set[x_set['Date'] < end_date].set_index('Date')
        y_train = y_set[y_set['Date'] < end_date].set_index('Date')
        m_days = x_set[x_set['Date'] >= end_date][x_set['Date'] <= end_date + timedelta(days = m)].set_index('Date')
        model_type.fit(x_train, y_train)
        y_pred_array = model_type.predict(m_days)
        y_pred = pd.DataFrame(data=y_pred_array, index = [end_date]*51*(m+1), columns=["party"])
        y_train = y_train.append(y_pred)
    return y_train, x_set[x_set['Date'] <= end_date + timedelta(days = m)].set_index('Date')

# Define a result table as a DataFrame
result_table_6f = pd.DataFrame(columns=['model_type','m_values','fpr','tpr','auc'])
m_values = [5,10]
# Train the models and record the results
for m in m_values:
        y, x = search_k_6f(m,model_kfold, X_6b3, Y_6b3)
        y_prob = model_kfold.predict_proba(x)[:,1]
        auc = roc_auc_score(y, y_prob)
        fpr, tpr, threshold = roc_curve(y, model_kfold.predict_proba(x)[:, 1])
        result_table_6f = result_table_6f.append({'model_type':'improved_model',
                                            'm_values' : m,
                                            'fpr' : fpr,
                                            'tpr' : tpr,
                                            'auc':auc}, ignore_index=True)
for m in m_values:
        y, x = search_k_6f(m,model, X, Y)
        y_prob = model.predict_proba(x)[:,1]
        auc = roc_auc_score(y, y_prob)
        fpr, tpr, threshold = roc_curve(y, model.predict_proba(x)[:, 1])
        result_table_6f = result_table_6f.append({'model_type':'baseline_model',
                                            'm_values' : m,
                                            'fpr' : fpr,
                                            'tpr' : tpr,
                                            'auc':auc}, ignore_index=True)  

# +
fig = plt.figure(figsize=(8,6))

for i in result_table_6f.index:
    plt.plot(result_table_6f.loc[i]['fpr'], 
             result_table_6f.loc[i]['tpr'], 
             label= result_table_6f.loc[i]['model_type'][:8] + ", " + str(result_table_6f.loc[i]['m_values']) + " days, " + "AUC={:.3f}".format(result_table_6f.loc[i]['auc']))

plt.xticks(np.arange(0.0, 1.1, step=0.1))
plt.xlabel("False Positive Rate", fontsize=15)
plt.yticks(np.arange(0.0, 1.1, step=0.1))
plt.ylabel("True Positive Rate", fontsize=15)

plt.title('ROC Curve Analysis of Improved vs Baseline Model', fontsize=15)
plt.legend(prop={'size':13}, loc='lower right')

plt.show()

StarcoderdataJupyterScriptsDedupFiltered