import pandas as pd
import numpy as np
import datetime
import dateutil
import time
import math
from sklearn import linear_model
def init(context):
    set_benchmark('000001.SH')
    log.info('策略开始运行,初始化函数全局只运行一次')
    set_commission(PerShare(type='stock',cost=0.0002))
    set_slippage(PriceSlippage(0.005))
    set_volume_limit(0.25,0.5)
    context.n = 40
    get_iwencai('沪深 A 股')
    g.qtr_cc = ['20000']
    g.jd = ['ss']
    g.js = [1]
    g.rti_pd = [1]

def before_trading(context):
    date1 = get_datetime()
    if int(date1.strftime('%w')) == 1 or len(g.qtr_cc) == 1:
        date2=date1.strftime("%Y%m%d")
        year=int(date1.strftime('%Y'))
        m = date2[4:8]
        month = int(m)
        if 100 < month < 501:
            t=4
            year = year -1
        elif 500 < month < 901:
            t = 1
        elif 900 < month < 1101:
            t=2
        elif 1100 < month < 1233:
            t = 3
        time = []
        for i in range(2008,2021):
            for j in range(1,5):
                tt = '%sq%s'%(i,j)
                time.append(tt)
        sj = []
        gang_sj = []
        t_list = ['0101','0331','0401','0630','0701','0930','1001','1231']
        for i in range(13):
            for j in range(0,7,2):
                x = 2008+i
                y = t_list[0+j]
                z = t_list[1+j]
                t1 = '%s%s'%(x,y)
                t2 = '%s%s'%(x,z)
                kk = get_trade_days(t1,t2).strftime('%Y%m%d').tolist()
                kk1 = get_trade_days(t1,t2).strftime('%Y-%m-%d').tolist()
                sj.append(kk[0])
                sj.append(kk[-1])
                gang_sj.append(kk1[0])
                gang_sj.append(kk1[-1])
        qtr_now ='%sq%s'%(year,t)
        if qtr_now != g.qtr_cc[-1]:
            g.qtr_cc.append(qtr_now)
            qtr_index = time.index(qtr_now)
            date_index = qtr_index * 2 + 1
            time_new = time[qtr_index-11:qtr_index+1]
            sj_new = sj[date_index-23:date_index+1]
            gsj_new = gang_sj[date_index-23:date_index+1]
            sjnew = sj_new[1:24:2]
            gang_sjn = gsj_new[1:24:2]
            stc_list = []
            for stc in context.iwencai_securities:
                if get_security_info(stc).start_date <date1-datetime.timedelta(days=1095):
                    stc_list.append(stc)
            df_roe = pd.DataFrame()
            q_roe = query(profit_sq.symbol,profit_sq.roe).filter(profit_sq.symbol.in_(stc_list))
            for i in time_new:
                roe = get_fundamentals(q_roe,statDate=i)
                roe['date']=i
                df_roe = df_roe.append(roe)
            df_roe = df_roe.sort_values(by=['profit_sq_stat_symbol','date'],ascending=[True,True])
            df_roe.columns=['symbol','roe','date']
            df_roe = df_roe.dropna()
            lirun_df = pd.DataFrame()
            for i in time:
                lr = get_fundamentals(query(income_sq.symbol,income_sq.profit_before_tax).filter(income_sq.symbol.in_(stc_list)),statDate=i)
                lr['date']=i
                lirun_df = lirun_df.append(lr)
            lirun_df = lirun_df.sort_values(by=['income_sq_stat_symbol','date'],ascending=[True,True])
            lirun_df.columns=['symbol','profit_tax','date']
            lirun_df = lirun_df.dropna()
            total_df = pd.DataFrame()
            for i in time_new:
                total = get_fundamentals(query(balance.symbol,balance.total_assets
).filter(balance.symbol.in_(stc_list)),statDate=i)
                total['date']=i
                total_df = total_df.append(total)
            total_df = total_df.sort_values(by=['balance_stat_symbol','date'],ascending=[True,True])
            total_df.columns=['symbol','total_asse','date']
            total_df = total_df.dropna()
            gpa = pd.merge(lirun_df,total_df,on=['symbol','date']) #合并利润，和资产
            gpa['gpa'] = gpa['profit_tax']/gpa['total_asse']
            gpa = gpa.drop(['profit_tax','total_asse'],axis=1) #删去多余列
            gpa = gpa.dropna()    
            cash_df = pd.DataFrame() #获取经营性现金流净值
            for i in time_new:
                cash = get_fundamentals(query(cashflow_sq.symbol,cashflow_sq.net_cash_flows_from_opt_act
).filter(cashflow_sq.symbol.in_(stc_list)),statDate=i)
                cash['date'] = i
                cash_df = cash_df.append(cash)
            cash_df = cash_df.sort_values(by=['cashflow_sq_stat_symbol','date'],ascending=[True,True])
            cash_df.columns=['symbol','cash','date']
            cash_df = cash_df.dropna()
            ape = pd.merge(cash_df,total_df,on=['symbol','date'])
            ape['ape'] = ape['cash']/ape['total_asse']
            ape = ape.drop(['cash','total_asse'],axis=1)
            ape = ape.dropna()
            pb_df = pd.DataFrame() #
            for i in stc_list:
                pb = get_factors(query(factor.date,factor.pb).filter(factor.symbol == i,factor.date.in_(gang_sjn)))
                pb['symbol'] = i
                pb_df = pb_df.append(pb)
                pb_df = pb_df.dropna()
            for x in range(0,12):
                pb_df = pb_df.replace(gang_sjn[x],time_new[x])
                pb_df['bm']=1/pb_df['factor_pb']
                pb_df.columns=['date','drop','symbol','bm']
                pb_df = pb_df.drop('drop',axis=1)
                pb_df = pb_df.dropna()
                gl_df = pd.DataFrame()
            for i in time:
                guli = get_fundamentals(query(cashflow_sq.symbol,cashflow_sq.dividend_etc_cp
).filter(cashflow_sq.symbol.in_(stc_list)),statDate=i)
                guli['date']=i
                gl_df = gl_df.append(guli)
                gl_df = gl_df.sort_values(by=['cashflow_sq_stat_symbol','date'],ascending=[True,True])
                gl_df.columns=['symbol','guli','date']
                gl_df = gl_df.dropna()
                tm_df = pd.DataFrame()
            for i in time:
                mv = get_fundamentals(query(asharevalue.symbol,asharevalue.total_mv).filter(asharevalue.symbol.in_(stc_list)),statDate=i)
                mv['date']=i
                tm_df = tm_df.append(mv)
            tm_df = tm_df.sort_values(by=['asharevalue_stat_symbol','date'],ascending=[True,True])
            tm_df.columns=['symbol','mv','date']
            tm_df = tm_df.sort_values(by=['asharevalue_stat_symbol','date'],ascending=[True,True])
            tm_df.columns=['symbol','mv','date']
            tm_df = tm_df.append(mv)
            tm_df = tm_df.dropna()
            npy = pd.merge(gl_df,tm_df,on=['symbol','date'])
            npy['npy'] = npy['guli']/npy['mv']
            npy = npy.drop(['guli','mv'],axis=1)
            npy = npy.dropna()
            p = pd.DataFrame()
            for i in sj_new:
                price = get_price(stc_list,None,i,'1d',['close'],True,None,1,is_panel=1)
                p = p.append(price['close'])
            p = p.pct_change()
            p['ind'] = p.index.values
            p = p[p['ind'].isin(sjnew)]
            p=p.fillna(0)
            p= p.drop(['ind'],axis=1)
            p1 = p.stack().unstack(0).stack()
            tuplee = p1.index.values
            close_r = p1.values.tolist()
            close_sj = []
            close_code = []
            for i in tuplee:
                ii = pd.to_datetime(i[1]).strftime('%Y-%m-%d')
                close_sj.append(ii)
                close_code.append(i[0])
            dictt = {'symbol':close_code,'date':close_sj,'r':close_r}
            close_df = pd.DataFrame(dictt)
            close_df = close_df.sort_values(by=['symbol','date'],ascending=[True,True])
            for x in range(0,12):
                close_df = close_df.replace(gang_sjn[x],time_new[x])
            all = pd.merge(df_roe,gpa,on=['symbol','date'])
            all = pd.merge(all,ape,on=['symbol','date'])
            all = pd.merge(all,npy,on=['symbol','date'])
            all = pd.merge(all,close_df,on=['symbol','date'])
            all = pd.merge(all,pb_df,on=['symbol','date'])
            all = all.dropna() #去空
            count = all['symbol'].values.tolist()
            kd = []
            for x in count:
                gp = all[all['symbol'] == x]
                if len(gp) == 12:
                    kd_st = gp['symbol'].values.tolist()[1]
                    kd.append(kd_st)
            kd = list(set(kd))
            all = all[all['symbol'].isin(kd)]
            all['new'] = all['r'].shift(-1)
            for zb in ['gpa','ape']:
                s_l = all[all[zb] < 0 ]
                s_l = list(set(s_l['symbol'].values.tolist()))
                all = all[~all['symbol'].isin(s_l)]
            all['nroe'] = all['roe'].rolling(window=2).mean()
            all['ngpa'] = all['gpa'].rolling(window=2).mean()
            all['nape'] = all['ape'].rolling(window=2).mean()
            all['nnpy'] = all['npy'].rolling(window=2).mean()
            all['nbm'] = all['bm'].rolling(window=2).mean()
            all = all.drop(['roe','gpa','ape','npy','bm','r'],axis=1)
            all = all.dropna()
            tccc = [time_new[-1]]
            dao = all[all['date'].isin(tccc)]
            all = all[~all['date'].isin(tccc)]
            single = dao['symbol'].values.tolist()
            k = []
            single_c = []
            for x in single:
                zg = all[all['symbol'] == x]
                if len(zg) > 0:
                    code = zg['symbol'].values.tolist()[0]
                    single_c.append(code)
                    reg = linear_model.LinearRegression()
                    reg.fit(zg.drop(['symbol','date','new'],axis=1),zg['new'])
                    k.append(reg.coef_)
            name = ['broe','bgpa','bape','bnpy','bbm']
            dao1 = pd.DataFrame(k,columns=name)
            dao1['symbol'] = single_c
            fdf = pd.merge(dao,dao1,on='symbol')
            fdf = fdf.dropna()
            fdf['fir'] = fdf['broe']*fdf['nroe'] +fdf['bgpa']*fdf['ngpa'] +
            fdf['bape']*fdf['nape']+fdf['bnpy']*fdf['nnpy']+fdf['bbm']*fdf['nbm']
            fdf = fdf.sort_values(by=['fir'],ascending=False)
            fir_list = fdf['fir'].values.tolist()
            symbol_list = fdf['symbol'].values.tolist()
            fir_dict = {'fir':fir_list,'symbol':symbol_list}
            fir_df = pd.DataFrame(fir_dict)
            g.jd = fir_df
        if int(date1.strftime('%w')) == 1 or len(g.js) == 1:
            datetime11 = get_last_datetime()
            date=datetime11.strftime('%Y%m%d')
            stock_list=context.iwencai_securities
            a=history(stock_list, ['close'], bar_count=375, fre_step='1d', skip_paused = False, fq = 'pre', is_panel=1)
            b=a['close']
            c=list((b.iloc[-1]/b.iloc[0])*100)
            df=pd.DataFrame()
            df['symbol']=stock_list
            df['r']=c
            df1=df[df['r']>100]
            g.df1=df1.sort_values(by='r',ascending=False)
        if int(date1.strftime('%w')) == 1 or len(g.js) == 1:
            fdf = g.jd
            fdf=fdf[fdf['fir']>0].head(70)
            stockn = pd.merge(fdf,g.df1,on=['symbol'])
            stockm=stockn.sort_values(by=['fir','r'],ascending=(False,False)).head(40)
            g.stock_list=stockm['symbol'].values.tolist()
            log.info('选出的股票池为%s'%(g.stock_list))
            last_date=get_last_datetime()
            date=pd.to_datetime(last_date)
            day = last_date.strftime('%Y%m%d')
            ddf = get_price('000300.SH',None,day,'1d',['close'],True,None,5,is_panel=1)
            datelist = pd.to_datetime(ddf.index.values).strftime('%Y%m%d')
            std = pd.DataFrame(columns=g.stock_list)
            for i in datelist:
                stt ='%s %s'%(i,'9:30')
                end ='%s %s'%(i,'14:55')
                df = get_price(g.stock_list,stt,end,'5m',['close'],is_panel=0)
                bd_list = []
                for st in g.stock_list:
                    ff = df[st]
                    ff['lnc'] = ff['close'].apply(lambda close:math.log(close)).diff()
                    ff['lnc2'] = ff['lnc'].apply(lambda lnc:lnc ** 2)
                    bd = ff['lnc2'].sum()
                    bd_list.append(bd)
                std.loc[len(std)] = bd_list
                std_all_list = []
                std_all = pd.DataFrame(columns=g.stock_list)
                for stc in g.stock_list:
                    stdall = std[stc].mean()
                    std_all_list.append(stdall)
                std_all.loc[0] = std_all_list
                std_all['sum']=std_all.apply(lambda x:x.sum(),axis=1)
                g.sm=std_all['sum']
                g.stdd = std_all
            def handle_bar(context, bar_dict):
                date111 = get_datetime()
                lastdate=get_last_datetime().strftime('%Y%m%d')
                if len(list(context.portfolio.stock_account.positions))>0:
                    securities=list(context.portfolio.stock_account.positions)
                    for stock in securities:
                        price1=history(stock, ['close'], bar_count=1, fre_step='1d', skip_paused = False, fq = 'pre', is_panel=1)
                        price=price1['close']
                        if (context.portfolio.positions[stock].cost_basis/price[0])-1<-0.09:
                            order_target(stock,0)
                            log.info('%s 止损：%s'%(lastdate,stock))
                        price_bench= history('000300.SH',['close'],5,'1d',skip_paused = False, fq = None)
                        if (price_bench['close'][-1]/price_bench['close'][-5])-1 <-0.13:
                            if len(list(context.portfolio.stock_account.positions.keys()))>0:
                            for stock in list(context.portfolio.stock_account.positions.keys()):
                                order_target(stock,0)
                                log.info('%s 大盘突然下跌,卖出%s' %(lastdate,stock))
                        stock_list = []
                        for stc in g.stock_list:
                            if bar_dict[stc].is_st == 0 and bar_dict[stc].is_paused == 0: #去除停牌和 st
                            stock_list.append(stc)
                            if int(date111.strftime('%w')) == 1 or len(g.js) == 1:
                                if len(stock_list)>0:
                                if len(list(context.portfolio.stock_account.positions.keys())) > 0:
                                    for stock in list(context.portfolio.stock_account.positions.keys()):
                                    if stock not in stock_list:
                                        order_target(stock, 0)
                            if len(g.js) == 1 or g.rti_pd[-1] == 1:
                                if len(stock_list) > 0:
                                for stock in stock_list:
                                    if stock not in list(context.portfolio.stock_account.positions.keys()):
                                        if len(list(context.portfolio.stock_account.positions.keys())) <
context.n :
                                        q=float(g.stdd[stock]/g.sm)
                        order_value(stock,(context.portfolio.stock_account.total_value)*q)
                        else:
                        order_value(stock,context.portfolio.stock_account.available_cash)