# -*- coding: utf-8 -*-
"""
Created on Sat Aug 15 18:15:00 2020

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

path = r'...\replication package'
os.chdir(path)

## step 1 calculate aggregated China mobility
# import residents mobility data from Baidu
china_mobility = pd.read_csv(r'China_City Movement Intensity 0101-0502.csv')

# to match the English and Chinese city name
city_code = pd.read_csv(r'Index_City_CH_EN.csv')

# import the city-level population data.
population = pd.read_excel(r'Population_City_China2.xlsx')
cities = population['数据来源'].to_list() #sign the name to every city

c = re.compile('.+市')
#check the number of cities
cities = cities[1:]
cities2 = [c.findall(i) for i in cities]
cities2 = [i[0] for i in cities2]
cities2 = list(set(cities2)) #299

type_p = population['Unnamed: 0'].to_list()
type_p_com = [f[-2:] for f in type_p]

population['type'] = type_p_com
population2 = population[1:]
population2['city'] = cities2

population2 = population2.loc[population2.type=='常住']


population2.to_excel(r'test.xlsx',index = False)
### I print out the data to careful check the data availability. and code the availability year
population3 = pd.read_excel(r'test.xlsx')
population_tomatch = population3[['city', 'year of data', 'population']]#281

city_code_tomatch = city_code[['GbCity','City_CH','City_EN']]
city_code_tomatch = city_code_tomatch.rename(columns ={'City_CH':'city'})#342

share_city = pd.merge(population_tomatch,city_code_tomatch, on = 'city', how = 'inner' ) #244

china_mobility_city = china_mobility.GbCity_EN.to_list()
c2 = re.compile('.+\_')
china_mobility_city_simp = [c2.findall(f) for f in china_mobility_city]
china_mobility_city_simp = [f[0] for f in china_mobility_city_simp]
china_mobility_city_simp = [f[:-1] for f in china_mobility_city_simp]
china_mobility_city_simp2 = []
for i in china_mobility_city_simp:
    if len(i)>16:
        china_mobility_city_simp2.append(i)
    else:
        china_mobility_city_simp2.append(i[:-3])

china_mobility['city'] = china_mobility_city_simp2

share_city2 = share_city[['City_EN', 'year of data', 'population']]
share_city2 = share_city2.rename(columns = {'City_EN':'city'})

he = [f for f in share_city2 if f in china_mobility_city_simp2]# this is all the cities I'll work on 239

china_mobility_selected = china_mobility.loc[china_mobility.city.isin(he)]
china_mobility_selected[china_mobility_selected.city.duplicated()]

resident_act = pd.merge(china_mobility_selected,share_city2, on = 'city', how='inner' ) #248
resident_act.to_excel(r'test2.xlsx')

resident_act1 = pd.read_excel(r'test2.xlsx')
resident_act1['weight'] = resident_act1['population']/resident_act1['population'].sum()

results1=pd.DataFrame()
for index, row in resident_act1.iterrows():
    a =row[4:-2]*row[-1]
    results1 = results1.append(a)
    
results2 = results1.sum()
base = results2[:10].median()

mobility_deviation = results2[10:]
mobility_deviation = (mobility_deviation/base-1)*100 #make it percent deviation

mobility_deviation = pd.DataFrame(data = {'percent_deviation':mobility_deviation})
mobility_deviation = mobility_deviation.reset_index()
mobility_deviation = mobility_deviation.rename(columns = {'index':'date'})

mobility_deviation['date'] = pd.to_datetime(mobility_deviation['date'])
#mobility_deviation.to_excel(r'china mobility percent deviation from baseline.xlsx', index = False)
