# -*- coding: utf-8 -*-
"""
Created on Sat Jan  2 19:11:18 2021

@author: dingq
"""
import os
import pandas as pd
import matplotlib.pyplot as plt


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

advanced_weighted = pd.read_excel(r'advanced weighted.xlsx')
emerging_weighted = pd.read_excel(r'emerging weighted.xlsx')

advanced_weighted = advanced_weighted[['date', 'w_workplace_percent_change_from_baseline']]
emerging_weighted = emerging_weighted[['date', 'w_workplace_percent_change_from_baseline']]

j_s_s = pd.read_csv(r'g_equity_mob_death_npi.csv')
j_s_s = j_s_s.loc[j_s_s.country.isin(['Japan', 'Sweden', 'South Korea'])]

j_s_s = j_s_s[['country', 'date', 'workplace_percent_change_from_baseline']]

j = j_s_s.loc[j_s_s.country == 'Japan']
sw = j_s_s.loc[j_s_s.country == 'Sweden']
sk = j_s_s.loc[j_s_s.country == 'South Korea']
china = pd.read_excel(r'china mobility percent deviation from baseline.xlsx')
china['date'] = pd.to_datetime(china['date'])
dates = pd.to_datetime(j['date'])
china = china.loc[china.date.isin(dates)]

## plot 
plt.rcParams.update({'font.size': 28})
fig, ax = plt.subplots(figsize=(20,10))

ax.plot(advanced_weighted['date'], advanced_weighted['w_workplace_percent_change_from_baseline'], c = 'royalblue', linestyle='--', linewidth = 3.0, marker = 'x')
ax.plot(emerging_weighted['date'], emerging_weighted['w_workplace_percent_change_from_baseline'], c = 'darkorange', linewidth = 3.0)
plt.grid(color = 'lavender', alpha = 0.5)
plt.xticks(rotation=40)

ax.plot(pd.to_datetime(j['date']), j['workplace_percent_change_from_baseline'], c = 'gold', linewidth = 3.0)
ax.plot(pd.to_datetime(sw['date']), sw['workplace_percent_change_from_baseline'], c = 'skyblue', linewidth = 3.0)
ax.plot(pd.to_datetime(sk['date']), sk['workplace_percent_change_from_baseline'], c = 'mediumseagreen', linewidth = 3.0)
ax.plot(pd.to_datetime(china['date']), china['percent_deviation'], c = 'firebrick', linewidth = 3.0, linestyle = '--')

ax.set_ylim( -70, 20)
#ax.set_xticklabels(pd.to_datetime(advanced_weighted['date']).dt.strftime('%m/%d/%Y').to_list())
plt.grid(color = 'lavender', alpha = 0.9)
plt.xticks(rotation=40)
ax.legend(['Advanced Economy', 'EMDE', 'Japan', 'Sweden', 'South Korea', 'China'],bbox_to_anchor=[0.2,-0.5], loc='lower left',
           ncol=2, borderaxespad=0.)

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



