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

############### import and merge all raw data #########################

######## Use this list to filter every raw dataset #######
######## 34 economies in the sample (exclude China)
countries = ['Spain', 'Switzerland', 'Belgium', 'Poland', 'South Korea', 'France', 
           'Germany', 'Japan', 'Greece','Canada', 'Ireland', 'Singapore', 'Slovenia', 
           'United States', 'Netherlands', 'Sweden', 'Argentina', 'Thailand', 'Croatia',
           'Hungary', 'Brazil', 'Qatar', 'India', 'Romania', 'Turkey', 'Chile', 
           'Kazakhstan', 'South Africa', 'Mexico', 'Malaysia', 'Taiwan', 'New Zealand', 'Australia','United Kingdom']

### 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']
close_base = close_base.groupby('country').mean()
close_base = pd.DataFrame(close_base)
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)

# calculate the accumulated return from Feb. 17 2020
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

# filter the sample dates
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)
equity = equity.sort_values(['country', 'date'])


#### Step 2: import Google mobility data ###
mobility = pd.read_excel(r'workplace mobility.xlsx')
mobility = mobility.rename(columns = {'country_region': 'country'})

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

# filter the sample countries
mobility = mobility.loc[mobility.country.isin(countries)]
mobility = mobility[['country','date', 'workplace_percent_change_from_baseline']]

## merge equity index and mobility
equity_mobility = pd.merge(equity, mobility, how = 'left', on = ['country','date'])

weight = pd.read_excel(r'market cap.xlsx')
weight = weight.loc[weight.country.isin(countries)]
weight['weight'] = weight['market_cap']/weight['market_cap'].sum()
weight = weight[['country', 'weight']]

# order of the countries by the market cap:
# United States, Japan, France, India, Canada, Switzerland
# Germany, Australia, South Korea, Netherlands, Brazil, South Africa
# Taiwan, Spain, Singapore, Sweden, Thailand, Malaysia
# Mexico, Belgium, Chile, Qatar, Poland, Turkey
# Romania, Ireland, New Zealand, Argentina, Greece, Kazakhstan
# Hungary, Croatia, Slovenia

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'])

cases['deaths'] = cases['new death']

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

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

# the unit of raw data is thousand
# calculate death per million
cases['death'] = cases['new death']/(cases['population']/1000) # so that THE population can be million

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

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)

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

equity_mob_death_npi = equity_mob_death_npi.groupby('country').apply(lambda group: group.interpolate(method='linear'))

equity_mob_death_npi.drop(equity_mob_death_npi[(equity_mob_death_npi.country=='United States')&(equity_mob_death_npi.date == '2020-02-17')].index, inplace = True)
equity_mob_death_npi.drop(equity_mob_death_npi[(equity_mob_death_npi.country=='Canada')&(equity_mob_death_npi.date == '2020-02-17')].index, inplace = True)
equity_mob_death_npi.drop(equity_mob_death_npi[(equity_mob_death_npi.country=='South Africa')&(equity_mob_death_npi.date == '2020-02-17')].index, inplace = True)

equity_mob_death_npi.to_csv(r'data_individual plots.csv', index = False)

