#工具函数
#工具函数
import time
from datetime import datetime, timedelta
import jqdata
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

#设置画图样式
plt.style.use('ggplot')

#输入起止日期，返回所有自然日日期
def get_date_list(begin_date, end_date):
    dates = []
    dt = datetime.strptime(begin_date,"%Y-%m-%d")
    date = begin_date[:]
    while date <= end_date:
        dates.append(date)
        dt += timedelta(days=1)
        date = dt.strftime("%Y-%m-%d")
    return dates

#获取日期列表
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
        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 
        
def ret_se(start_date='2018-6-1',end_date='2018-7-1',stock_pool=None,weight=0):
    pool = stock_pool
    if len(pool) != 0:
        #得到股票的历史价格数据
        df = get_price(list(pool),start_date=start_date,end_date=end_date,fields=['close']).close
        df = df.dropna(axis=1)
        #获取列表中的股票流通市值对数值
        df_mkt = get_fundamentals(query(valuation.code,valuation.circulating_market_cap).filter(valuation.code.in_(df.columns)))
        df_mkt.index = df_mkt['code'].values
        fact_se =pd.Series(df_mkt['circulating_market_cap'].values,index = df_mkt['code'].values)
        fact_se = np.log(fact_se)
    else:
        df = get_price('000001.XSHG',start_date=start_date,end_date=end_date,fields=['close'])
        df['v'] = [1]*len(df)
        del df['close']
    #相当于昨天的百分比变化
    pct = df.pct_change()+1
    pct.iloc[0,:] = 1
    if weight == 0:
        #等权重平均收益结果
        se = pct.cumsum(axis=1).iloc[:,-1]/pct.shape[1]
        return se
    else:
        #按权重的方式计算
        se = (pct*fact_se).cumsum(axis=1).iloc[:,-1]/sum(fact_se)
        return se
    
#获取所有分组pct
def get_all_pct(pool_dict,trade_list,groups=5):
    num = 1
    for s,e in zip(trade_list[:-1],trade_list[1:]):
        stock_list = pool_dict[s]
        stock_num = len(stock_list)//groups
        if num == 0:
            pct_se_list = []
            for i in range(groups):
                pct_se_list.append(ret_se(start_date=s,end_date=e,stock_pool=stock_list[i*stock_num:(i+1)*stock_num]))
            pct_df1 = pd.concat(pct_se_list,axis=1)
            pct_df = pd.concat([pct_df,pct_df1],axis=0)
        else:
            pct_se_list = []
            for i in range(groups):
                pct_se_list.append(ret_se(start_date=s,end_date=e,stock_pool=stock_list[i*stock_num:(i+1)*stock_num]))
            pct_df = pd.concat(pct_se_list,axis=1)    
            num = 0
    return pct_df

def tradedays_before(date,count):#获取指定交易日往前推count天交易日
    date = get_price('000001.XSHG',end_date=date,count=count+1).index[0]
    return date

#进行新股、St股过滤，返回筛选后的股票
def filter_stock(stockList,date,days=21*3):
    
    #去除上市距beginDate不足3个月的股票
    def delect_stop(stocks,beginDate,n=days):
        stockList=[]
        beginDate = datetime.datetime.strptime(beginDate, "%Y-%m-%d")
        for stock in stocks:
            start_date=get_security_info(stock).start_date
            if start_date<(beginDate-datetime.timedelta(days=n)).date():
                stockList.append(stock)
        return stockList
    
    #剔除ST股
    st_data=get_extras('is_st',stockList, count = 1,end_date=date)
    stockList = [stock for stock in stockList if not st_data[stock][0]]
    
    #剔除停牌、新股及退市股票
    stockList=delect_stop(stockList,date)
    '''
    #剔除开盘涨跌停
    #如果需要收盘涨跌停可以改字段即可
    df = get_price(stockList,end_date=date,fields=['open','high_limit','low_limit'],count=1).iloc[:,0,:]
    df['h_limit']=(df['open']==df['high_limit'])
    df['l_limit']=(df['open']==df['low_limit'])
    stockList = [df.index[i] for i in range(len(df)) if not (df.h_limit[i] or df.l_limit[i])] #过滤涨跌停股票
    '''
    return stockList

#第一部分：数据获取，X,Y
#参数及工具函数部分
from jqfactor import *
import warnings  
warnings.filterwarnings('ignore') 

#因子数据获取函数
#获取时间为date的全部因子数据
def get_factor_data(stock,date):
    data=pd.DataFrame(index=stock)
    q = query(valuation,balance,cash_flow,income,indicator).filter(valuation.code.in_(stock))
    df = get_fundamentals(q, date)
    df['market_cap']=df['market_cap']*100000000
    factor_data=get_factor_values(stock,['roe_ttm','roa_ttm','total_asset_turnover_rate',\
                               'net_operate_cash_flow_ttm','net_profit_ttm','net_profit_ratio',\
                              'cash_to_current_liability','current_ratio',\
                             'gross_income_ratio','non_recurring_gain_loss',\
                            'operating_revenue_ttm','net_profit_growth_rate'],end_date=date,count=1)
    factor=pd.DataFrame(index=stock)
    for i in factor_data.keys():
        factor[i]=factor_data[i].iloc[0,:]
    df.index = df['code']
    data['code'] = df['code']
    del df['code'],df['id']
    #合并得大表
    df=pd.concat([df,factor],axis=1)
    #总市值取对数
    data['size_lg']=np.log(df['market_cap'])
    #净利润(TTM)/总市值
    data['EP']=df['net_profit_ttm']/df['market_cap']
    #净资产/总市值
    data['BP']=1/df['pb_ratio']
    #营业收入(TTM)/总市值
    data['SP']=1/df['ps_ratio']
    #净现金流(TTM)/总市值
    data['NCFP']=1/df['pcf_ratio']
    #经营性现金流(TTM)/总市值
    data['OCFP']=df['net_operate_cash_flow_ttm']/df['market_cap']
    #净利润同比增长率
    data['net_g'] = df['net_profit_growth_rate']
    #净利润(TTM)同比增长率/PE_TTM
    data['G/PE']=df['net_profit_growth_rate']/df['pe_ratio']
    #ROE_ttm
    data['roe_ttm']=df['roe_ttm']
    #ROE_YTD
    data['roe_q']=df['roe']
    #ROA_ttm
    data['roa_ttm']=df['roa_ttm']
    #ROA_YTD
    data['roa_q']=df['roa']
    #净利率
    data['netprofitratio_ttm'] = df['net_profit_ratio']
    #毛利率TTM
    data['grossprofitmargin_ttm']=df['gross_income_ratio']
    #毛利率YTD
    data['grossprofitmargin_q']=df['gross_profit_margin']

    #扣除非经常性损益后净利润率YTD
    data['profitmargin_q']=df['adjusted_profit']/df['operating_revenue']
    #资产周转率TTM
    data['assetturnover_ttm']=df['total_asset_turnover_rate']
    #总资产周转率YTD 营业收入/总资产
    data['assetturnover_q']=df['operating_revenue']/df['total_assets']
    #经营性现金流/净利润TTM
    data['operationcashflowratio_ttm']=df['net_operate_cash_flow_ttm']/df['net_profit_ttm']
    #经营性现金流/净利润YTD
    data['operationcashflowratio_q']=df['net_operate_cash_flow']/df['net_profit']
    #净资产
    df['net_assets']=df['total_assets']-df['total_liability']
    #总资产/净资产
    data['financial_leverage']=df['total_assets']/df['net_assets']
    #非流动负债/净资产
    data['debtequityratio']=df['total_non_current_liability']/df['net_assets']
    #现金比率=(货币资金+有价证券)÷流动负债
    data['cashratio']=df['cash_to_current_liability']
    #流动比率=流动资产/流动负债*100%
    data['currentratio']=df['current_ratio']
    #总市值取对数
    data['ln_capital']=np.log(df['market_cap'])
    #TTM所需时间
    his_date = [pd.to_datetime(date) - datetime.timedelta(90*i) for i in range(0, 4)]
    tmp = pd.DataFrame()
    tmp['code']=list(stock)
    for i in his_date:
        tmp_adjusted_dividend = get_fundamentals(query(indicator.code, indicator.adjusted_profit, \
                                                     cash_flow.dividend_interest_payment).
                                               filter(indicator.code.in_(stock)), date = i)
        tmp=pd.merge(tmp,tmp_adjusted_dividend,how='outer',on='code')

        tmp=tmp.rename(columns={'adjusted_profit':'adjusted_profit'+str(i.month), \
                                'dividend_interest_payment':'dividend_interest_payment'+str(i.month)})
    tmp=tmp.set_index('code')
    tmp_columns=tmp.columns.values.tolist()
    tmp_adjusted=sum(tmp[[i for i in tmp_columns if 'adjusted_profit'in i ]],1)
    tmp_dividend=sum(tmp[[i for i in tmp_columns if 'dividend_interest_payment'in i ]],1)
    #扣除非经常性损益后净利润(TTM)/总市值
    data['EPcut']=tmp_adjusted/df['market_cap']
    #近12个月现金红利(按除息日计)/总市值
    data['DP']=tmp_dividend/df['market_cap']
    #扣除非经常性损益后净利润率TTM
    data['profitmargin_ttm']=tmp_adjusted/df['operating_revenue_ttm']
    #营业收入(YTD)同比增长率
    #_x现在 _y前一年
    his_date = pd.to_datetime(date) - datetime.timedelta(365)
    name=['operating_revenue','net_profit','net_operate_cash_flow','roe']
    temp_data=df[name]
    his_temp_data = get_fundamentals(query(valuation.code, income.operating_revenue,income.net_profit,\
                                            cash_flow.net_operate_cash_flow,indicator.roe).
                                      filter(valuation.code.in_(stock)), date = his_date)
    his_temp_data=his_temp_data.set_index('code')
    #重命名 his_temp_data last_year
    for i in name:
        his_temp_data=his_temp_data.rename(columns={i:i+'last_year'})

    temp_data =pd.concat([temp_data,his_temp_data],axis=1)
    #营业收入(YTD)同比增长率
    data['sales_g_q']=temp_data['operating_revenue']/temp_data['operating_revenuelast_year']-1
    #净利润(YTD)同比增长率
    data['profit_g_q']=temp_data['net_profit']/temp_data['net_profitlast_year']-1
    #经营性现金流(YTD)同比增长率
    data['ocf_g_q']=temp_data['net_operate_cash_flow']/temp_data['net_operate_cash_flowlast_year']-1
    #ROE(YTD)同比增长率
    data['roe_g_q']=temp_data['roe']/temp_data['roelast_year']-1
    
    
    #计算beta部分
    #辅助线性回归的函数
    def linreg(X,Y,columns=3):
        X=sm.add_constant(array(X))
        Y=array(Y)
        if len(Y)>1:
            results = regression.linear_model.OLS(Y, X).fit()
            return results.params
        else:
            return [float("nan")]*(columns+1)
    #个股60个月收益与上证综指回归的截距项与BETA
    stock_close=get_price(list(stock), count = 12*20+1, end_date=date, frequency='daily', fields=['close'])['close']
    SZ_close=get_price('000001.XSHG', count = 12*20+1, end_date=date, frequency='daily', fields=['close'])['close']
    stock_pchg=stock_close.pct_change().iloc[1:]
    SZ_pchg=SZ_close.pct_change().iloc[1:]
    beta=[]
    stockalpha=[]
    for i in stock:
        temp_beta, temp_stockalpha = stats.linregress(SZ_pchg, stock_pchg[i])[:2]
        beta.append(temp_beta)
        stockalpha.append(temp_stockalpha)
    #此处alpha beta为list
    #data['alpha']=stockalpha
    data['beta']=beta
    
    #反转
    data['reverse_1m']=stock_close.iloc[-21]/stock_close.iloc[-1]-1
    data['reverse_3m']=stock_close.iloc[-63]/stock_close.iloc[-1]-1
    
    #波动率（一个月、三个月标准差）
    data['std_1m']=stock_close[-20:].std()
    data['std_3m']=stock_close[-60:].std()
    
    #换手率
    #tradedays_1m = get_tradeday_list(start=date,end=date,frequency='day',count=21)#最近一个月交易日
    tradedays_3m = get_tradeday_list(start=date,end=date,frequency='day',count=63)#最近一个月交易日
    data_turnover_ratio=pd.DataFrame()
    data_turnover_ratio['code']=list(stock)
    for i in tradedays_3m:
        q = query(valuation.code,valuation.turnover_ratio).filter(valuation.code.in_(stock))
        temp = get_fundamentals(q, i)
        data_turnover_ratio=pd.merge(data_turnover_ratio, temp,how='left',on='code')
        data_turnover_ratio=data_turnover_ratio.rename(columns={'turnover_ratio':i})
    data['turn_3m']= (data_turnover_ratio.set_index('code').T).mean()
    data['turn_1m']= (data_turnover_ratio.set_index('code').T)[-21:].mean()    
    '''
    #技术指标部分
    date_1 = tradedays_before(date,1)
    data['PSY']=pd.Series(PSY(stock, date_1, timeperiod=20))
    data['RSI']=pd.Series(RSI(stock, date_1, N1=20))
    data['BIAS']=pd.Series(BIAS(stock,date_1, N1=20)[0])
    dif,dea,macd=MACD(stock, date_1, SHORT = 10, LONG = 30, MID = 15)
    #data['DIF']=pd.Series(dif)
    #data['DEA']=pd.Series(dea)
    data['MACD']=pd.Series(macd)
    '''
    return data

#多期获取数据
#输入：指数名称,统计时间列表
#输出：字典，key为时间，vlaues为因子值数据

factor_data_dict = {}

#设置统计数据区间
index = '000905.XSHG' #中证500
date_start = '2009-1-1'
date_end   = '2019-1-1'

#获取交易日列表，每月首个交易日
date_list = get_tradeday_list(start=date_start,end=date_end,frequency='month',count=None)

#循环获取因子数据，并将值存在字典里面(耗时较长)
#并计算个期股票收益情况，加入指数涨跌幅
import time  
t1 = time.time()
for date_1,date_2 in zip(date_list[:-1],date_list[1:]):
    pool = get_index_stocks(index,date = date_1)
    pool = filter_stock(pool,str(date_1)[:10],days=21*3) #进行新股、ST股票过滤
    #计算指数涨跌幅
    df_1 = get_price(index,end_date=date_1,fields=['close'],count = 1)['close']
    df_2 = get_price(index,end_date=date_2,fields=['close'],count = 1)['close']
    index_pct = df_2.values[0]/df_1.values[0] - 1#具体数值
    
    #计算各股票涨跌幅
    df_1 = get_price(pool,end_date=date_1,fields=['close'],count = 1)['close']
    df_2 = get_price(pool,end_date=date_2,fields=['close'],count = 1)['close']
    df_3 = pd.concat([df_1,df_2],axis=0).T #进行合并
    stock_pct = df_3.iloc[:,1]/df_3.iloc[:,0] - 1 #计算pct，series
    
    #记录因子值
    factor_df = get_factor_data(pool,date=date_1) #进行因子值获取
    factor_df['init_pct'] = stock_pct
    factor_df['pct'] = stock_pct-index_pct
    del factor_df['code']
    
    #对数据进行处理、标准化、去极值、中性化
    factor_df = winsorize_med(factor_df, scale=3, inclusive=True, inf2nan=True, axis=0) #中位数去极值处理
    factor_df = standardlize(factor_df, inf2nan=True, axis=0) #对每列做标准化处理
    factor_df = neutralize(factor_df, how=['sw_l1', 'market_cap'], date=date_1, axis=0,fillna='sw_l1')#中性化
    factor_df['pct_init'] =  stock_pct-index_pct
    factor_df['pct_'] =  stock_pct
    factor_data_dict[date_1] = factor_df
t2 = time.time()
print('计算因子耗时：{0}'.format(t2-t1))
factor_data_dict[date_1].head(5)

#将计算出来的因子值进行存储
#使用pickle模块将数据对象保存到文件
import pickle
pkl_file = open('factor_data_10year.pkl', 'wb')
pickle.dump(factor_data_dict, pkl_file, 0)
pkl_file.close()

#读取计算出来的因子值
import pickle
pkl_file = open('factor_data_10year.pkl', 'rb')
factor_data_dict = pickle.load(pkl_file)
pkl_file.close()

#数据合并，用于模型训练
#将字典组成一个无日期的大表
all_data = pd.DataFrame()
for d in date_list[:60]:
    all_data = pd.concat([all_data,factor_data_dict[d]],axis=0) #进行纵向拼接
print(d)
print(all_data.shape)
all_data.head(5)

for colu in all_data.columns:
    if sum(all_data[colu].isnull())> 0.2*len(all_data[colu]):
        print('因子：%s缺失值较多，不做为统计特征处理'%colu)
        del all_data[colu]
all_data.shape

#进行空值处理
for colu in all_data.columns:
    if sum(all_data[colu].isnull())> 0:#0.05*len(all_data[colu]):
        all_data[colu].fillna(all_data[colu].mean(),inplace=True)
        #del all_data[colu]
all_data.shape

#第二部分：决策树模型
from sklearn import tree
from sklearn.model_selection import train_test_split #对数据进行随机划分

#进行因子值预测（分类方式）
df_y = pd.DataFrame()
df_score = pd.DataFrame(index = ['score_p','score_r'])
colus = factor_data_dict[date_list[60]].columns[:-4]
df_features_c = pd.DataFrame(index=colus)

for i in range(60,len(date_list)-1):
    print('正在计算{0}...数据'.format(date_list[i]))

    #数据合并，用于模型训练
    #将字典组成一个无日期的大表
    all_data = pd.DataFrame()
    for d in date_list[i-60:i]:
        all_data = pd.concat([all_data,factor_data_dict[d]],axis=0) #进行纵向拼接

    #空值NAN值处理
    for colu in all_data.columns:
        if sum(all_data[colu].isnull())> 0.2*len(all_data[colu]):
            print('因子：%s缺失值较多，不做为统计特征处理'%colu)
            del all_data[colu]
    #进行空值处理
    for colu in all_data.columns:
        if sum(all_data[colu].isnull())> 0:#0.05*len(all_data[colu]):
            all_data[colu].fillna(all_data[colu].mean(),inplace=True)
            #del all_data[colu]

    #获取前三分之一和后三分之一的数据
    all_data = all_data.sort_values('pct_init')
    #print(all_data.shape)
    all_data['num'] = list(range(len(all_data)))
    all_data = all_data[(all_data['num']<=len(all_data)/3) | (all_data['num']>=2*len(all_data)/3)]

    #print(all_data.shape)

    #记录标签y
    y1 = (all_data["pct_init"]>0).astype("int") #走完流程在改？
    y  = all_data['pct_init']
    x = all_data.iloc[:,:-5]
    
    X_train, X_test, y_train, y_test = train_test_split(x,y1.values, test_size=0.3)
    #建立最大深度为2的决策树，并用测试数据来训练这颗树
    clf = tree.DecisionTreeClassifier(max_depth = 6
                                         )
    clf = clf.fit(X_train, y_train) #模型训练
    s_p = clf.score(X_test,y_test)#样本内得分
    #print(clf.score(X_test,y_test)) #模型得分
    
    #模型预测
    factor_df_temp = factor_data_dict[date_list[i]].iloc[:,:-4]
    factor_df_temp1 = factor_df_temp.dropna(axis=0)
    
    df_temp = pd.DataFrame(index=factor_df_temp1.index)
    df_temp[date_list[i]] = [i[1] for i in clf.predict_proba(factor_df_temp1)]
    
    #记录模型真实数据得分
    y_t = factor_data_dict[date_list[i]].dropna()
    x_s = y_t.iloc[:,:-4]
    y_s = (y_t["pct_init"]>0).astype("int")
    s_r = clf.score(x_s,y_s)
    
    df_score[date_list[i]] = [s_p,s_r]
    #记录特征重要性
    df_features_c[date_list[i]] = clf.feature_importances_
    
    df_y = pd.concat([df_y,df_temp],axis=1) #进行横向拼接
df_y

factor_df = df_y.T

'''
选取头部股票计算超额收益
输入：factor_df，index为日期，column是股票名字，values是进行排序的因子值
'''
#选取头部股票构造组合进行回测
return_df = pd.DataFrame()
for d1,d2 in zip(factor_df.index[:-1],factor_df.index[1:]):
    #获取头部股票
    df_temp = factor_df.loc[d1,:].sort_values() #mo默认从小到大排序
    #pool_temp = df_temp.index[:50]   
    pool_temp = df_temp.index[-50:]
    #计算组合收益
    df1 = get_price(list(pool_temp),end_date=d1,count=1,fields=['close'])['close'] #index为日期，columns为股票名称
    df1 = df1.dropna(axis=1) #去掉NAN值，删除列
    df2 = get_price(list(df1.columns),end_date=d2,count=1,fields=['close'])['close']
    ret = (df2.values/df1.values - 1).mean() #计算组合收益均值
    #计算同期指数收益率
    df_index1 = get_price('000905.XSHG',end_date=d1,count=1,fields=['close'])['close']
    df_index2 = get_price('000905.XSHG',end_date=d2,count=1,fields=['close'])['close']
    index_ret = df_index2.values[-1]/df_index1.values[-1]-1
    
    return_df[d1] = [ret-index_ret]  #记录超额收益
    #print(ret)
return_df = return_df.T
return_df.columns = ['return']
print(return_df.head(5))
(return_df+1).cumprod().plot(figsize=(12,6))