
# -*- coding: utf-8 -*-

"""
Created on Thu Jun  4 13:16:37 2020

This script is to make plots for "Stock Prices Lockdowns and Economic Activity in the Time of Coronavirus" 

@author: Dingqian (Olivia) Liu American University Ph.D. Candidate in Economics
"""

# import modules

import os
import numpy as np
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)


###############to make aggregated level##################################################################################

### Step 1: import equity index ####
equity = pd.read_csv(r'equity index.csv')
equity = equity.rename(columns = {'Country': 'country'})
equity = equity.rename(columns = {'Date': 'date'})

## change some countries' name to be consistent across all the datasets
equity['country'] = equity['country'].replace('Korea, Republic Of', 'South Korea')
equity['country'] = equity['country'].replace('Czech Republic', 'Czechia')
equity['country'] = equity['country'].replace('Bosnia And Herzegowina', 'Bosnia and Herzegovina')

equity['date'] = pd.to_datetime(equity['date'])

equity = equity.loc[equity.country.isin(countries)]

equity = equity[['country','date', 'Real_Close' ]]

close_base = equity.loc[equity.date =='2020-02-17'] # remember to delete U.S. Canada and South Africa
close_base = close_base[['country', 'Real_Close']]

close_base = close_base.reset_index()
close_base = close_base.rename(columns = {'Real_Close': 'close_base'})
equity = pd.merge(equity,close_base , on = 'country')

equity['Real_Close'] = equity['Real_Close'].astype(float)
equity['log_close_deviate'] = np.log(equity['Real_Close']/equity['close_base'])
equity['stock market percentage change from Feb.17'] = equity['log_close_deviate']*100

start_date = '2020-02-17'
end_date = '2020-05-21'
equity = equity.loc[(equity['date'] >= start_date) &(equity['date']<=end_date)]
equity = equity.reset_index(drop = True)

mobility = pd.read_excel(r'workplace mobility.xlsx')
mobility = mobility.rename(columns = {'country_region': 'country'})

mobility['date'] = pd.to_datetime(mobility['date'])

len(equity)
len(mobility)

mobility = mobility.loc[mobility.country.isin(countries)]

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


## for country-level time tracking stock/mobility plot
equity_mobility = pd.merge(equity, mobility, how = 'left', on = ['country','date'])
equity_mobility = equity_mobility.interpolate(method='linear')
weight = pd.read_excel(r'market cap.xlsx')

weight['weight'] = weight['market_cap']/weight['market_cap'].sum()

weight = weight[['country', 'weight']]

equity_mobility = pd.merge(equity_mobility, weight, on = 'country')


## merge with death cases

cases = pd.read_csv(r'deaths.csv')
cases['date'] = pd.to_datetime(cases['date'])
cases = cases.sort_values(['country','date'])

population = pd.read_excel(r'population.xlsx')

cases = pd.merge(cases, population, on = 'country', how = 'outer')

cases['death'] = cases['new death']/(cases['population']/1000) # so that population can be million
cases = cases[['country', 'date', 'death']]

equity_mob_death = pd.merge(equity_mobility, cases, how = 'left', on = ['country', 'date'])
equity_mob_death = equity_mob_death.interpolate(method='linear')

npi = pd.read_csv(r'NPI_OxCGRT_latest.csv')

npi = npi[['CountryName', 'Date', 'GovernmentResponseIndex', 'EconomicSupportIndex', 'StringencyIndex']]
npi = npi.rename(columns = {'Date': 'date', 'CountryName': 'country'})
npi['date'] = pd.to_datetime(npi['date'], format = '%Y%m%d' )
npi = npi.loc[npi.country.isin(countries)]

equity_mob_death_npi = pd.merge(equity_mob_death, npi, how ='left', on = ['country', 'date'])
equity_mob_death_npi = equity_mob_death_npi.reset_index(drop = True)

union_dates = pd.DataFrame({'date':equity.date.unique()})

union_dates = union_dates.sort_values('date')
union_dates = union_dates.reset_index(drop = True)

########################################################
# This part should be changed when adding countries, replace 34 with the number of countries in the  ##
union_dates = pd.concat([union_dates]* 34)
union_dates = union_dates.reset_index()
countries_l = countries*69
countries_l.sort()
#######################################################

countries_l = pd.Series(countries_l, name = 'country')

len(union_dates) == len(countries_l)
union_dates = union_dates.merge(countries_l, left_index = True, right_index = True)
union_dates = union_dates.drop(columns = 'index')

union_dates['date'] = pd.to_datetime(union_dates['date'])

g_equity_mob_death_npi = pd.merge(equity_mob_death_npi, union_dates, on = ['date', 'country'], how = 'right')
g_equity_mob_death_npi = g_equity_mob_death_npi.sort_values(['country', 'date'])

g_equity_mob_death_npi.to_csv(r'test.csv', index = False) #2070 observations with many days gaps.

g_equity_mob_death_npi = g_equity_mob_death_npi.interpolate(method='linear')

g_equity_mob_death_npi = g_equity_mob_death_npi.sort_values(['country', 'date'])
g_equity_mob_death_npi = g_equity_mob_death_npi.reset_index(drop = True)

g_equity_mob_death_npi = g_equity_mob_death_npi.rename(columns = {'GovernmentResponseIndex': 'Government Response Index'})
g_equity_mob_death_npi = g_equity_mob_death_npi.rename(columns = {'EconomicSupportIndex': 'Economic Support Index'})
g_equity_mob_death_npi = g_equity_mob_death_npi.rename(columns = {'StringencyIndex': 'Stringency Index'})

g_equity_mob_death_npi.to_csv(r'g_equity_mob_death_npi.csv')



# to make weighted index

weight = pd.read_excel(r'market cap.xlsx')

## Advanced economies20: Taiwan, Spain, Switzerland, Belgium, Poland,South Korea, France, Germany, Japan, Greece, Canada, Ireland, Singapore, Slovenia, United States, Netherland, Sweden, Australia, New Zealand
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']
advanced_weight = weight.loc[weight.country.isin(advanced)]

## Emerging Economies14: Argentina, Thailand, Croatia, Hungary, Brazil, Qatar, India, Romania, Turkey, Chile, Kazakhstan, South Africa, Mexico, Malaysia
emerging = ['Argentina', 'Thailand', 'Croatia', 'Hungary', 'Brazil', 'Qatar', 'India', 'Romania', 'Turkey', 'Chile', 'Kazakhstan', 'South Africa', 'Mexico', 'Malaysia']
emerging_weight = weight.loc[weight.country.isin(emerging)]



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

weight = weight[['country', 'weight']]
advanced_weight = advanced_weight[['country', 'weight']]
emerging_weight = emerging_weight[['country', 'weight']]


###### interpolate dates ####
g_equity_mob_death_npi = g_equity_mob_death_npi[['country', 'date', 'workplace_percent_change_from_baseline',
       'stock market percentage change from Feb.17', 'death',
       'Government Response Index', 'Economic Support Index',
       'Stringency Index']]
## global #34
w_equity_mob_cases_npi = pd.merge(g_equity_mob_death_npi, weight, on = 'country', how = 'outer')

w_equity_mob_cases_npi = w_equity_mob_cases_npi.drop_duplicates()
#w_equity_mob_cases_npi.to_csv('w_equity_mob_cases_npi_wTaiwan.csv')

# advanced 20
a_equity_mob_cases_npi = pd.merge(g_equity_mob_death_npi, advanced_weight, on = 'country', how = 'outer')

a_countries = list(advanced_weight.country.unique())
a_equity_mob_cases_npi = a_equity_mob_cases_npi.loc[a_equity_mob_cases_npi.country.isin(a_countries)]
a_equity_mob_cases_npi = a_equity_mob_cases_npi.drop_duplicates()

# emerging 14
e_equity_mob_cases_npi = pd.merge(g_equity_mob_death_npi, emerging_weight, on = 'country', how = 'outer')

e_countries = list(emerging_weight.country.unique())
e_equity_mob_cases_npi = e_equity_mob_cases_npi.loc[e_equity_mob_cases_npi.country.isin(e_countries)]
e_equity_mob_cases_npi = e_equity_mob_cases_npi.drop_duplicates()




#w_equity_mob_cases_npi['w_close'] = w_equity_mob_cases_npi['Real_Close']*w_equity_mob_cases_npi['weight'].astype(float)
#w_equity_mob_cases_npi['w_close_base'] = w_equity_mob_cases_npi['close_base']*w_equity_mob_cases_npi['weight']
#w_equity_mob_cases_npi['log_close_deviate'] = w_equity_mob_cases_npi['w_close']/w_equity_mob_cases_npi['w_close_base'])

# calculated market-cap weighted equity deviation
w_equity_mob_cases_npi['stock market percentage change from Feb.17'] = w_equity_mob_cases_npi['stock market percentage change from Feb.17']*w_equity_mob_cases_npi['weight']
a_equity_mob_cases_npi['stock market percentage change from Feb.17'] = a_equity_mob_cases_npi['stock market percentage change from Feb.17']*a_equity_mob_cases_npi['weight']
e_equity_mob_cases_npi['stock market percentage change from Feb.17'] = e_equity_mob_cases_npi['stock market percentage change from Feb.17']*e_equity_mob_cases_npi['weight']

# calculated market-cap weighted mobility deviation
w_equity_mob_cases_npi['w_workplace_percent_change_from_baseline'] = w_equity_mob_cases_npi[ 'workplace_percent_change_from_baseline']*w_equity_mob_cases_npi['weight']
a_equity_mob_cases_npi['w_workplace_percent_change_from_baseline'] = a_equity_mob_cases_npi[ 'workplace_percent_change_from_baseline']*a_equity_mob_cases_npi['weight']
e_equity_mob_cases_npi['w_workplace_percent_change_from_baseline'] = e_equity_mob_cases_npi[ 'workplace_percent_change_from_baseline']*e_equity_mob_cases_npi['weight']

# calculate death and governemnt response index
w_equity_mob_cases_npi['death'] = w_equity_mob_cases_npi['death']*w_equity_mob_cases_npi['weight']
w_equity_mob_cases_npi['Government Response Index'] = w_equity_mob_cases_npi['Government Response Index']*w_equity_mob_cases_npi['weight']
w_equity_mob_cases_npi['Economic Support Index'] = w_equity_mob_cases_npi['Economic Support Index']*w_equity_mob_cases_npi['weight']
w_equity_mob_cases_npi['Stringency Index'] = w_equity_mob_cases_npi['Stringency Index']*w_equity_mob_cases_npi['weight']


a_equity_mob_cases_npi['death'] = a_equity_mob_cases_npi['death']*a_equity_mob_cases_npi['weight']
a_equity_mob_cases_npi['Government Response Index'] = a_equity_mob_cases_npi['Government Response Index']*a_equity_mob_cases_npi['weight']
a_equity_mob_cases_npi['Economic Support Index'] = a_equity_mob_cases_npi['Economic Support Index']*a_equity_mob_cases_npi['weight']
a_equity_mob_cases_npi['Stringency Index'] = a_equity_mob_cases_npi['Stringency Index']*a_equity_mob_cases_npi['weight']

e_equity_mob_cases_npi['death'] = e_equity_mob_cases_npi['death']*e_equity_mob_cases_npi['weight']
e_equity_mob_cases_npi['Government Response Index'] = e_equity_mob_cases_npi['Government Response Index']*e_equity_mob_cases_npi['weight']
e_equity_mob_cases_npi['Economic Support Index'] = e_equity_mob_cases_npi['Economic Support Index']*e_equity_mob_cases_npi['weight']
e_equity_mob_cases_npi['Stringency Index'] = e_equity_mob_cases_npi['Stringency Index']*e_equity_mob_cases_npi['weight']



##################
global_weighted = w_equity_mob_cases_npi[['stock market percentage change from Feb.17', 'w_workplace_percent_change_from_baseline', 'death','Government Response Index', 'date','Economic Support Index' ,'Stringency Index' ]].groupby('date').sum()
advanced_weighted = a_equity_mob_cases_npi[['stock market percentage change from Feb.17', 'w_workplace_percent_change_from_baseline', 'death','Government Response Index', 'date', 'Economic Support Index', 'Stringency Index'  ]].groupby('date').sum()
emerging_weighted = e_equity_mob_cases_npi[['stock market percentage change from Feb.17', 'w_workplace_percent_change_from_baseline', 'death','Government Response Index', 'date', 'Economic Support Index', 'Stringency Index'  ]].groupby('date').sum()

##global_weighted['date'] = global_weighted.index
#advanced_weighted['date'] = advanced_weighted.index
#emerging_weighted['date'] = emerging_weighted.index


#global_weighted = global_weighted.loc[global_weighted['stock market percentage change from Feb.17']!=0]
global_weighted.to_excel(r'global weighted.xlsx')

#advanced_weighted = advanced_weighted.loc[advanced_weighted['stock market percentage change from Feb.17']!=0]
advanced_weighted.to_excel(r'advanced weighted.xlsx')

#emerging_weighted = emerging_weighted.loc[emerging_weighted['stock market percentage change from Feb.17']!=0]
emerging_weighted.to_excel(r'emerging weighted.xlsx')

advanced_weighted = advanced_weighted.reset_index()
emerging_weighted = emerging_weighted.reset_index()



### scatterplot global equity work --- Advanced and emerging, figure 6.a
plt.rcParams.update({'font.size': 17.5})

fig, ax = plt.subplots(nrows=1, ncols=2,  sharey = True,  sharex = True,
                        figsize=(14,8))        
  #  ax0 = fig.add_subplot(111)
#plt.xlabel("Percent Workplace … Percent Change …")
fig.text(0.5, 0,"Percent Workplace Mobility Deviation from Baseline", ha = 'center', fontsize=22)
fig.text(0.0, 0.1,"Percent Stock Market Deviation from Feb.17", ha = 'center', rotation = 'vertical', fontsize= 23)
plt.tight_layout()

plt.axis([-65, 15, -50, 10])

ax[0].set_title('Advanced Economy')
ax[1].set_title('EMDE')
for ax0 in ax.flat:
    ax0.label_outer()

advanced_weighted = advanced_weighted.fillna(0)
advanced_weighted['week'] = advanced_weighted.date.dt.week

advanced_weighted_week = advanced_weighted.groupby(['week']).mean()

y0= advanced_weighted_week['stock market percentage change from Feb.17']
x0 = advanced_weighted_week["w_workplace_percent_change_from_baseline"]
xa0 = sm.add_constant(x0)
ols0 = sm.OLS(y0,xa0 )
res0 = ols0.fit()
slope0 = res0.params[1]
intercept0 = res0.params[0]
rsquared_adj0 = res0.rsquared_adj
s_err_i0 = res0.bse[0]
s_err_s0 = res0.bse[1]
ax[0].plot(x0, y0, 'o')
ax[0].plot(x0, intercept0+slope0*x0, '-', c = 'black')

tex01 = 'y = ' +str(round(intercept0, 2)) + ' + ' + str(round(slope0, 2)) + '*x' 
tex02 = '      (' + str(round(s_err_i0, 2)) + ' )      (' + str(round(s_err_s0, 2)) +')'
tex03 = "adj. R-square = " + str(round(rsquared_adj0, 2))
ax[0].text(-58,5, tex01, horizontalalignment='left', size='large', color='black', weight='semibold')
ax[0].text(-58, 2, tex02, horizontalalignment='left', size='large', color='black', weight='semibold')
ax[0].text(-58, -4, tex03, horizontalalignment='left', size='large', color='black', weight='semibold')


emerging_weighted = emerging_weighted.fillna(0)
emerging_weighted['week'] = emerging_weighted.date.dt.week

emerging_weighted_week = emerging_weighted.groupby(['week']).mean()

y1= emerging_weighted_week['stock market percentage change from Feb.17']
x1 = emerging_weighted_week["w_workplace_percent_change_from_baseline"]
xa1 = sm.add_constant(x1)
ols1 = sm.OLS(y1,xa1 )
res1 = ols1.fit()
slope1 = res1.params[1]
intercept1 = res1.params[0]
rsquared_adj1 = res1.rsquared_adj
s_err_i1 = res1.bse[0]
s_err_s1 = res1.bse[1]
ax[1].plot(x1, y1, 'o')
ax[1].plot(x1, intercept1 + slope1*x1, '-', c = 'black')

tex11 = 'y = ' +str(round(intercept1, 2)) + ' + ' + str(round(slope1, 2)) + '*x' 
tex12 = '      (' + str(round(s_err_i1, 2)) + ' )      (' + str(round(s_err_s1, 2)) +')'
tex13 = "adj. R-square = " + str(round(rsquared_adj1, 2))
ax[1].text(-58,5, tex11, horizontalalignment='left', size='large', color='black', weight='semibold')
ax[1].text(-58, 2, tex12, horizontalalignment='left', size='large', color='black', weight='semibold')
ax[1].text(-58, -4, tex13, horizontalalignment='left', size='large', color='black', weight='semibold')


### scatterplot aggregated equity death --- Advanced and emerging figure 6.d
plt.rcParams.update({'font.size': 17.5})

fig, ax = plt.subplots(nrows=1, ncols=2,  sharey = True,  sharex = True,
                        figsize=(14,8))        
  #  ax0 = fig.add_subplot(111)
#plt.xlabel("Percent Workplace … Percent Change …")
fig.text(0.5, 0,"Deaths per Million", ha = 'center', fontsize=22)
fig.text(0.0, 0.1,"Percent Stock Market Deviation from Feb.17", ha = 'center', rotation = 'vertical', fontsize= 23)
plt.tight_layout()

plt.axis([-1, 6, -50, 10])

ax[0].set_title('Advanced Economy')
ax[1].set_title('EMDE')
for ax0 in ax.flat:
    ax0.label_outer()


advanced_weighted = advanced_weighted.fillna(0)
advanced_weighted['week'] = advanced_weighted.date.dt.week

advanced_weighted_week = advanced_weighted.groupby(['week']).mean()

y0= advanced_weighted_week['stock market percentage change from Feb.17']
x0 = advanced_weighted_week["death"]
xa0 = sm.add_constant(x0)
ols0 = sm.OLS(y0,xa0 )
res0 = ols0.fit()
slope0 = res0.params[1]
intercept0 = res0.params[0]
rsquared_adj0 = res0.rsquared_adj
s_err_i0 = res0.bse[0]
s_err_s0 = res0.bse[1]
ax[0].plot(x0, y0, 'o')
ax[0].plot(x0, intercept0+slope0*x0, '-', c = 'black')

tex01 = 'y = ' +str(round(intercept0, 2)) + ' + ' + str(round(slope0, 2)) + '*x' 
tex02 = '      (' + str(round(s_err_i0, 2)) + ' )      (' + str(round(s_err_s0, 2)) +')'
tex03 = "adj. R-square = " + str(round(rsquared_adj0, 2))
ax[0].text(0.5,5, tex01, horizontalalignment='left', size='large', color='black', weight='semibold')
ax[0].text(0.5, 2, tex02, horizontalalignment='left', size='large', color='black', weight='semibold')
ax[0].text(0.5, -4, tex03, horizontalalignment='left', size='large', color='black', weight='semibold')


emerging_weighted = emerging_weighted.fillna(0)
emerging_weighted['week'] = emerging_weighted.date.dt.week

emerging_weighted_week = emerging_weighted.groupby(['week']).mean()

y1= emerging_weighted_week['stock market percentage change from Feb.17']
x1 = emerging_weighted_week["death"]
xa1 = sm.add_constant(x1)
ols1 = sm.OLS(y1,xa1 )
res1 = ols1.fit()
slope1 = res1.params[1]
intercept1 = res1.params[0]
rsquared_adj1 = res1.rsquared_adj
s_err_i1 = res1.bse[0]
s_err_s1 = res1.bse[1]
ax[1].plot(x1, y1, 'o')
ax[1].plot(x1, intercept1 + slope1*x1, '-', c = 'black')

tex11 = 'y = ' +str(round(intercept1, 2)) + ' + ' + str(round(slope1, 2)) + '*x' 
tex12 = '      (' + str(round(s_err_i1, 2)) + ' )      (' + str(round(s_err_s1, 2)) +')'
tex13 = "adj. R-square = " + str(round(rsquared_adj1, 2))
ax[1].text(0.5,5, tex11, horizontalalignment='left', size='large', color='black', weight='semibold')
ax[1].text(0.5, 2, tex12, horizontalalignment='left', size='large', color='black', weight='semibold')
ax[1].text(0.5, -4, tex13, horizontalalignment='left', size='large', color='black', weight='semibold')





### scatterplot global equity economic support --- Advanced and emerging figure 6.c
plt.rcParams.update({'font.size': 17.5})

fig, ax = plt.subplots(nrows=1, ncols=2,  sharey = True,  sharex = True,
                        figsize=(14,8))        
  #  ax0 = fig.add_subplot(111)
#plt.xlabel("Percent Workplace … Percent Change …")
fig.text(0.5, 0,"Economic Support Index", ha = 'center', fontsize=22)
fig.text(0.0, 0.1,"Percent Stock Market Deviation from Feb.17", ha = 'center', rotation = 'vertical', fontsize= 23)
plt.tight_layout()

plt.axis([0, 100, -50, 10])

ax[0].set_title('Advanced Economy')
ax[1].set_title('EMDE')
for ax0 in ax.flat:
    ax0.label_outer()


advanced_weighted = advanced_weighted.fillna(0)
advanced_weighted['week'] = advanced_weighted.date.dt.week

advanced_weighted_week = advanced_weighted.groupby(['week']).mean()

y0= advanced_weighted_week['stock market percentage change from Feb.17']
x0 = advanced_weighted_week["Economic Support Index"]
xa0 = sm.add_constant(x0)
ols0 = sm.OLS(y0,xa0 )
res0 = ols0.fit()
slope0 = res0.params[1]
intercept0 = res0.params[0]
rsquared_adj0 = res0.rsquared_adj
s_err_i0 = res0.bse[0]
s_err_s0 = res0.bse[1]
ax[0].plot(x0, y0, 'o')
ax[0].plot(x0, intercept0+slope0*x0, '-', c = 'black')

tex01 = 'y = ' +str(round(intercept0, 2)) + ' + ' + str(round(slope0, 2)) + '*x' 
tex02 = '      (' + str(round(s_err_i0, 2)) + ' )      (' + str(round(s_err_s0, 2)) +')'
tex03 = "adj. R-square = " + str(round(rsquared_adj0, 2))
ax[0].text(10,5, tex01, horizontalalignment='left', size='large', color='black', weight='semibold')
ax[0].text(10, 2, tex02, horizontalalignment='left', size='large', color='black', weight='semibold')
ax[0].text(10, -4, tex03, horizontalalignment='left', size='large', color='black', weight='semibold')


emerging_weighted = emerging_weighted.fillna(0)
emerging_weighted['week'] = emerging_weighted.date.dt.week

emerging_weighted_week = emerging_weighted.groupby(['week']).mean()

y1= emerging_weighted_week['stock market percentage change from Feb.17']
x1 = emerging_weighted_week["Economic Support Index"]
xa1 = sm.add_constant(x1)
ols1 = sm.OLS(y1,xa1 )
res1 = ols1.fit()
slope1 = res1.params[1]
intercept1 = res1.params[0]
rsquared_adj1 = res1.rsquared_adj
s_err_i1 = res1.bse[0]
s_err_s1 = res1.bse[1]
ax[1].plot(x1, y1, 'o')
ax[1].plot(x1, intercept1 + slope1*x1, '-', c = 'black')

tex11 = 'y = ' +str(round(intercept1, 2)) + ' + ' + str(round(slope1, 2)) + '*x' 
tex12 = '      (' + str(round(s_err_i1, 2)) + ' )      (' + str(round(s_err_s1, 2)) +')'
tex13 = "adj. R-square = " + str(round(rsquared_adj1, 2))
ax[1].text(10,5, tex11, horizontalalignment='left', size='large', color='black', weight='semibold')
ax[1].text(10, 2, tex12, horizontalalignment='left', size='large', color='black', weight='semibold')
ax[1].text(10, -4, tex13, horizontalalignment='left', size='large', color='black', weight='semibold')


### scatterplot global equity stringency --- Advanced and emerging   figure 6.b
plt.rcParams.update({'font.size': 17.5})

fig, ax = plt.subplots(nrows=1, ncols=2,  sharey = True,  sharex = True,
                        figsize=(14,8))        
  #  ax0 = fig.add_subplot(111)
#plt.xlabel("Percent Workplace … Percent Change …")
fig.text(0.5, 0,"Stringency Index", ha = 'center', fontsize=22)
fig.text(0.0, 0.1,"Percent Stock Market Deviation from Feb.17", ha = 'center', rotation = 'vertical', fontsize= 23)
plt.tight_layout()

plt.axis([0, 100, -50, 10])

ax[0].set_title('Advanced Economy')
ax[1].set_title('EMDE')
for ax0 in ax.flat:
    ax0.label_outer()


advanced_weighted = advanced_weighted.fillna(0)
advanced_weighted['week'] = advanced_weighted.date.dt.week

advanced_weighted_week = advanced_weighted.groupby(['week']).mean()

y0= advanced_weighted_week['stock market percentage change from Feb.17']
x0 = advanced_weighted_week["Economic Support Index"]
xa0 = sm.add_constant(x0)
ols0 = sm.OLS(y0,xa0 )
res0 = ols0.fit()
slope0 = res0.params[1]
intercept0 = res0.params[0]
rsquared_adj0 = res0.rsquared_adj
s_err_i0 = res0.bse[0]
s_err_s0 = res0.bse[1]
ax[0].plot(x0, y0, 'o')
ax[0].plot(x0, intercept0+slope0*x0, '-', c = 'black')

tex01 = 'y = ' +str(round(intercept0, 2)) + ' + ' + str(round(slope0, 2)) + '*x' 
tex02 = '      (' + str(round(s_err_i0, 2)) + ' )      (' + str(round(s_err_s0, 2)) +')'
tex03 = "adj. R-square = " + str(round(rsquared_adj0, 2))
ax[0].text(10,5, tex01, horizontalalignment='left', size='large', color='black', weight='semibold')
ax[0].text(10, 2, tex02, horizontalalignment='left', size='large', color='black', weight='semibold')
ax[0].text(10, -4, tex03, horizontalalignment='left', size='large', color='black', weight='semibold')


emerging_weighted = emerging_weighted.fillna(0)
emerging_weighted['week'] = emerging_weighted.date.dt.week

emerging_weighted_week = emerging_weighted.groupby(['week']).mean()

y1= emerging_weighted_week['stock market percentage change from Feb.17']
x1 = emerging_weighted_week["Stringency Index"]
xa1 = sm.add_constant(x1)
ols1 = sm.OLS(y1,xa1 )
res1 = ols1.fit()
slope1 = res1.params[1]
intercept1 = res1.params[0]
rsquared_adj1 = res1.rsquared_adj
s_err_i1 = res1.bse[0]
s_err_s1 = res1.bse[1]
ax[1].plot(x1, y1, 'o')
ax[1].plot(x1, intercept1 + slope1*x1, '-', c = 'black')

tex11 = 'y = ' +str(round(intercept1, 2)) + ' + ' + str(round(slope1, 2)) + '*x' 
tex12 = '      (' + str(round(s_err_i1, 2)) + ' )      (' + str(round(s_err_s1, 2)) +')'
tex13 = "adj. R-square = " + str(round(rsquared_adj1, 2))
ax[1].text(10,5, tex11, horizontalalignment='left', size='large', color='black', weight='semibold')
ax[1].text(10, 2, tex12, horizontalalignment='left', size='large', color='black', weight='semibold')
ax[1].text(10, -4, tex13, horizontalalignment='left', size='large', color='black', weight='semibold')



#### Figure 4. aggregated time trace for global, emerging and advanced ####
plt.figure(figsize=(12,5.8))
plt.rcParams.update({'font.size': 17.5})

fig, ax = plt.subplots(nrows=3, ncols=1, sharex = True, sharey = True,  
                        figsize=(14,18))        

plt.xlim( -60, 10)
plt.ylim( -50, 10)

fig.text(0.5, 0,"Percent Workplace Mobility Deviation from Baseline", ha = 'center', fontsize=22)
fig.text(0.0, 0.3,"Percent Stock Market Deviation from 17 February 2020", ha = 'center', rotation = 'vertical', fontsize= 22)
plt.tight_layout()

#plt.axis([-60, 15, -50, 9])

ax[0].set_title('Global Economy')
ax[1].set_title('Advanced Economy')
ax[2].set_title('Emerging Market and Developing Economys (EMDEs)')
for ax0 in ax.flat:
    ax0.label_outer()

data0 = global_weighted.reset_index()
x0 = data0["w_workplace_percent_change_from_baseline"]
y0 = data0['stock market percentage change from Feb.17']
#left0 = data.index[data["w_workplace_percent_change_from_baseline"] == data["w_workplace_percent_change_from_baseline"].min()]
xl0 = data0["w_workplace_percent_change_from_baseline"].min()
yl0 = data0.loc[data0["w_workplace_percent_change_from_baseline"] == xl0]['stock market percentage change from Feb.17']
zl0 = data0.loc[data0["w_workplace_percent_change_from_baseline"] == xl0]['date']
zl0 = zl0.dt.strftime('%m-%d').values[0]
xs0 = data0.loc[data0["Stringency Index"]>=70]["w_workplace_percent_change_from_baseline"].iloc[0]
ys0 = data0.loc[data0["Stringency Index"]>=70]['stock market percentage change from Feb.17'].iloc[0]
zs0 = data0.loc[data0["Stringency Index"]>=70]['date'].dt.strftime('%m-%d').iloc[0]

data1 = advanced_weighted.reset_index()
x1 = data1["w_workplace_percent_change_from_baseline"]
y1 = data1['stock market percentage change from Feb.17']
#left0 = data.index[data["w_workplace_percent_change_from_baseline"] == data["w_workplace_percent_change_from_baseline"].min()]
xl1 = data1["w_workplace_percent_change_from_baseline"].min()
yl1 = data1.loc[data1["w_workplace_percent_change_from_baseline"] == xl1]['stock market percentage change from Feb.17']
zl1 = data1.loc[data1["w_workplace_percent_change_from_baseline"] == xl1]['date']
zl1 = zl1.dt.strftime('%m-%d').values[0]
xs1 = data1.loc[data1["Stringency Index"]>=70]["w_workplace_percent_change_from_baseline"].iloc[0]
ys1 = data1.loc[data1["Stringency Index"]>=70]['stock market percentage change from Feb.17'].iloc[0]
zs1 = data1.loc[data1["Stringency Index"]>=70]['date'].dt.strftime('%m-%d').iloc[0]


data2 = emerging_weighted.reset_index()
x2 = data2["w_workplace_percent_change_from_baseline"]
y2 = data2['stock market percentage change from Feb.17']
#left0 = data.index[data["w_workplace_percent_change_from_baseline"] == data["w_workplace_percent_change_from_baseline"].min()]
xl2 = data2["w_workplace_percent_change_from_baseline"].min()
yl2 = data2.loc[data2["w_workplace_percent_change_from_baseline"] == xl2]['stock market percentage change from Feb.17']
zl2 = data2.loc[data2["w_workplace_percent_change_from_baseline"] == xl2]['date']
zl2 = zl2.dt.strftime('%m-%d').values[0]
xs2 = data2.loc[data2["Stringency Index"]>=70]["w_workplace_percent_change_from_baseline"].iloc[0]
ys2 = data2.loc[data2["Stringency Index"]>=70]['stock market percentage change from Feb.17'].iloc[0]
zs2 = data2.loc[data2["Stringency Index"]>=70]['date'].dt.strftime('%m-%d').iloc[0]

ax[0].scatter(xl0,yl0, marker = 'D', s =400, c = 'orange')
ax[0].scatter(xs0,ys0, marker = 'X', s =400, c = 'Lime')
ax[1].scatter(xl1,yl1, marker = 'D', s =400, c = 'orange')
ax[1].scatter(xs1,ys1, marker = 'X', s =400, c = 'Lime')
ax[2].scatter(xl2,yl2, marker = 'D', s =400, c = 'orange')
ax[2].scatter(xs2,ys2, marker = 'X', s =400, c = 'Lime')

ax[0].plot(x0, y0)
ax[0].scatter(x0, y0)

ax[1].plot(x1, y1)
ax[1].scatter(x1, y1)

ax[2].plot(x2, y2)
ax[2].scatter(x2, y2)


dates0 = data0.date
dates0 = dates0.dt.strftime('%m-%d')

ax[0].annotate(dates0.iloc[0], (x0.iloc[0], y0.iloc[0]), textcoords="offset points",
         xytext=(0,10),ha='center')
ax[0].annotate(dates0.iloc[-1], (x0.iloc[-1], y0.iloc[-1]), textcoords="offset points",
         xytext=(0,10),ha='center')
ax[0].annotate(dates0.iloc[18], (x0.iloc[18], y0.iloc[18]), textcoords="offset points",
         xytext=(0,10),ha='center')
ax[0].annotate(zl0, (xl0, yl0), textcoords="offset points",
         xytext=(0,15),ha='center')
ax[0].annotate(zs0, (xs0, ys0), textcoords="offset points",
         xytext=(-40,-45),ha='center' ,arrowprops=dict(arrowstyle="->",
                            connectionstyle="arc3"))
ax[0].annotate(dates0.iloc[25], (data0.iloc[25]['w_workplace_percent_change_from_baseline'], data0.iloc[25]['stock market percentage change from Feb.17']), textcoords="offset points",
         xytext=(20,35),ha='center' ,arrowprops=dict(arrowstyle="->",
                            connectionstyle="arc3"))

dates1 = data1.date
dates1 = dates1.dt.strftime('%m-%d')

ax[1].annotate(dates1.iloc[0], (x1.iloc[0], y1.iloc[0]), textcoords="offset points",
         xytext=(0,10),ha='center')
ax[1].annotate(dates1.iloc[-1], (x1.iloc[-1], y1.iloc[-1]), textcoords="offset points",
         xytext=(0,10),ha='center')
ax[1].annotate(dates1.iloc[18], (x1.iloc[18], y1.iloc[18]), textcoords="offset points",
         xytext=(0,10),ha='center')
ax[1].annotate(zl1, (xl1, yl1), textcoords="offset points",
         xytext=(0,15),ha='center')
ax[1].annotate(zs1, (xs1, ys1), textcoords="offset points",
         xytext=(-40,-45),ha='center' ,arrowprops=dict(arrowstyle="->",
                            connectionstyle="arc3"))
ax[1].annotate(dates1.iloc[25], (data1.iloc[25]['w_workplace_percent_change_from_baseline'], data1.iloc[25]['stock market percentage change from Feb.17']), textcoords="offset points",
         xytext=(20,35),ha='center' ,arrowprops=dict(arrowstyle="->",
                            connectionstyle="arc3"))

dates2 = data2.date
dates2 = dates2.dt.strftime('%m-%d')

ax[2].annotate(dates2.iloc[0], (x2.iloc[0], y2.iloc[0]), textcoords="offset points",
         xytext=(0,10),ha='center')
ax[2].annotate(dates2.iloc[-1], (x2.iloc[-1], y2.iloc[-1]), textcoords="offset points",
         xytext=(0,10),ha='center')
ax[2].annotate(dates2.iloc[18], (x2.iloc[18], y2.iloc[18]), textcoords="offset points",
         xytext=(0,10),ha='center')
ax[2].annotate(zl2, (xl2, yl2), textcoords="offset points",
         xytext=(0,15),ha='center')
ax[2].annotate(zs2, (xs2, ys2), textcoords="offset points",
         xytext=(-40,-45),ha='center' ,arrowprops=dict(arrowstyle="->",
                            connectionstyle="arc3"))
ax[2].annotate(dates2.iloc[25], (data2.iloc[25]['w_workplace_percent_change_from_baseline'], data2.iloc[25]['stock market percentage change from Feb.17']), textcoords="offset points",
         xytext=(20,35),ha='center' ,arrowprops=dict(arrowstyle="->",
                            connectionstyle="arc3"))

plt.savefig(r'plots\aggregated\Figure 4.jpg', dpi=100, bbox_inches='tight')
plt.tight_layout()


#### time series of advanced EMDE, 4 independent variables

## advanced workplace mobility
plt.rcParams.update({'font.size': 15})
plt.figure(figsize=(10,8))
p = sns.relplot(x = "date", y='w_workplace_percent_change_from_baseline', data = advanced_weighted, kind = 'line')
plt.title("Advanced Economy")
plt.ylim( -60, 10)
plt.ylabel("Workplace Mobility, Percent Change from Baseline")

p.set_xticklabels( rotation=45)
p.savefig(r'timeseries_advanced_work.jpg')
plt.clf()


## EMDE workplace mobility
plt.rcParams.update({'font.size': 15})
plt.figure(figsize=(10,8))
p = sns.relplot(x = "date", y='w_workplace_percent_change_from_baseline', data = emerging_weighted, kind = 'line')
plt.title("EMDE")
plt.ylim( -60, 10)
plt.ylabel("Workplace Mobility, Percent Change from Baseline")
p.set_xticklabels( rotation=45)
p.savefig(r'timeseries_emerging_work.jpg')
plt.clf()

