# 导入需要的数据库
from jqfactor import *
from jqdata import *
import pandas as pd
import warnings  
warnings.filterwarnings('ignore') 
from tqdm import tqdm_notebook # 进度条工具
from tqdm import tqdm
import time # 用于计时
from datetime import date, datetime, timedelta
from sklearn import preprocessing, linear_model

# 获取日期列表（War：较get_trade_day()，可配置交易频率）
def get_tradeday_list(start,end,frequency=None,count=None):
    if count != None:
        df = get_price('000001.XSHG',end_date=end,count=count)
    else:
        df = get_price('000001.XSHG',start_date=start,end_date=end)
    if frequency == None or frequency =='day':
        return df.index
    else:
        df['year-month'] = [str(i)[0:7] for i in df.index]
        if frequency == 'month':
            return df.drop_duplicates('year-month').index # War：drop_duplicates默认保留第一个出现的值
        elif frequency == 'quarter':
            df['month'] = [str(i)[5:7] for i in df.index]
            df = df[(df['month']=='01') | (df['month']=='04') | (df['month']=='07') | (df['month']=='10') ]
            return df.drop_duplicates('year-month').index
        elif frequency =='halfyear':
            df['month'] = [str(i)[5:7] for i in df.index]
            df = df[(df['month']=='01') | (df['month']=='06')]
            return df.drop_duplicates('year-month').index 

# 设置起止时间（为提高运算速度，取13年以后数据，同样覆盖熊牛市）
start='2015-01-01'
end='2019-05-31'
# 设置调仓周期，月/季/半年
periods=(5,10,20)
# 设置分层数量
quantiles=10
# 股票池（为缩短计算时间，这里取中证500）
stkpool = '000905.XSHG'
# 并剔除选股日 ST/PT、上市不满一年、动量计算期间及选股日停牌）
# 因子预处理及中性化，按聚宽一级行业和市值中性化
nutural_way = ['jq_l1', 'market_cap']
# 获取日期列表
date_list = get_tradeday_list(start=start,end=end,count=None) # 获取回测日期间的所有交易日
date_list[1]

# 定义基础动量因子——动量因子（org）和结合均线的趋势动量因子（MA）
def factor_cal(pool,date,N,how=None):
    '''
    说明：N为往前获得N天的股票价格数据
    how='org'或者为空，则计算常规动量因子
    how='MA'，则计算结合均线的趋势动量因子
    '''
    df0 = get_price(pool,end_date=date,fq='post',count=N,fields=['close','paused'])
    # 动量计算期间若存在停牌，则剔除
    dfp = df0.iloc[1,:,:].sum() # 若停牌，paused=1
    d_stk = dfp[dfp>0].index.tolist() # 因期间停牌需要剔除的股票列表
#     print(str(date)+' 期间停牌股票为 '+str(d_stk[0:4])+'等')
    df =df0['close'].drop(d_stk,axis=1)
#     print('剔除期间停牌股票后，剩余股数'+str(df.shape[1]))
    if N==0:
        print('输入错误：统计周期N不可以为0')
    else:
        if how=='org' or how==None:
            far = df.iloc[-1,:]/df.iloc[0,:] - 1
        elif how=='MA':
            far = df.iloc[:,:].sum()/N/df.iloc[-1,:]
    return far
factor_cal(['000006.XSHE','000008.XSHE','000009.XSHE'],'2017-12-28',21,'MA')

# 定义剔除上市不满一年、选股日 ST/PT股票的函数
def screener(pool_org, date):
    '''
    注：此处date必须为string格式的日期
    '''
    pool = []
#     print('剔除上市时间断、ST/PT前股票数量'+str(len(pool_org)))
    # 剔除上市不满一年
    for s in pool_org:
        ipo_date = get_security_info(s).start_date
        period = datetime.strptime(date,'%Y-%m-%d').date()-ipo_date
        if period.days>=360:
            pool.append(s)
#     print('剔除上市不满一年股票后股票数量'+str(len(pool)))
    # 剔除选股日 ST/PT
    is_st = get_extras('is_st',pool,count=1,end_date=date)
    pool = [s for s in pool if not is_st[s][0]]
#     print('剔除选股日ST/PT股票后股票数量'+str(len(pool)))
    return pool

# 预设空的dataframe以存放因子值
factor_df_org = pd.DataFrame()
factor_df_MA = pd.DataFrame()
# 循环计算给定日期范围的因子值
mark = 1
for d in tqdm(date_list):
    pool_org = get_index_stocks(stkpool,date=d)
    pool = screener(pool_org,str(d.date()))
    far_org = factor_cal(pool,d,21,'org')
    far_MA = factor_cal(pool,d,21,'MA')
    if mark == 1:
        factor_df_org = far_org
        factor_df_MA = far_MA
        mark = 0
    else:
        factor_df_org = pd.concat([factor_df_org,far_org],axis=1,join='outer') # War：原模板此处写反了
        factor_df_MA = pd.concat([factor_df_MA,far_MA],axis=1)
#将columns更改为可以日期标签
factor_df_org.columns = date_list
factor_df_MA.columns = date_list
factor_df_MA.head(3)

factor_df_org.head(3)

# 对动量因子（org）做标准化及中性化处理
for date in tqdm(date_list):
    se = standardlize(factor_df_org[date], inf2nan=True) # 截面标准化处理，即对每列做标准化
    se = neutralize(se, how=nutural_way, date=date) # 按研报剔除原始因子值与市值、行业相关的部分
    factor_df_org[date] = se
# 进行转置，调整为分析可用的格式（列为股票代码，行为日期）
factor_df_org = factor_df_org.T
factor_df_org.head(3)

# 对结合均线的趋势动量因子（MA）做标准化及中性化处理
for date in tqdm(date_list):
    se = standardlize(factor_df_MA[date], inf2nan=True)
    se = neutralize(se, how=nutural_way, date=date)
    factor_df_MA[date] = se
factor_df_MA = factor_df_MA.T
factor_df_MA.head(3)

# 对动量因子（org）进行单因子分析
t0 = time.time() # 记录开始时间
far = analyze_factor(factor=factor_df_org, start_date=date_list[0], end_date=date_list[-1], weight_method='avg', industry='jq_l1', quantiles=quantiles, periods=periods,max_loss=0.3)
t1 = time.time()
print('完成单因子分析，耗时 %.2f 分钟' %round((t1-t0)/60,2))
far.create_full_tear_sheet(demeaned=False, group_adjust=False, by_group=False, turnover_periods=None, avgretplot=(5, 15), std_bar=False)

# 对趋势动量因子（MA）进行单因子分析
t0 = time.time()
far = analyze_factor(factor=factor_df_MA, start_date=date_list[0], end_date=date_list[-1], weight_method='avg', industry='jq_l1', quantiles=quantiles, periods=periods,max_loss=0.3)
t1 = time.time()
print('完成单因子分析，耗时 %.2f 分钟' %round((t1-t0)/60,2))
far.create_full_tear_sheet(demeaned=False, group_adjust=False, by_group=False, turnover_periods=None, avgretplot=(5, 15), std_bar=False)

# 在此仅对动量因子（org）剥离流动性进行分析
factor_df_org = factor_df_org.T # 将因子矩阵进行转置，以保持中性化代码结构一致
for date in tqdm(date_list):
    se = neutralize(factor_df_org[date], how='liquidity', date=date)
    factor_df_org[date] = se
factor_df_org = factor_df_org.T
factor_df_org.head(3)

# 对提纯动量因子进行单因子分析
t0 = time.time() # 记录开始时间
far = analyze_factor(factor=factor_df_org, start_date=date_list[0], end_date=date_list[-1], weight_method='avg', industry='jq_l1', quantiles=quantiles, periods=periods,max_loss=0.3)
t1 = time.time()
print('完成单因子分析，耗时 %.2f 分钟' %round((t1-t0)/60,2))
# 打印信息比率（IC）相关表
far.plot_information_table(group_adjust=False, method='rank')

# 对提纯动量因子进行单因子分析——各分位数平均收益图
far.plot_quantile_returns_bar(by_group=False, demeaned=0, group_adjust=False)

# 对提纯动量因子进行单因子分析——换手率表
far.plot_turnover_table()

# 对提纯动量因子进行单因子分析——累计收益走势图
# 计算指定调仓周期的各分位数每日累积收益
df = far.calc_cumulative_return_by_quantile(period=10)
# 进行数据展示
df.plot(figsize=(15,6))

# 获取单日Fama-French因子
def FFfactor(date,mkt_bench='000905.XSHG',rf=0.04):
    '''
    默认以中证500作为市场基准，可修改为其他指数
    默认年化无风险收益率为4%
    '''
    pool = get_index_stocks(mkt_bench,date=date)
    q = query(
        valuation.code,
        valuation.market_cap,
        (balance.total_owner_equities/valuation.market_cap).label("BTM")
    ).filter(valuation.code.in_(pool))
    cap_BTM = get_fundamentals(q,date)
    # 股票池中各股前一天收益率（注：count=3，避免未来函数）
    p = get_price(pool,end_date=date,fq='post',count=3,fields=['close'])['close'].iloc[0:2,:]
    r = pd.DataFrame(p.iloc[1,:]).T # 框架用以存放各股票单日收益率
    r.iloc[0,:]=np.diff(p,axis=0)/np.array(p.iloc[0,:])
    # 筛选出最小市值的1/3股票和最大市值的1/3股票
    small_cap = cap_BTM.sort_values('market_cap',ascending=True)['code'][:int(len(cap_BTM)/3)].tolist()
    big_cap = cap_BTM.sort_values('market_cap',ascending=True)['code'][len(cap_BTM)-int(len(cap_BTM)/3):].tolist()
    # 筛选出最小B/P的1/3股票和最大B/P的1/3股票
    LBP = cap_BTM.sort_values('BTM',ascending=True)['code'][:int(len(cap_BTM)/3)].tolist()
    HBP = cap_BTM.sort_values('BTM',ascending=True)['code'][len(cap_BTM)-int(len(cap_BTM)/3):].tolist()
    # 计算因子值SMB、HMI、MP
    SMB = r[small_cap].mean(axis=1)-r[big_cap].mean(axis=1)
    HMI = r[LBP].mean(axis=1)-r[HBP].mean(axis=1)
    bench_price = get_price(mkt_bench,end_date=date,fq='post',count=3).iloc[0:2,:]['close']
    MP = np.diff(bench_price)/bench_price[0]-rf/360
    return pd.DataFrame({"MP":MP,"SMB":SMB,"HMI":HMI})
# 演示
FFfactor('2019-08-08') 
# Tips：输出数据为查询日前一天的数据，日期同
# 若输入交易日当天日期，交易时间同输入日期，且每次执行返回不同值，则说明引入未来函数

# 在剥离流动性后提纯动量因子基础上剔除Fama-French三因子
factor_df = factor_df_org

# 设中证500作为市场基准，用以计算三因子
mkt_bench='000905.XSHG'
# 设定回归天数（即因子-因变量组数）为21天
reg_period=21

# 获取所有交易日三因子的值    
FF_factor = pd.DataFrame()
for i in tqdm(factor_df.index):
    FF_factor = FF_factor.append(FFfactor(i,mkt_bench=mkt_bench))
FF_factor.head(3)

# 计算剔除三因子影响后的残差
# 注意：输入的factor_df要求格式列为股票代码，行为日期
reg = linear_model.LinearRegression()
postFFfactor = pd.DataFrame()

for s in tqdm(factor_df.columns):
    fct = []
    for i in range(reg_period,len(factor_df)+1):
        y = array(factor_df[s][i-reg_period:i]) # 选取(i+1)-reg_period到i日因子，共reg_preriod天
        x = FF_factor.iloc[i-reg_period:i,:].values
        x = preprocessing.scale(x, axis=1)
        # 对y进行缺失值处理
        if np.isnan(y[reg_period-1]): # 若因子当天为nan，则将当日残差因子也设为nan
            fct = fct+[np.nan]
        else: # 若期间y取值有nan，则将nan用均值填充
            y = array(pd.DataFrame(y).fillna(np.nanmean(y))[0])
            reg.fit(x,y)
            fct = fct+[y[reg_period-1]-reg.predict(x[reg_period-1].reshape(1,-1))[0]] # 单样本predict需reshape
    postFFfactor[s] = fct
    postFFfactor.index = factor_df.index[20:] # 注：剔除FF因子后，因子列表会损失的行数等于回归天数
postFFfactor.head(3)

# 对残差动量因子进行单因子分析
t0 = time.time() # 记录开始时间
far = analyze_factor(factor=postFFfactor, start_date=date_list[0], end_date=date_list[-1], weight_method='avg', industry='jq_l1', quantiles=quantiles, periods=periods,max_loss=0.3)
t1 = time.time()
print('完成单因子分析，耗时 %.2f 分钟' %round((t1-t0)/60,2))
# 打印信息比率（IC）相关表
far.plot_information_table(group_adjust=False, method='rank')

# 对残差动量因子进行单因子分析——各分位数平均收益图
far.plot_quantile_returns_bar(by_group=False, demeaned=0, group_adjust=False)

# 对残差动量因子进行单因子分析——累计收益走势图
# 计算指定调仓周期的各分位数每日累积收益
df = far.calc_cumulative_return_by_quantile(period=10)
# 进行数据展示
df.plot(figsize=(15,6))

# 定义单只股票改造K线因子计算函数
def fct_alterK(date,s,N=21,frq='5m'):
    '''
    说明：N为交易日往前N天作为观察期，默认21天
    frq为获取交易额的时间段颗粒度，支持如下周期：'1m', '5m', '15m', '30m', '60m'，默认'5m'
    '''
    if get_price(s, count=1, end_date=date, fields=['paused'])['paused'].values[0]==1:
        fct_alterK = pd.DataFrame({"avg_return":np.nan,"prob_up":np.nan},index=[date])
    else:
        k_cnt = int(N*240/float(frq.split('m')[0])) # 获取观察周期颗粒度K线总数
        df0 = get_price([s], count = k_cnt, end_date=date, frequency=frq, fields=['open','close','money'],fq='post')
        # 对每只股票获取其改造后每根K线对应的成交额、收益率及涨跌符号
        limit = df0['money'][s].sum()/N # 定义阈值
        k_info = pd.DataFrame()
        agr_money = 0
        p_open = df0['open'][s][0]
        for i in range(k_cnt):
            agr_money = agr_money +df0['money'][s][i]
            if i<k_cnt-1:
                if agr_money+df0['money'][s][i+1]>limit:
                    p_close = df0['close'][s][i]
                    return_rate = p_close/p_open-1
                    k_info = k_info.append({'agr_money':agr_money,'return_rate':return_rate,
                                            'sign':np.sign(return_rate)},ignore_index=True)
                    agr_money = 0
                    p_open = p_close
            elif i==k_cnt-1:
                p_close = df0['close'][s][i]
                return_rate = p_close/p_open-1
                k_info = k_info.append({'agr_money':agr_money,'return_rate':return_rate,
                                        'sign':np.sign(return_rate)},ignore_index=True)
        # 统计N根改造K线的平均收益率，以及N根改造K线的上涨概率
        avg_return = np.mean(k_info['return_rate'])
        prob_up = len(k_info.loc[k_info['sign']>0])/len(k_info)
        fct_alterK = pd.DataFrame({"avg_return":avg_return,"prob_up":prob_up},index=[date])
    return fct_alterK
fct_alterK('2017-12-28','000006.XSHE',N=21,frq='5m')

# 计算每只股票每个交易日的改造K线因子
factor_df_alterK = pd.DataFrame()
# 循环计算给定日期范围的因子值
# 特别说明：for准点交稿，这边就暂时只跑改造K线的平均收益率因子，太费时了(T＿T)
fct_avg_r = pd.DataFrame()
#--fct_prob_up = pd.DataFrame()
# 特别说明：由于改造K线因子计算耗时，在此仅试算近150个交易日的因子值（算是包含了小牛熊市场吧）
for d in tqdm(date_list[-150:]):
    pool_org = get_index_stocks(stkpool,date=d)
    pool = screener(pool_org,str(d.date()))
    avg_r = pd.DataFrame()
#--    prob_up = pd.DataFrame()
    for s in pool:
        # 特别说明：考虑到计算耗时太久，将观察期临时设为一周（即5个交易日），并将K线颗粒度改为30min
        fct = fct_alterK(d,s,5,'30m')
        avg_r = avg_r.append(fct[['avg_return']])
#--        prob_up = prob_up.append(fct[['prob_up']])
    avg_r.index = pool
    avg_r.columns = [d]
#--    prob_up.index = pool
#--    prob_up.columns = [d]
    fct_avg_r = fct_avg_r.append(avg_r.T)
#--    fct_prob_up = fct_prob_up.append(prob_up.T)
fct_avg_r.head(3)

# 对改造K线因子进行单因子分析
t0 = time.time() # 记录开始时间
far = analyze_factor(factor=fct_avg_r, start_date=date_list[0], end_date=date_list[-1], weight_method='avg', industry='jq_l1', quantiles=quantiles, periods=periods,max_loss=0.5)
t1 = time.time()
print('完成单因子分析，耗时 %.2f 分钟' %round((t1-t0)/60,2))
# 打印信息比率（IC）相关表
far.plot_information_table(group_adjust=False, method='rank')

# 对改造K线因子进行单因子分析——各分位数平均收益图
far.plot_quantile_returns_bar(by_group=False, demeaned=0, group_adjust=False)

# 对改造K线因子进行单因子分析——累计收益走势图
# 计算指定调仓周期的各分位数每日累积收益
df = far.calc_cumulative_return_by_quantile(period=20)
# 进行数据展示
df.plot(figsize=(15,6))