
# -*- 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)

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


#### 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()

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']
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']
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']
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\Figure 4.jpg', dpi=100, bbox_inches='tight')
plt.tight_layout()
