# -*- coding: utf-8 -*-
"""
Created on Sun Sep 20 14:17:51 2020

@author: dingq
"""
# -*- coding: utf-8 -*-
"""
Created on Fri Aug  7 17:35:52 2020

@author: dingq
"""

#import yfinance as yf
import pandas as pd
import numpy as np
import os

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

# to seperate AH_A, AH_H, A not AH, H not AH
a_share = pd.read_csv(r'A shares info.csv')

ah_share = a_share[['Stkcd', 'Commnt(AH交叉码)', 'Markettype']]

ah_share = ah_share.loc[ah_share['Commnt(AH交叉码)'].notnull()]#128

#2 and 8 are B shares
ah_share = ah_share.loc[ah_share.Markettype!=2] #127

ah_share = ah_share.loc[ah_share.Markettype!=8] #123

ah_share_a = ah_share.Stkcd.to_list()#123
ah_share_a = ['0' * (6-len(str(i))) + str(i) for i in ah_share_a] # This is all ah_a share codes

ah_share_h = ah_share['Commnt(AH交叉码)'].to_list()
ah_share_h = ['0' * (4-len(str(i)[:-2])) + str(i)[:-2] for i in ah_share_h] #This is all ah_h share codes
ah_share_h = pd.DataFrame({'code': ah_share_h})

## find only A shares
a_n_ah = a_share[['Stkcd', 'Commnt(AH交叉码)', 'Markettype']]
a_n_ah =  a_n_ah.loc[a_n_ah['Commnt(AH交叉码)'].isnull()]
a_n_ah = a_n_ah.loc[a_n_ah.Markettype!=2] #3986
a_n_ah = a_n_ah.loc[a_n_ah.Markettype!=8] #3930
a_n_ah = a_n_ah.Stkcd.to_list()

a_n_ah = ['0' * (6-len(str(i))) + str(i) for i in a_n_ah] # This is all ah_a share codes
a_n_ah = pd.DataFrame({'code': a_n_ah})

a_all = pd.read_csv(r'A shares price.csv') #4027

a_all = a_all.loc[(a_all.Markettype!=2)&(a_all.Markettype!=8)] #to delete b shares still 3932

# to generate a new variable ah, 1 for yes it is ah share, 0 for not
a_all = a_all.rename(columns = {'Clsprc':'close', 'Stkcd':'code', 'Trddt':'date', 'Dsmvosd':'marketcap'})

a_all['ah'] = 0
a_all.loc[a_all.code.isin(ah_share_a),'ah']= 1

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

a_all_cap1 = a_all.loc[a_all.date =='2020-01-02'] #119 here
# as 600849 is changed to 601607 after aquisation and merging
# 600956 新天绿能，first traded on 2020-06-29
# 688180 君实生物，first traded on 2020-07-15
# 688505 复旦张江，first traded on 2020-06-19
final_ah_a = a_all_cap1[['code','ah']]
final_ah_a = final_ah_a.reset_index(drop = True)
#final_ah_a.to_csv(r'final_a_code.csv', index = False)

#a_all_cap2 = a_all.loc[a_all.date =='2020-08-03'] #119 here
#a_all_cap2 = a_all_cap2.loc[a_all_cap2.code.isin(final_ah_a.code)]
a_all = a_all.loc[a_all.code.isin(final_ah_a.code.to_list())]

# to calculate weighted close
a_all_total_cap = a_all.loc[a_all.date == '2019-12-26', 'marketcap'].sum()
ah_all_total_cap = a_all.loc[(a_all.date == '2019-12-26')&(a_all.ah==1), 'marketcap'].sum()
nah_all_total_cap = a_all.loc[(a_all.date == '2019-12-26')&(a_all.ah==0), 'marketcap'].sum()

## for AH share return
a_ah_weighted = a_all.copy()
a_ah_weighted.loc[a_ah_weighted.ah==1, 'weight'] = a_ah_weighted.loc[a_ah_weighted.ah==1, 'marketcap']/ah_all_total_cap
a_ah_weighted.loc[a_ah_weighted.ah==0, 'weight'] = a_ah_weighted.loc[a_ah_weighted.ah==0, 'marketcap']/nah_all_total_cap
a_ah_weighted['close_w'] = a_ah_weighted['close']*a_ah_weighted['weight']

a_ah_weighted = a_ah_weighted[['date','ah','close_w']].groupby(['ah','date']).sum()
a_ah_weighted = a_ah_weighted.reset_index()

## for only A share daily price deviation from Jan.13
a_ah_baseline = a_ah_weighted.loc[(a_ah_weighted.ah==0) & (a_ah_weighted.date == '2020-01-13'), 'close_w']
a_ah_baseline = a_ah_baseline.to_list()[0]

a_ah_weighted.loc[a_ah_weighted.ah==0, 'stock market deviation'] = np.log(a_ah_weighted.loc[a_ah_weighted.ah==0, 'close_w']) - np.log(a_ah_baseline) 
a_ah_weighted = a_ah_weighted.loc[a_ah_weighted.ah == 0]
a_ah_weighted = a_ah_weighted.loc[a_ah_weighted.date>='2020-01-02']
a_ah_weighted = a_ah_weighted[['date', 'stock market deviation']]

add_dates0 = pd.date_range(start='1/24/2020', periods=1, freq='D')
add_dates1 = pd.date_range(start='1/27/2020', periods=5, freq='D')

add_dates0 = add_dates0.append(add_dates1)
add_dates = pd.DataFrame({'date':add_dates0})
add_dates['stock market deviation'] = np.nan

a_ah_weighted = a_ah_weighted.append(add_dates)
a_ah_weighted = a_ah_weighted.sort_values(by = 'date')
a_ah_weighted = a_ah_weighted.reset_index(drop = True)
a_ah_weighted['stock market deviation'] = a_ah_weighted['stock market deviation'] *100
a_ah_weighted.to_csv(r'only A shares stock market deviation.csv', index = False)

