from datetime import datetime,timedelta
import pandas as pd
import re

STRATEGY_NAME = "mom"

PRINT = 0 #控制打印

#下单函数,order_flag=11买多，10平多，21卖空，20平空
def exe_order(context,symbol,order_num,price,order_type,order_flag):

    rv = 0  #下单结果
    if(order_flag==11):
        rv = context.buy_open(symbol, order_num, price, order_type=order_type)
        msg =  "开多  for " + symbol + " "+ str(price) + " rv:" + str(rv)
    elif(order_flag==10):
        rv = context.sell_close(symbol, order_num, price, order_type=order_type)
        msg =  "平多  for " + symbol + " "+ str(price) + " rv:" + str(rv)
    elif(order_flag==21):
        rv = context.sell_open(symbol, order_num, price, order_type=order_type)
        msg =  "开空  for " + symbol + " "+ str(price) + " rv:" + str(rv)
    elif(order_flag==20):
        rv = context.buy_close(symbol, order_num, price, order_type=order_type)
        msg =  "平空  for " + symbol + " "+ str(price) + " rv:" + str(rv)
    else:
        msg =  "下单标识错误，请遵照order_flag=11买多，10平多，21卖空，20平空！"
    
    context.write_log(msg, stdout=PRINT)
    #输出错误信息
    if(rv!=0):
        context.write_log(context.get_error_msg(rv),stdout=PRINT)
    return rv



#初始化函数，只执行一次
def initialize(context):
    """初始化"""
    context.write_log("initialize", stdout=PRINT)  
    start_date = context.get_conf_param("start_date")#从传入参数中获取需要交易的合约
    end_date = context.get_conf_param("end_date")#下单手数
    instruments = context.get_conf_param("instruments")#下单手数
    #订阅交易合约
    context.subscribe(instruments)
    context.hold_days = 22
    context.nowdays = -1 #第一天要执行
    outer_data = context.options['data'].read_pickle()
    context.winners = outer_data['winners']
    context.losers = outer_data['losers']
    # 设置杠杆率
    context.strategy_leverage = 3
    #主力合约表
    context.dom = DataSource('dominant_CN_FUTURE').read(start_date=start_date,end_date=end_date)
    context.set_enable_auto_planed_order(1)
    
def before_trading(context, data):
    """盘前处理"""
    context.write_log("\nbefore_trading:=================================={}".format(data.current_dt), stdout=PRINT)
    return
    
def handle_data(context, data):
    """Bar行情推送"""
    context.nowdays += 1 
    today = data.current_dt.strftime('%Y-%m-%d') # 当前交易日期
    #判断当前持仓是否是当前主力合约，如果不是全部平掉，在当前主力合约上交易===================
    positions = context.get_account_positions()
    if {}!=positions:
        symbolkeys = list(positions.keys())
        for symbol in symbolkeys:
            future_variety = ''.join(re.findall(r'[A-Za-z]', symbol.split('.')[0]))
            #获取主力合约
            main_symbol = context.get_dominant(future_variety)
            if(main_symbol != symbol):#不是当前主力合约说明换合约了
                context.write_log(">>>>>>>>>>>当前日期>>>>>>{} 老合约{} 新合约{}".format(today,symbol,main_symbol), stdout=PRINT)
                # 分别获取多头持仓，和空头持仓
                position_long = context.get_position(symbol, Direction.LONG)
                position_short = context.get_position(symbol, Direction.SHORT)
                old_price = data.current(symbol, 'close') 
                main_price = data.current(main_symbol, 'close') 
                #平旧多开新多
                if(position_long.avail_qty)!=0:
                    context.write_log(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>主力切换平旧多{} 开新多 {}".format(symbol,main_symbol), stdout=PRINT)
                    exe_order(context,symbol,position_long.avail_qty,old_price,OrderType.MARKET,10)#平旧多
                    #如果合约切换和换仓是同一天，且新的买入合约中包含要切换的新主力合约则买入，否则不买
                    if(context.nowdays % context.hold_days == 0):
                        future_index = future_variety+"0000."+symbol.split('.')[1] #RB0000.SHF
                        if(future_index in context.winners[today]):#如果新开仓的有此合约则开仓，否则就不用开仓
                            exe_order(context,main_symbol,position_long.avail_qty,main_price,OrderType.MARKET,11)#开新多
                    else:
                        exe_order(context,main_symbol,position_long.avail_qty,main_price,OrderType.MARKET,11)#开新多
                #平旧空开新空
                if(position_short.avail_qty)!=0: 
                    context.write_log(">>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>>主力切换平旧空{} 开新空 {}".format(symbol,main_symbol), stdout=PRINT)
                    exe_order(context,symbol,position_short.avail_qty,old_price,OrderType.MARKET,20) #平旧空
                    #如果合约切换和换仓是同一天，且新的买入合约中包含要切换的新主力合约则买入，否则不买
                    if(context.nowdays % context.hold_days == 0):
                        future_index = future_variety+"0000."+symbol.split('.')[1] #RB0000.SHF
                        if(future_index in context.losers[today]):#如果新开仓的有此合约则开仓，否则就不用开仓
                            exe_order(context,main_symbol,position_short.avail_qty,main_price,OrderType.MARKET,21) #开新空
                    else:
                        exe_order(context,main_symbol,position_short.avail_qty,main_price,OrderType.MARKET,21) #开新空
    #===================================================================================
    
    # 不在换仓日就return,相当于后面的代码只会一个月运行一次，买入的股票会持有1个月
    if  context.nowdays % context.hold_days != 0:
        return 
    context.write_log(">>>>>>>>>>>到调仓日 当前日期>>>>>>{}".format(today), stdout=1)
    #有仓位先平
    positions = context.get_account_positions()
    if {}!=positions:
        context.write_log(">>>>>>>>>>>调仓日平仓", stdout=PRINT)
        for symbol in positions.keys():
            # 获取当前价格
            price = data.current(symbol, "close")
            # 分别获取多头持仓，和空头持仓
            position_long = context.get_position(symbol, Direction.LONG)
            position_short = context.get_position(symbol, Direction.SHORT)
            #平多
            if(position_long.avail_qty)!=0:
                exe_order(context,symbol,position_long.avail_qty,price,OrderType.MARKET,10)

            #平空
            if(position_short.avail_qty)!=0: 
                exe_order(context,symbol,position_short.avail_qty,price,OrderType.MARKET,20)
    
    
    #获取赢者和输者
    winners = context.winners[today]
    losers = context.losers[today]
   
    # 等权重分配仓位
    weight =  1 / (len(winners)+len(losers))
    portfolio_value = context.portfolio.portfolio_value
    book_value = portfolio_value * context.strategy_leverage  * weight
    
    #计算开仓需要的参数
    def cal_lots(ins):
        #找出对应的主力合约
        future_variety = ''.join(re.findall(r'[A-Za-z]', ins.split('.')[0]))  # RB
        #获取主力合约
        trade_ins = context.get_dominant(future_variety)
        #计算开仓手数
        unit = context.get_contract(future_variety).multiplier
        price = data.current(trade_ins, 'close')
        lots = math.floor(book_value/(unit*price))
        return trade_ins,price,lots
    #做多强者
    context.write_log(">>>>>>>>>>>本周期做多的合约{}".format(winners), stdout=1)
    for ins in winners:  # ['RB0000.SHF']
        trade_ins,price,lots = cal_lots(ins)
        exe_order(context,trade_ins,lots,price,OrderType.MARKET,11) #开多
    #做空弱者
    context.write_log(">>>>>>>>>>>本周期做空的合约{}".format(losers), stdout=1)
    for ins in losers:  # ['RB0000.SHF']
        trade_ins,price,lots = cal_lots(ins)
        exe_order(context,trade_ins,lots,price,OrderType.MARKET,21) #开空
            
def handle_order(context, order):
    """委托回报推送"""
    msg = "handle_order:" + order.log_str()
    context.write_log(msg, stdout=PRINT) #输出关键日志

def handle_trade(context, trade):
    """成交回报推送"""
    msg = "handle_trade:" + trade.log_str()
    context.write_log(msg, stdout=PRINT) #输出关键日志
    # 分别获取多头持仓，和空头持仓
    position_long = context.get_position(trade.symbol, Direction.LONG)
    position_short = context.get_position(trade.symbol, Direction.SHORT)
    msg = "当前多头持仓:"+ str(position_long) + "当前空头持仓:"+ str(position_short)
    context.write_log(msg, stdout=PRINT) #输出关键日志


start_date = "2020-01-01"
end_date = "2021-07-01"

df_basic = DataSource('basic_info_CN_FUTURE').read(start_date=start_date,end_date=end_date)
df_basic = df_basic[~df_basic.instrument.str.contains("CFX|INE")] #去除股指期货，国债等
# df_basic = df_basic[~df_basic.instrument.str.contains("AP|JD|SM|SF|UR|CJ")]#只测试有夜盘的品种

#找出所有指数合约
inx_lst = df_basic[df_basic["instrument"].str.contains("0000")]["instrument"].to_list()

#往前取N个交易日计算指标
NDAYS=120
df = DataSource('all_trading_days').read(end_date=end_date)
data_start_date = df[(df.country_code=="CN")&(df.date<=start_date)].iloc[-NDAYS-10]["date"]
data_start_date = datetime.strftime(data_start_date,'%Y-%m-%d')
#指数合约日线数据
df_idx = DataSource('bar1d_CN_FUTURE').read(instruments=inx_lst,start_date=data_start_date,end_date=end_date)

#计算涨幅和持仓量均值
def calc(df):
    df['mom'] = df['close'] / df['close'].shift(NDAYS) - 1
    df['roll_open_intl'] = df['open_intl'].rolling(22).mean()
    return df 
mom_df  = df_idx.groupby('instrument').apply(calc)
#过滤掉持仓量小于***的
mom_df = mom_df[mom_df['roll_open_intl']>=100000]
inx_lst = mom_df["instrument"].unique()

#只取n个合约
symbol_nums = 5
winners = mom_df.groupby('date').apply(lambda x:x.dropna().sort_values('mom', ascending=False)['instrument'][:symbol_nums].tolist())
losers = mom_df.groupby('date').apply(lambda x:x.dropna().sort_values('mom', ascending=False)['instrument'][-1*symbol_nums:].tolist())
  
instruments = [i[0:i.find("0000")] for i in inx_lst]#RB,JM等，用于注册给框架
# print("instruments=====================",instruments)
ds = DataSource.write_pickle({'winners':winners, 'losers':losers})

#需要交易者传入的参数
strategy_setting = [
    {
        "instruments" : instruments,
        "start_date" : start_date,
        "end_date" : end_date,
    }    
]

md = M.hfbacktest.v1(start_date=start_date,
                     end_date=end_date,
                     instruments=instruments,
                     options_data=ds,
                     capital_base=50000000,
                     product_type=Product.FUTURE,
                     frequency=Frequency.DAILY,
                     initialize=initialize,
                     before_trading_start=before_trading,
                     handle_data=handle_data,
                     handle_order=handle_order,
                     handle_trade=handle_trade,
                     plot_charts=True,
                     volume_limit=0,
                     disable_cache=0,
                     strategy_setting=strategy_setting,
                     slippage_type=SlippageType.PERCENT,#滑点模式
                     slippage_value=0,#买卖双向滑点
                     m_deps=np.random.rand())