# imports
import time
from datetime import datetime, timedelta
import jqdata as jq
import numpy as np
import pandas as pd
import math
from statsmodels import regression
import statsmodels.api as sm
import matplotlib.pyplot as plt
from jqfactor import get_factor_values
import datetime
#from jqlib.technical_analysis import *
from scipy import stats

logger = logging.getLogger()
logger.setLevel("INFO")

import matplotlib.pyplot as plt
%matplotlib inline

today = pd.to_datetime('2019-04-12').date()
date0 = '1990-01-01'
sec1= '000001.XSHE'
date_valid = '2005-03-04'
df_sec = get_all_securities()
df_sec.loc[sec1]

#取得复权数据
list_field = ['open', 'close', 'low', 'high', 'volume', 'money', 'factor', 'high_limit','low_limit', 'avg', 'pre_close', 'paused']
df_px_pre = get_price(sec1,start_date=date0, end_date=today, fields=list_field, fq='pre')
df_px_none = get_price(sec1,start_date=date0, end_date=today, fields=list_field, fq=None)
df_px_post = get_price(sec1,start_date=date0, end_date=today, fields=list_field, fq='post')

# 计算前复权因子
print(date_valid)
px_none = df_px_none.close.loc[date_valid]
px_pre = df_px_pre.close.loc[date_valid]
adj_valid = px_pre/px_none
print('jqdata 前复权='+str(adj_valid))
px_google = 2.05
adj_google = px_google/px_none
print('google 前复权='+str(adj_google))

# 数据来源作图
df_px_pre.close.plot(title=sec1+'_复权=pre')
px_valid = df_px_pre.loc[date_valid].close
print("{t} clsoe price={x}".format(t=date_valid, x=px_valid))

# jqdata 前复权因子作图
df_px_pre.factor.plot()

#jqdata 前复权factor 变化日期
df_adj = df_px_pre.factor.pct_change().replace(0, np.nan).dropna()
df_adj

# 读取事件
q=jq.query(jq.finance.STK_XR_XD).filter(jq.finance.STK_XR_XD.code==sec1)
df_div_split_data  = jq.finance.run_query(q)

#combine a_xr_date and dividend_arrival_date
df_xr_temp = df_div_split_data['a_xr_date']
df_xr_temp.update(df_div_split_data['dividend_arrival_date'])
df_div_split_data['a_xr_date']  = df_xr_temp
list_field_basic = ['code', 'board_plan_pub_date', 'board_plan_bonusnote', 'implementation_bonusnote']
list_field_div =  [ 'a_bonus_date','bonus_ratio_rmb']
list_field_split =  [ 'a_xr_date','dividend_ratio', 'transfer_ratio', 'dividend_arrival_date']
df_div_data = df_div_split_data[list_field_basic + list_field_div]
df_split_data = df_div_split_data[list_field_basic + list_field_split]

print('股息事件')
df_div = df_div_data[df_div_data.a_bonus_date.notnull()].set_index('a_bonus_date')
df_div = df_div.join(df_px_none[['close']])
df_div = df_div.join(df_px_none['close'].shift(-1).to_frame('close_preday'))
#df_div['div_yield'] = (df_div['bonus_ratio_rmb']/(df_div['bonus_ratio_rmb']+df_div['close']*10.))
df_div['div_yield'] = (df_div['bonus_ratio_rmb']/(df_div['close']*10.))
df_div.sort_index(ascending=True)

print('拆股事件')
df_split = df_split_data[df_split_data.a_xr_date.notnull()].set_index('a_xr_date')
df_split = df_split.join(df_px_none[['close']])
df_split['dividend_ratio'] = df_split['dividend_ratio'].fillna(0)
df_split['transfer_ratio'] = df_split['transfer_ratio'].fillna(0)
#df_div['div_yield'] = (df_div['bonus_ratio_rmb']/(df_div['bonus_ratio_rmb']+df_div['close']*10.))
df_split['split_ratio'] = (df_split['dividend_ratio']+df_split['transfer_ratio'])/10.
df_split.drop('dividend_arrival_date', axis=1).sort_index(ascending=True)

# div_yield + split_ratio
df_adj_data = pd.concat([df_div[['div_yield']], df_split[['split_ratio']]], axis=1).fillna(0).sort_index(ascending=True)
df_adj_valid = df_adj_data
df_adj_valid['adj_ratio'] = df_adj_data['div_yield'] + df_adj_data['split_ratio']
#df_adj_valid = df_adj_valid['2005-01-01':]
df_adj_valid['factor'] = (df_adj_valid.adj_ratio+1).cumprod()
df_adj_valid['factor'] = df_adj_valid['factor']/df_adj_valid['factor'].iloc[-1]
print(date_valid+'使用factor为0.2318， 验证jqdata')
print(df_adj_data[:date_valid].iloc[-1].factor)
df_adj_data['2000-01-01':]

# split_ratio only
df_adj_data = pd.concat([df_div[['div_yield']], df_split[['split_ratio']]], axis=1).fillna(0).sort_index(ascending=True)
df_adj_google = df_adj_data
df_adj_google['adj_ratio'] =  df_adj_google['split_ratio']
#df_adj_valid = df_adj_valid['2005-01-01':]
df_adj_google['factor'] = (df_adj_google.adj_ratio+1).cumprod()
df_adj_google['factor'] = df_adj_google['factor']/df_adj_google['factor'].iloc[-1]
print(date_valid+'若不考虑股息排放，factor为0.25292')
print(df_adj_google[:date_valid].iloc[-1].factor)
print('与google factor(0.32)相差40%，考虑应为google忽略了某些事件，数据错误')
df_adj_google['2000-01-01':]