
import pandas as pd
import numpy as np
import os

path = r'...\replication package'
path = r'C:\Users\dingq\Documents\200_academic\220_Academic_My dissertation\220_papaer_preperation\covid\replication package'
os.chdir(path)

# to generate market_cap weighted average

####### import and merge all the data ###
### here are all the sample countries we use in this project
### Use this list to filter every dataset
## 34 countries in samples
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', 'Australia','New Zealand', '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'] # 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'])

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

cases['deaths'] = cases['new death']
cases['deaths'] = cases['deaths'].astype('int')

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

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

cases['death'] = cases['deaths']/(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)

dt = pd.read_excel(r'doubling time till Aug 31.xlsx')
dt = dt[['country', 'date', 'log_dt7']]
dt['date'] = pd.to_datetime(dt['date'])

dt = dt.loc[(dt.date>='2020-02-17') & (dt.date<='2020-05-21')]
dt = dt.loc[dt.country.isin(countries)]
dt = dt.reset_index(drop = True)
equity_mob_death_npi_dt = pd.merge(equity_mob_death_npi, dt, on = ['country', 'date'], how = 'inner')

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

equity_mob_death_npi_dt = equity_mob_death_npi_dt.groupby('country').apply(lambda group: group.interpolate(method='linear'))
#### these three countries didn't trade on 2-17
#equity_mob_death_npi_dt.drop(equity_mob_death_npi_dt[(equity_mob_death_npi_dt.country=='United States')&(equity_mob_death_npi_dt.date == '2020-02-17')].index, inplace = True)
#equity_mob_death_npi_dt.drop(equity_mob_death_npi_dt[(equity_mob_death_npi_dt.country=='Canada')&(equity_mob_death_npi_dt.date == '2020-02-17')].index, inplace = True)
#equity_mob_death_npi_dt.drop(equity_mob_death_npi_dt[(equity_mob_death_npi_dt.country=='South Africa')&(equity_mob_death_npi_dt.date == '2020-02-17')].index, inplace = True)


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

union_dates = union_dates.sort_values('date')
union_dates = union_dates.reset_index(drop = True)
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_dt = pd.merge(equity_mob_death_npi_dt, union_dates, on = ['date', 'country'], how = 'right')
g_equity_mob_death_npi_dt = g_equity_mob_death_npi_dt.sort_values(['country', 'date'])

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

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

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

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

#g_equity_mob_death_npi_dt.to_csv(r'g_equity_mob_death_npi_dt.csv', index = False)

# to make weighted index

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

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

###### interpolate dates ####
g_equity_mob_death_npi_dt = g_equity_mob_death_npi_dt[['country', 'date', 'workplace_percent_change_from_baseline',
       'stock market percentage change from Feb.17', 'death',
       'Government Response Index', 'Economic Support Index',
       'Stringency Index','log_dt7']]

## global
w_equity_mob_cases_npi_dt = pd.merge(g_equity_mob_death_npi_dt, weight, on = 'country', how = 'outer')

w_equity_mob_cases_npi_dt = w_equity_mob_cases_npi_dt.drop_duplicates()
#w_equity_mob_cases_npi_dt.to_csv('w_equity_mob_cases_npi_dt.csv')

# calculated market-cap weighted equity deviation
w_equity_mob_cases_npi_dt['stock market percentage change from Feb.17'] = w_equity_mob_cases_npi_dt['stock market percentage change from Feb.17']*w_equity_mob_cases_npi_dt['weight']
# calculated market-cap weighted mobility deviation
w_equity_mob_cases_npi_dt['w_workplace_percent_change_from_baseline'] = w_equity_mob_cases_npi_dt[ 'workplace_percent_change_from_baseline']*w_equity_mob_cases_npi_dt['weight']

# calculate death and governemnt response index
w_equity_mob_cases_npi_dt['death'] = w_equity_mob_cases_npi_dt['death']*w_equity_mob_cases_npi_dt['weight']
w_equity_mob_cases_npi_dt['log_dt7'] = w_equity_mob_cases_npi_dt['log_dt7']*w_equity_mob_cases_npi_dt['weight']

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


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

global_weighted['date'] = global_weighted.index

global_weighted.to_excel(r'global weighted for regression.xlsx')
global_weighted = global_weighted.rename(columns = {'stock market percentage change from Feb.17':'g_smd', 
                                         'w_workplace_percent_change_from_baseline': 'g_wmd',
                                         'death': 'g_death', 'Government Response Index':'g_gov',
                                         'Economic Support Index': 'g_econ','Stringency Index': 'g_strin', "log_dt7":"g_log_dt7"})

global_weighted = global_weighted.reset_index(drop = True)
equity_mob_death_npi0_dt = equity_mob_death_npi0_dt.rename(columns = {'stock market percentage change from Feb.17': 'smd', 
                                                    'workplace_percent_change_from_baseline': 'wmd', 'Economic Support Index':'econ',
                                                   'Stringency Index':'strin' })


equity_mob_death_npi_dt_global = pd.merge(equity_mob_death_npi0_dt, global_weighted, on = 'date', how = 'outer')

equity_mob_death_npi_dt_global = equity_mob_death_npi_dt_global.sort_values(['country', 'date'])

equity_mob_death_npi_dt_global = equity_mob_death_npi_dt_global[['country', 'date','smd', 'wmd', 'econ', 'log_dt7',
                                                                             'g_log_dt7', 'death','strin', 'g_death', 'g_smd',
                                                                             'g_wmd', 'g_gov', 'g_econ', 'g_strin']]

equity_mob_death_npi_dt_global.to_csv(r'table2 regression data.csv', index = False)

equity_mob_death_npi_dt_global = equity_mob_death_npi_dt_global.drop_duplicates()

equity_mob_death_npi_dt_global  = equity_mob_death_npi_dt_global [['date', 'country', 'smd', 'wmd']]
equity_mob_death_npi_dt_global.to_csv(r'table1 regression data_short period.csv', index = False)


