# 选择起止时间
start='2008-01-01'
end='2018-12-01'

# 选择指数
index=['000300.XSHG']

# 选择统计事件触发后多久的走势
later_days=15

from jqdata import *
import jqdata
import pandas as pd
import datetime 

# 全局变量
# 交易日期列表
all_trade_days=[i.strftime('%Y-%m-%d') for i in list(get_all_trade_days())]
# 作图尺寸
figsize=(16,8)

# 判断是否触及布林线上轨或下轨
def boll_filter(stk_list,date,up_down='down',ix=-20,k_std=2,field='close',fqy='1d'):
    df=get_price(security=stk_list,count=-ix,end_date=date,fields=[field],frequency=fqy)[field,:,:].iloc[1:,:]
    df=pd.DataFrame({'now':df.loc[date],'mean':df.mean(),'std':df.std()})
    df['up']=df['mean']+k_std*df['std']
    df['down']=df['mean']-k_std*df['std']
    if up_down=='down':
        rlt=list(df[df['now']<=df['down']].index)
    elif up_down=='up':
        rlt=list(df[df['now']>=df['up']].index)
    else:
        print('wrong up_down')
    return rlt

# 统计事件触发后股价变化
def after_story(totalLog,start,end,ahead_days=0,later_days=15,baseIndex=''):
    # 选择要统计触发事件后多少天的股价变化情况
    n_limt=later_days


    # 获取起止时间内的交易日
    days=[i.strftime('%Y-%m-%d') for i in list(get_trade_days(start,end))]

    # 选择基准指数
    # 如果选择了基准指数则统计的股票涨跌幅是减去了该指数涨跌幅的相对涨跌幅
    # 如果想看单纯股票涨跌幅，可以赋值为空字符串''
    baseIndex=''

    # 把事件数据赋值给stk
    stk=totalLog
    # 初始化存储统计结果的变量
    totalSum=pd.DataFrame()

    # 循环触发了事件的每一天
    for i in stk.keys():
        if i<start or i>end:
            continue
    #     print(i)
        # 获得当天触发事件的股票
        stk_list=stk[i]
        # 如果选择了基准指数
        if baseIndex!='':
            # 获取当天触发事件后股票相对基准指数的涨跌幅
            t=get_pct_change_accum([baseIndex]+stk_list,n_limt,all_trade_days[all_trade_days.index(i)-ahead_days],fqy='1d')
            t.reset_index(drop=True,inplace=True)
            t=t.sub(t[baseIndex],axis=0)
            del t[baseIndex]
        # 如果没选基准指数
        else:
            # 获取当天触发事件后股票的涨跌幅
            t=get_pct_change_accum(stk_list,n_limt,all_trade_days[all_trade_days.index(i)-ahead_days],fqy='1d')
            t.reset_index(drop=True,inplace=True)
            t.columns=[i]*len(t.columns)
#             t.columns=[i[:4]]*len(t.columns)
        # 记录当天触发事件后股票的涨跌幅，存储在totalSum
        totalSum=pd.concat([totalSum,t],axis=1)
    # 展示统计结果
    # 格式为dataframe，字段为股票代码，索引为事件后第几天，当天为0，值为涨跌幅
    # 设置作图大小
    
    # 生成事件触发后股价平均涨跌幅走势图
    totalSum.T.mean().plot(figsize=figsize).grid()
#     totalSum.T.mean().pct_change().plot(figsize=figsize).grid()
    print(len(totalSum.columns))
    return totalSum

# 获取相对baseday的ix个交易日量价累计涨跌幅
def get_pct_change_accum(scu,ix,date,fqy='1d',field='close'):
    if ix<=0:
        df=get_price(security=scu,count=-ix+2,end_date=date,fields=[field],frequency=fqy)[field,:,:].iloc[1:,:]

    else:
        df=get_price(security=scu,count=ix+2,end_date=all_trade_days[all_trade_days.index(date)+ix],fields=[field],frequency=fqy)[field,:,:].iloc[1:,:]
    if df.index.empty==True:
        df=df.T
        df['0000-00-00']=[]
        df=df.T
        
    return df/df.iloc[0]

# 初始化存储统计结果的变量
totalLog=dict()
# 获取起止时间内的交易日
days=[i.strftime('%Y-%m-%d') for i in list(get_trade_days(start,end))]


# 循环统计起止时间内的交易日
for baseDay in days:

    stk_list=index
    date=baseDay
    
    
    stk_list=boll_filter(stk_list=stk_list,date=date,up_down='up')

    # 记录触发事件的股票，存在totalLog中
    if stk_list!=[]:
            totalLog[baseDay]=stk_list
# 对于连续的触发只保留第一次
t=totalLog.keys()
t.sort(reverse = True)
for i in t:
    w=all_trade_days[all_trade_days.index(i)-1]
    if w in t:
        del totalLog[i]
totalLog   

# 统计之后的平均走势
w=after_story(totalLog,start,end,later_days=later_days)

# 分年度，统计之后的平均走势
t=w
f=lambda x : x[:4]
t.columns=map(f,t.columns)
t.T.groupby(t.columns).mean()
t.T.groupby(t.columns).mean().T.plot(figsize=figsize).grid()

# 初始化存储统计结果的变量
totalLog=dict()
# 获取起止时间内的交易日
days=[i.strftime('%Y-%m-%d') for i in list(get_trade_days(start,end))]


# 循环统计起止时间内的交易日
for baseDay in days:
    
    stk_list=index
    date=baseDay
    
    
    stk_list=boll_filter(stk_list=stk_list,date=date,up_down='down')

    # 记录触发事件的股票，存在totalLog中
    if stk_list!=[]:
            totalLog[baseDay]=stk_list
# 对于连续的触发只保留第一次
t=totalLog.keys()
t.sort(reverse = True)
for i in t:
    w=all_trade_days[all_trade_days.index(i)-1]
    if w in t:
        del totalLog[i]
totalLog

# 统计之后的平均走势
w=after_story(totalLog,start,end,later_days=later_days)

# 分年度，统计之后的平均走势
t=w
f=lambda x : x[:4]
t.columns=map(f,t.columns)
t.T.groupby(t.columns).mean()
t.T.groupby(t.columns).mean().T.plot(figsize=figsize).grid()