import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import time
import warnings
from jqdata import *
warnings.filterwarnings('ignore')
sns.set_style("whitegrid")
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号

t_begin = time.time()
freq = 5 #计算时采用的分钟数据频次
N = 60*4/freq
frequency = str(freq) + 'm'
n = 5 #后续计算平滑值的窗口期
zz500 = get_index_stocks('000905.XSHG') #取出中证500股票
start_date = '2013-01-01'
end_date = '2019-03-27'
m_end_date = '2019-03-28' #由于聚宽提取分钟数据的方式，需要将提取分钟数据的日期延后一天
#读取数据上首次可能会花费很长的时间，6年多的5分钟频次的数据量比较大
#后续每个部分都有计时可以参考
t0 = time.time()
print('开始读取数据')
st = get_extras('is_st', zz500, start_date=start_date, end_date=end_date, df=True)
open_paused = get_price(zz500, start_date=start_date, end_date=end_date, frequency='daily', fields=['paused','open'], fq='pre')
paused = open_paused['paused'] == 1
nostock = open_paused['paused'].isnull()
newstock = (nostock + 0).astype('int').rolling(252, min_periods=1).max() == 1
log_minute_close = np.log(get_price(zz500, start_date=start_date, end_date=m_end_date, frequency=frequency, fields=['close'], fq='pre')['close'])
open_ = open_paused['open']
log_minute_close['date'] = log_minute_close.index.date
log_minute_close['time'] = log_minute_close.index.time
log_minute_close.set_index(['date','time'],inplace=True)
minute_return = log_minute_close.groupby(level=0).apply(lambda df:df.diff(1))
out = ~(st | paused | nostock | newstock) #剔除停牌和st股以及未上市和新上市一年内的股票
future_ret1 = open_.pct_change(1).shift(-1-1) #开盘买入以及调仓，因此以开盘价计算1日未来收益率，后续用于计算分组净值
future_ret5 = open_.pct_change(5).shift(-5-1) #同上，5日未来收益率，用于计算后续的IC
t1 = time.time()
print('数据读取完成，耗时 %s 分钟' %round(((t1-t0)/60),3))
#这里定义了一个plot table的函数，方便后续查看一些表
def print_table(table, name=None):
    """
    将 dataframe 表格更好看地输出

    参数：
    ----------
    table : pd.Series or pd.DataFrame
        用于输出的表格.
    name : str, optional
        表格左上的名.
    """
    if isinstance(table, pd.Series):
        table = pd.DataFrame(table)

    if isinstance(table, pd.DataFrame):
        table.columns.name = name

    prev_option = pd.get_option('display.float_format')

    display(table)
                     

t0 = time.time()
print('因子计算开始')
RDVar = minute_return.pow(2).groupby(level=0).sum(axis=0) 
RDSkew = (np.sqrt(N) * minute_return.pow(3).groupby(level=0).sum(axis=0)) / RDVar.pow(3/2)
RDKurt = N * minute_return.pow(4).groupby(level=0).sum(axis=0) / RDVar.pow(2)
#mad法去极值处理的矩阵运算，将大于/小于 md +/- 3 * mad的部分按大小排序拉回[md+3*mad,md+3.5*mad]/[md-3.5*mad,md-3*mad]区间
def winsorize(df, nsigma=3):
    md = df.median(axis=1)
    mad = 1.483 * (df.sub(md, axis=0)).abs().median(axis=1)
    up = df.apply(lambda k: k > md + mad * nsigma)
    down = df.apply(lambda k: k < md - mad * nsigma)
    df[up] = df[up].rank(axis=1, pct=True).multiply(mad * 0.5, axis=0).add(md + mad * nsigma, axis=0)

    df[down] = df[down].rank(axis=1, pct=True).multiply(mad * 0.5, axis=0).add(md - mad * (0.5 + nsigma), axis=0)

    return df
RDVar = winsorize(RDVar)
RDSkew = winsorize(RDSkew)                          
RDKurt = winsorize(RDKurt)  
RVol = np.sqrt(RDVar.rolling(n+1).mean() * 242)  #根据研报计算，平滑用了前五天和当日的值，因此rolling窗口期用n+1
RSkew = RDSkew.rolling(n+1).mean()
RKurt = RDKurt.rolling(n+1).mean()   
#对研报中的条件进行股票剔除，并去掉所有股票都没有因子值的日期
RDVar = RDVar[out]
RDSkew = RDSkew[out]
RDKurt = RDKurt[out]
RVol = RVol[out]
RSkew = RSkew[out]
RKurt = RKurt[out]
RDVar.dropna(axis=0,how='all',inplace=True)
RDSkew.dropna(axis=0,how='all',inplace=True)
RDKurt.dropna(axis=0,how='all',inplace=True)
RVol.dropna(axis=0,how='all',inplace=True)
RSkew.dropna(axis=0,how='all',inplace=True)
RKurt.dropna(axis=0,how='all',inplace=True)
t1 = time.time()
print('因子计算完成 耗时 %s 分钟' %round(((t1-t0)/60),3))

t0 = time.time()
print('因子特征报告生成开始')
def create_characteristic(factor, fac_name): 
    """
    创建因子特征报告，包括因子分布一览和百分位走势一览

    参数：
    ----------
    factor : pd.DataFrame
        用于计算特征的因子.
    fac_name : str,
        因子名，用于绘图时的标题.
    """
    fig,axes = plt.subplots(figsize=(18,12),ncols=2,nrows=1)
    quantile = pd.DataFrame()
    quantiles = [90,75,50,25,10]
    for q in quantiles:
        quantile[q] = factor.quantile(q/100,axis=1)
    quantile.plot(ax=axes[1])
    axes[1].set_title(fac_name + ' 中证500因子百分位走势一览')
    sns.distplot(factor.stack(),ax=axes[0])
    axes[0].set_title(fac_name + ' 中证500因子分布一览')
    return axes

create_characteristic(factor=RVol, fac_name='RVol')
create_characteristic(factor=RSkew, fac_name='RSkew')
create_characteristic(factor=RKurt, fac_name='RKurt')
t1 = time.time()
print('因子特征报告生成完成 耗时 %s 秒' %round((t1-t0),3))

t0 = time.time()
print('因子选股分档表现报告开始')
groups = 5
hold_period = 5
def create_groupret_report(factor,groups,hold_period,fac_name):
    """
    创建因子选股分档报告，因子分档的净值走势

    参数：
    ----------
    factor : pd.DataFrame
        用于计算的因子.
    groups : int
        分档的组数
    hold_period : int
        持仓期（调仓周期）
    fac_name : str,
        因子名，用于绘图时的标题.
    """
    grouped = factor.rank(axis=1,pct=True) * groups // 1  #因子分层
    period_group = pd.DataFrame()
    #这里取出每隔调仓期日期的分组情况并重复5次，因为调仓后在下一次调仓前各组持有的股票是不变的
    to_append = grouped.iloc[::hold_period,:]
    for i in range(5):
        period_group = period_group.append(to_append)
    period_group = period_group.sort_index().iloc[:factor.shape[0],:]
    period_group.index = factor.index[:period_group.shape[0]]  #把index调整回来
    period_group = period_group.reindex_like(factor)
    #计算分组收益率
    group_rets = {}
    for i in range(5):
        group_rets[i] = future_ret1[period_group == i].mean(axis=1)
    group_ret = pd.DataFrame.from_dict(group_rets,'columns')
    group_ret.dropna(axis=0,how='all',inplace=True)
    #在起点对齐，计算净值
    group_ret.iloc[0,:] = 0
    net_value = (group_ret + 1).cumprod()
    net_value.columns = ['Q1','Q2','Q3','Q4','Q5']
    fig,ax = plt.subplots(figsize=(18,12))
    net_value.plot(ax=ax,lw=2)
    ax.set_title(fac_name + ' 中证500因子指标选股分档表现')
    ax.set_xlabel('date')
    ax.set_ylabel('net value')
    return ax

create_groupret_report(factor=RVol, groups=groups, hold_period=hold_period, fac_name='RVol')
create_groupret_report(factor=RSkew, groups=groups, hold_period=hold_period, fac_name='RSkew')
create_groupret_report(factor=RKurt, groups=groups, hold_period=hold_period, fac_name='RKurt')
t1 = time.time()
print('因子选股分档表现报告完成 耗时 %s 秒' %round((t1-t0),3))

def minus_p(sr):    #计算负值比例
    return sr[sr < 0].count() / sr.count()

def create_ic_report(factor,future_ret,fac_name):
    """
    创建因子IC报告，包括IC统计量和IC走势图

    参数：
    ----------
    factor : pd.DataFrame
        用于计算的因子.
    future_ret : pd.DataFrame
        用于计算IC的未来收益率，这里选用了5日未来收益率，计算5日IC值
    fac_name : str,
        因子名，用于绘图时的标题.
    """
    ic = factor.corrwith(future_ret,axis=1)
    ic.name = 'IC'
    ic = ic.to_frame()
    ic_statistic = ic.agg(['mean','std','max','min',minus_p]).T
    ic_statistic.columns = ['IC均值','IC标准差','IC最大值','IC最小值','负IC占比']
    name1 = fac_name + '因子中证500选股-IC表现'
    print_table(ic_statistic,name=name1)

    ic['IC均值（滚动12）'] = ic['IC'].rolling(12).mean()
    fig,ax = plt.subplots(figsize=(18,12))
    ic.plot(ax=ax,lw=2)
    ax.set_title('中证500选股 ' + fac_name + ' 因子IC值走势一览')
    #按年度的IC统计量
    ic['year'] = ic.index.year
    ic.set_index('year',inplace=True)
    ic_sta_byyear = ic['IC'].groupby('year').agg(['mean','std','max','min',minus_p])
    ic_sta_byyear.columns = ['IC均值','IC标准差','IC最大值','IC最小值','负IC占比']
    ic_sta_byyear.loc['整体',:] = ic_statistic.values
    name2 = fac_name + '因子中证500选股-IC分年度表现一览'
    print_table(ic_sta_byyear,name=name2)

t0 = time.time()
print('IC分析开始')
create_ic_report(factor=RSkew, future_ret=future_ret5, fac_name='RSkew')
t1 = time.time()
print('IC分析完成 耗时 %s 秒' %round((t1-t0),3))

t0 = time.time()
factor = RSkew.copy()
fac_name = 'RSkew'
hold_period = 5
print('中证500选股 多-空策略部分报告开始')
def back(sr):   #计算最大回撤率
    return format((1 - sr / sr.cummax()).max(),'.2%')
def cum_ret(sr):    #计算累计收益率
    return format((sr.iloc[-1] - sr.iloc[0]) / sr.iloc[0],'.2%')
def annual_std(ret):    #计算年化波动率
    return (ret.std() * np.sqrt(252))
def annual_ret(sr):     #计算年化收益率
    return ((sr.iloc[-1]/sr.iloc[0])**(252/len(sr)) - 1)

#以下和上面分档部分的计算相同
grouped = factor.rank(axis=1,pct=True) * groups // 1
period_group = pd.DataFrame()
to_append = grouped.iloc[::hold_period,:]
for i in range(5):
    period_group = period_group.append(to_append)
period_group = period_group.sort_index().iloc[:factor.shape[0],:]
period_group.index = factor.index[:period_group.shape[0]]
period_group = period_group.reindex_like(factor)
group_rets = {}
for i in range(5):
    group_rets[i] = future_ret1[period_group == i].mean(axis=1)
group_ret = pd.DataFrame.from_dict(group_rets,'columns')
group_ret.dropna(axis=0,how='all',inplace=True)
group_ret.iloc[0,:] = 0
net_value = (group_ret + 1).cumprod()
net_value.columns = ['Q1','Q2','Q3','Q4','Q5']
#计算中证500的净值走势    
zz500_bench = future_ret1.mean(axis=1).loc[net_value.index]
zz500_bench.iloc[0] = 0
zz500_bench = (zz500_bench + 1).cumprod()
#取出多头 空头 多-空 和中证500的净值    
top_bottom = net_value[['Q1','Q5']]
top_bottom['tb'] = ((group_ret.iloc[:,0] - group_ret.iloc[:,-1]) + 1).cumprod()
top_bottom['bench'] = zz500_bench.values
top_bottom.columns = ['多头净值','空头净值','多-空净值（右轴）','中证500指数']
#净值走势绘图，将多-空净值纵坐标绘制于右轴    
fig,ax = plt.subplots(figsize=(18,12))
ax1 = ax.twinx()
top_bottom[['多头净值','空头净值','中证500指数']].plot(ax=ax,lw=2)
top_bottom['多-空净值（右轴）'].plot(ax=ax1,lw=2,legend=True,color='purple')
ax.legend(loc='upper left')
ax1.legend(loc='upper right')
ax.set_title('中证500成分股' + fac_name + '因子选股多-空策略净值走势一览')

tb_ret = (group_rets[0] - group_rets[4]).dropna(how='all')
table_name = '中证500成分股' + fac_name + '因子选股多-空策略分年度表现'
def create_table(to_table_nv,col,ret):
    """
    创建策略年度表现表格，包括累计收益率 最大回撤率 年化收益率 年化波动率 信息比率

    参数：
    ----------
    to_table_nv : pd.DataFrame
        用于计算净值走势df.
    col : str
        净值走势df中的要计算年度表现策略的列名，用于将其取出
    ret : pd.Series
        策略的收益率序列，用于计算年化波动率
    """
    to_table_nv['year'] = to_table_nv.index.year   #调整index用于之后计算年度表现
    byyear = to_table_nv.set_index('year')[col]   #取出传入的净值走势df中需要的列



    ret.index = byyear.index
    ratio_table = byyear.groupby(level=0).agg([cum_ret,back,annual_ret],raw=False)
    ratio_table['annual_std'] = ret.groupby(level=0).apply(annual_std)
    ratio_table['information rate'] = ratio_table['annual_ret']/ratio_table['annual_std']
    ratio_table['annual_ret'] = ratio_table['annual_ret'].apply(lambda x:format(x,'.2%'))
    ratio_table['annual_std'] = ratio_table['annual_std'].apply(lambda x:format(x,'.2%')) #这里不在之前的函数中输出为百分数主要是为了计算信息比率
    #计算整体的表现
    allyear = byyear.agg([cum_ret,back,annual_ret],raw=False).values
    allyear = np.append(allyear,ret.std() * np.sqrt(252))
    allyear = np.append(allyear,allyear[2]/allyear[3])
    allyear[2] = format(allyear[2],'.2%')
    allyear[3] = format(allyear[3],'.2%')
    ratio_table.loc['整体',:] = allyear
    ratio_table.columns = ['累计收益率','最大回撤率','年化收益率','年化波动率','信息比率']
    #print_table(ratio_table,name=table_name)
    return ratio_table
table = create_table(to_table_nv=top_bottom, col='多-空净值（右轴）', ret=tb_ret)
print_table(table, name=table_name)
t1 = time.time()
print('中证500选股 多-空策略部分报告完成 耗时 %s 秒' %round((t1-t0),3))

t0 = time.time()
print('多-中证500策略部分报告开始')
#先计算用于绘图和后续表格输出的净值走势df，包括多头净值和多-中证500净值
top = net_value['Q1'].to_frame()   
top_alpha = (group_rets[0] - future_ret1.mean(axis=1).loc[net_value.index]).dropna(how='all')  #多头相对中证500的超额收益率
top['top-bench'] = top_alpha
top['top-bench'].iloc[0] = 0
top['top-bench'] = top['top-bench'].add(1).cumprod()
top.columns = ['多头净值','多头-中证500净值']
top_alpha.name = '超额收益率（右轴）'
#净值和超额收益率绘图，超额收益率纵坐标绘制于右轴
fig,ax = plt.subplots(figsize=(18,12))
ax1 = ax.twinx()
top.plot(ax=ax,lw=3)
top_alpha.plot(ax=ax1,legend=True,color='g',lw=0.8)
ax.legend(loc='upper left')
ax.set_title('中证500成分股' + fac_name + '因子选股多-中证500策略净值走势一览')

top_ret = group_rets[0].dropna()   #多头收益率
table1 = create_table(to_table_nv=top, col='多头净值', ret=top_ret)
#调整一下多头年度策略表格的层次列名
table1.columns = pd.MultiIndex.from_product([['多头表现'],['累计收益率','最大回撤率','年化收益率','年化波动率','信息比率']])

table2 = create_table(to_table_nv=top, col='多头-中证500净值', ret=top_alpha)
#同上
table2.columns = pd.MultiIndex.from_product([['多-中证500表现'],['累计收益率','最大回撤率','年化收益率','年化波动率','信息比率']])
#拼接两个表格并输出
print_table(pd.concat([table1,table2],axis=1))

t1 = time.time()
print('多-中证500策略部分报告完成 耗时 %s 秒' %round((t1-t0),3))

t0 = time.time()
print('换手率报告开始')
#这里计算换手率的结果于研报在最小值上有比较大的出入，可能是采用的方式与研报不同
def cal_turnover(grouped, quantile, hold_period):
    """
    创建策略年度换手率分析表格，包括均值 最大值 最小值 标准差 累计值

    参数：
    ----------
    grouped : pd.DataFrame
        因子选股分档后的df，值为某天某股票所在的分组.
    quantile : int
        要计算的组别，根据之前的分层方式，若计算第k组，输入k-1
    hold_period : int
        调仓周期
    """
    #取出每次调仓时选中的所有股票
    quant_name_sets=grouped.iloc[::hold_period].apply(lambda k: set(k[k == quantile].index), axis=1)
    #作差，得到新一次调仓时相对于之前新增的股票
    new_names = (quant_name_sets - quant_name_sets.shift(1)).dropna()
    #新增股票相对于之前持仓股票数的比例即换手率
    quant_turnover = new_names.apply(lambda x: len(x))/\
                    quant_name_sets.apply(lambda x: len(x))
    quant_turnover.name = quantile
    return quant_turnover

def max_turn(sr):  #最大换手率
    return format(sr.max(),'.2%')
def min_turn(sr):  #最小换手率
    return format(sr.min(),'.2%')
def std_turn(sr):  #换手率标准差
    return format(sr.std(),'.2%')
def mean_turn(sr):  #换手率均值
    return format(sr.mean(),'.2%')
def cumsum_turn(sr):  #换手率累计值
    return sr.sum()
top_turnover = cal_turnover(grouped=grouped, quantile=0, hold_period=hold_period).dropna()
top_turnover.index = top_turnover.index.year
turn_table = top_turnover.groupby(level=0).agg([mean_turn,max_turn,min_turn,std_turn,cumsum_turn])
turn_table.loc['整体',:] = top_turnover.agg([mean_turn,max_turn,min_turn,std_turn,cumsum_turn]).values
turn_table.columns = ['均值','最大值','最小值','标准差','累计值']
name = '中证500选股换手率分年度表现'
print_table(turn_table, name=name)
t1 = time.time()
print('换手率报告完成 耗时 %s 秒' %round((t1-t0),3))

t_end = time.time()
print('所有报告生成完成 耗时 %s 分钟' %round(((t_end-t_begin)/60),3))