# -*- coding: utf-8 -*-
"""
Created on Mon Jan  4 14:43:22 2021

@author: dingq
"""

import os
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.dates import DateFormatter

path = r'...\replication packages' ## replace the directory with the one on your own hard drive.
os.chdir(path)

stringency = pd.read_csv(r'NPI_OxCGRT_latest.csv')
stringency = stringency[['CountryName', 'Date', 'StringencyIndex']]
stringency = stringency.rename(columns = {'Date': 'date', 'CountryName': 'country', 'StringencyIndex':'Stringency Index'})
stringency['date'] = pd.to_datetime(stringency['date'], format = '%Y%m%d')
stringency = stringency.loc[stringency.date<='2020-05-21']

ch = stringency.loc[stringency.country == 'China']
sk = stringency.loc[stringency.country == 'South Korea']
ja = stringency.loc[stringency.country == 'Japan']
sw = stringency.loc[stringency.country == 'Sweden']
tw = stringency.loc[stringency.country == 'Taiwan']
nz = stringency.loc[stringency.country == 'New Zealand']

ch = ch.reset_index(drop = True)
sk = sk.reset_index(drop = True)
ja = ja.reset_index(drop = True)
sw = sw.reset_index(drop = True)
tw = tw.reset_index(drop = True)
nz = nz.reset_index(drop = True)


#20 advanced countries
advanced = ['Spain', 'Switzerland', 'Belgium', 'Poland', 'South Korea', 'France', 'Germany', 'Japan', 'Greece','Canada', 'Ireland', 'Singapore', 'Slovenia', 'United States', 'Netherlands', 'Sweden', 'Taiwan', 'Australia', 'New Zealand', 'United Kingdom']

weight = pd.read_excel(r'market cap.xlsx')
advanced_weight = weight.loc[weight.country.isin(advanced)]

advanced_weight['weight'] = advanced_weight['market_cap']/advanced_weight['market_cap'].sum()
advanced_weight = advanced_weight[['country', 'weight']]

advanced_stringency = pd.merge(stringency, advanced_weight, on ='country' )

advanced_stringency['Stringency Index'] = advanced_stringency['Stringency Index']*advanced_stringency['weight']
ad_strin = advanced_stringency.groupby('date').sum()

ad_strin = ad_strin[['Stringency Index']]
ad_strin = ad_strin.reset_index()
ad_strin = ad_strin.loc[ad_strin.date<='2020-05-21']
ad_strin['line'] = 70

#import matplotlib.dates as mdates
plt.rcParams.update({'font.size': 26})
fig, ax = plt.subplots(figsize=(24,11))

ax.plot(ad_strin['date'], ad_strin['Stringency Index'], c = 'dodgerblue', marker = 'x',linewidth = 3.0)
ax.plot(ad_strin['date'], ad_strin['line'], c = 'black',linestyle='--', linewidth = 2.0)

ax.plot(ch['date'], ch['Stringency Index'], c = 'yellowgreen', linewidth = 3.0)
ax.plot(sk['date'], sk['Stringency Index'], c = 'deepskyblue', linewidth = 3.0, marker = '^')

ax.plot(ja['date'], ja['Stringency Index'], c = 'firebrick', linewidth = 4.0, marker = '|')
ax.plot(sw['date'], sw['Stringency Index'], c = 'blueviolet', marker='o', linewidth = 4.0)
ax.plot(tw['date'], tw['Stringency Index'], c = 'darkorange', linewidth = 3.0)

ax.plot(nz['date'], nz['Stringency Index'], c = 'lime', linestyle='--', linewidth = 3.0)
ax.set_ylim( 0, 120)

plt.grid(color = 'lavender', alpha = 0.9)
plt.xticks(rotation=40)
date_form = DateFormatter("%m/%d/%Y")
ax.xaxis.set_major_formatter(date_form)

ax.legend(['Advanced Economy','China', 'South Korea', 'Japan', 'Sweden', 'Taiwan', 'New Zealand'],bbox_to_anchor=(0.05,-0.4, 1, 1), loc='lower left',
           ncol=4, borderaxespad=0.)

plt.savefig(r'Figure 3.jpg' , dpi=100, bbox_inches='tight')
plt.clf()



