import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

df = pd.read_csv('https://jonghank.github.io/ase1302/files/owid-covid-data-nov2021-noindex.csv', index_col=3)
df

# your code here

location_set = set(df['location'])
number_of_countries = len(location_set)
location_list = sorted(list(location_set))

print(f'number of countries: {number_of_countries}')

count = 0
for country in location_list:
  print(country, end=' | ')
  count += 1
  if count%6==0:
    print()
print('\n')


# your code here

kor_df = df[df['location']=='South Korea']

plt.figure(figsize=(15,3), dpi=100)
kor_df['new_cases'].plot()
plt.legend()
plt.title('New confirmed cases in Korea')
plt.grid()
plt.show()

plt.figure(figsize=(15,3), dpi=100)
kor_df['new_deaths'].plot()
plt.legend()
plt.title('New deaths in Korea')
plt.grid()
plt.show()

plt.figure(figsize=(15,3), dpi=100)
kor_df['total_cases'].plot()
plt.title('Total confirmed cases in Korea')
plt.grid()
plt.show()

plt.figure(figsize=(15,3), dpi=100)
kor_df['total_deaths'].plot()
plt.title('Total deaths in Korea')
plt.grid()
plt.show()

G20_countries = ['Australia', 'Canada', 'Saudi Arabia', 'United States', \
                 'India', 'Russia', 'South Africa', 'Turkey', \
                 'Argentina', 'Brazil', 'Mexico', 'Spain', \
                 'France', 'Germany', 'Italy', 'United Kingdom', \
                 'China', 'Indonesia', 'Japan', 'South Korea']

# your code here

df['case_fatality_rate'] = df['total_deaths']/df['total_cases']

cfr_20211130 = []

plt.figure(figsize=(15,8), dpi=100)
for country in G20_countries[:-1]:
  df[df['location']==country]['case_fatality_rate'].plot(label=country, alpha=0.5)
  cfr_20211130.append(df[df['location']==country]['case_fatality_rate'][-1])
df[df['location']=='South Korea']['case_fatality_rate'].plot(label='South Korea', linewidth=3, color='k')
cfr_20211130.append(df[df['location']=='South Korea']['case_fatality_rate'][-1])
plt.title('Case fatality rate of G20 countries')
plt.legend()
plt.grid()
plt.show()

# your code here 

plt.figure(figsize=(15,6), dpi=100)
plt.bar(G20_countries, cfr_20211130)
plt.xticks(rotation='vertical')
plt.title('Case fatality rate as of Nov. 30th, 2021')
plt.grid()
plt.show()

# your code here

kor_df = kor_df.assign(new_cases_7=np.nan, new_cases_30=np.nan)

for i in range(len(kor_df['new_cases'])):
  if i<7:
    kor_df['new_cases_7'].iloc[i] = kor_df['new_cases'][:i].mean()
  else:
    kor_df['new_cases_7'].iloc[i] = kor_df['new_cases'][i-7:i].mean()
  if i<30:
    kor_df['new_cases_30'].iloc[i] = kor_df['new_cases'][:i].mean()
  else:
    kor_df['new_cases_30'].iloc[i] = kor_df['new_cases'][i-30:i].mean()

plt.figure(figsize=(15,6), dpi=100)
kor_df['new_cases'].plot(alpha=0.5, label='New confirmed cases')
kor_df['new_cases_7'].plot(linewidth=2, label='New confirmed cases (week averaged)')
kor_df['new_cases_30'].plot(linewidth=2, label='New confirmed cases (month averaged)')
plt.legend()
plt.grid()
plt.xticks(rotation='vertical')
plt.show()


df_us = pd.read_csv('https://jonghank.github.io/ase1302/files/Rates_of_COVID-19_Cases_or_Deaths_by_Age_Group_and_Vaccination_Status.csv', index_col=3)
df_us_vac = df_us[ (df_us['outcome']=='death') & (df_us['Age group']=='all_ages_adj') ]

df_us_phz = df_us_vac[df_us_vac['Vaccine product']=='Pfizer']
df_us_mod = df_us_vac[df_us_vac['Vaccine product']=='Moderna']
df_us_jan = df_us_vac[df_us_vac['Vaccine product']=='Janssen']

plt.figure(figsize=(15,6), dpi=100)
phz = df_us_phz['Vaccinated with outcome']/df_us_phz['Fully vaccinated population']*1000000
mod = df_us_mod['Vaccinated with outcome']/df_us_mod['Fully vaccinated population']*1000000
jan = df_us_jan['Vaccinated with outcome']/df_us_jan['Fully vaccinated population']*1000000
unv = df_us_phz['Unvaccinated with outcome']/df_us_phz['Unvaccinated population']*1000000
phz.plot(label='Vaccinated (Pfizer)')
mod.plot(label='Vaccinated (Moderna)')
jan.plot(label='Vaccinated (Janssen)')
unv.plot(label='Unvaccinated')
plt.ylabel('Deaths per million')
plt.title('Vaccine effectiveness (United States)')
plt.legend()
plt.grid()
plt.show()

# your code here

df_us = pd.read_csv('https://jonghank.github.io/ase1302/files/Rates_of_COVID-19_Cases_or_Deaths_by_Age_Group_and_Vaccination_Status.csv', index_col=3)
df_us_vac = df_us[ (df_us['outcome']=='death') & (df_us['Age group']=='18-29') ]

plt.figure(figsize=(15,6), dpi=100)
vac = df_us_vac['Vaccinated with outcome']/df_us_vac['Fully vaccinated population']*1000000
unv = df_us_vac['Unvaccinated with outcome']/df_us_vac['Unvaccinated population']*1000000
vac.plot(label='Vaccinated')
unv.plot(label='Unvaccinated')
plt.ylabel('Deaths per million')
plt.title('Vaccine effectiveness (United States)')
plt.legend()
plt.grid()
plt.show()