# 克隆自聚宽文章：https://www.joinquant.com/post/36687
# 标题：股指期货-简化版蜘蛛网策略
# 作者：宋兵乙

# 导入函数库
from six import BytesIO
import numpy as np
import pandas as pd
import talib as ta
from jqdata import finance

from datetime import datetime, timedelta

## 初始化函数，设定基准等等
def initialize(context):
    
    # 设定沪深300作为基准
    set_benchmark('000300.XSHG')
    # 开启动态复权模式(真实价格)
    set_option('use_real_price', True)
    # 过滤掉order系列API产生的比error级别低的log
    log.set_level('order', 'error')
    # 输出内容到日志 log.info()
    log.info('初始函数开始运行且全局只运行一次')
    set_subportfolios([SubPortfolioConfig(cash=context.portfolio.starting_cash, type='index_futures')])
    set_order_cost(OrderCost(open_commission=0.000023, close_commission=0.000023,close_today_commission=0.0004), type='index_futures')
    set_option('futures_margin_rate', 0.17)
    set_slippage(StepRelatedSlippage(4))
    
    g.close_long = []
    g.close_short = []
    g.open_long = []
    g.open_short = []
    g.times = 3

    run_daily(cal_signal, time='20:30')

    run_daily(trade, time='09:32')
    
## 开盘前运行函数
def set_parameter():
    ## 设置交易品种
    g.symbol_list = [
        'IH',
        ]
        
    ## 设置品种乘数
    g.muliti_dict = {
        ## 中金所
        'IF':300, 'TF':10000, 'IC':200, 'IH':300,
        }

## 开盘前20:30:00运行函数
def cal_signal(context):
    # 参数设置
    set_parameter()
    # 创建保持计算结果的DataFrame
    g.signal_df = pd.DataFrame(columns=['期货代码','交易信号','手数'])
    ## 获取当前持仓
    g.short_holds = context.portfolio.short_positions.keys()
    g.long_holds = context.portfolio.long_positions.keys()
    for symbol in g.symbol_list:
        ## 获取主力合约
        symbol_dominant = get_dominant_future(symbol)
        
        if len(symbol_dominant) == 0:
            continue
        
        # # 获取品种持仓状况
        if  symbol_dominant in g.short_holds:
            hold_signal = -1
        # 查看多单仓位情况    
        elif  symbol_dominant in g.long_holds:
            hold_signal = 1
        else:
            hold_signal = 0
        
        ## 现在的自然日日期
        date = context.current_dt.strftime('%Y%m%d')
    
        ## 获取主力合约价格
        current_time = context.current_dt
        price_data = get_price(symbol_dominant, end_date=current_time, frequency='1d', 
                                    fields=['close'], count = 10)
        dominant_price = price_data['close'][-1]
        # 获取所有合约
        contract_all = get_future_contracts(symbol, current_time)
        # print('contract_all: ', contract_all)

        long_hold_delta = 0
        short_hold_delta = 0
        
        for contract in contract_all:
        
            q=query(finance.FUT_MEMBER_POSITION_RANK.day,
            finance.FUT_MEMBER_POSITION_RANK.code,
            finance.FUT_MEMBER_POSITION_RANK.rank_type,
            finance.FUT_MEMBER_POSITION_RANK.rank,
            finance.FUT_MEMBER_POSITION_RANK.member_name,
            finance.FUT_MEMBER_POSITION_RANK.indicator,
            finance.FUT_MEMBER_POSITION_RANK.indicator_increase).\
            filter(finance.FUT_MEMBER_POSITION_RANK.code==contract,
                   finance.FUT_MEMBER_POSITION_RANK.day==date)
            
            df=finance.run_query(q)
            
            #数据分离，分别按多单持有量与空单持有量排序
            df1 = df[df['rank_type']=='持买单量排名']
            df2 = df[df['rank_type']=='持卖单量排名']
            
            #分别计算多、空持仓前二十会员的多仓增量与空仓增量
            long_hold_delta = long_hold_delta  + df1[:20]['indicator_increase'].sum()
            short_hold_delta = short_hold_delta + df2[:20]['indicator_increase'].sum()
            # print(contract)
            # print('long_hold_delta: ', long_hold_delta)
            # print('short_hold_delta: ', short_hold_delta)
            
        try:
            if hold_signal == 0:
                if long_hold_delta>0 and short_hold_delta<0:
                    signal = 1
                elif long_hold_delta<0 and short_hold_delta>0:
                    signal = -1
                else: signal = 0
                    
            if hold_signal > 0:
                if long_hold_delta <= 0 or short_hold_delta >= 0:
                    signal = 0
                else: 
                    signal = 1
              
            if hold_signal < 0:
                if long_hold_delta >= 0 or short_hold_delta <= 0:
                    signal = 0
                else: 
                    signal = -1   
        except:
            log.info('信号计算报错'+str(symbol_dominant))
            
        # 按照两倍杠杆计算持仓手数
        total_value = context.portfolio.total_value
        multi = g.muliti_dict[symbol]
        amount = int(total_value / dominant_price / multi /len(g.symbol_list) * g.times)
        print('total_value: ', total_value)
        print('dominant_price: ', dominant_price)
        print('multi: ', multi)
        print('g.symbol_list: ', len(g.symbol_list))
        print('amount: ', amount)
        # 开仓数量过滤
        if amount < 1:
            amount = 1
            
        g.signal_df = g.signal_df.append({'期货代码': symbol_dominant, 
                                '交易信号': signal,
                                '手数': amount,
                                },
                              ignore_index=True)
                          
    log.info('\n' + str(g.signal_df))
    signal_df_analysis(context)
    
## 夜盘交易函数
def trade(context):

    # 平仓
    for future in g.close_long:
        order_target(future, 0, side='long')
    for future in g.close_short:
        order_target(future, 0, side='short')

    # 开仓
    for future in g.open_long:
        amount = g.signal_df[g.signal_df['期货代码'] == future]['手数'].values[0]
        order_target(future, amount, side='long')
    for future in g.open_short:
        amount = g.signal_df[g.signal_df['期货代码'] == future]['手数'].values[0]
        order_target(future, amount, side='short') 
    
    
def signal_df_analysis(context):
    #交易信号
    target_long = g.signal_df[g.signal_df['交易信号'] > 0]
    target_short = g.signal_df[g.signal_df['交易信号'] < 0]
    
    target_long_list = list(target_long['期货代码'])  # 想持多仓的目标
    target_short_list = list(target_short['期货代码'])  # 想持空仓的目标
    
    # 下一个交易日需要处理多类的品种
    g.open_long = [i for i in target_long_list if i not in g.long_holds]
    g.close_long = [i for i in g.long_holds if i not in target_long_list]
    # 下一个交易日需要处理空类的品种
    g.open_short = [i for i in target_short_list if i not in g.short_holds]
    g.close_short = [i for i in g.short_holds if i not in target_short_list]