from jqdata import *

def initialize (context):
    """初始化函数"""
    
    # 交易周期计数起始值
    g.days=0
    # 调仓周期为
    g.refresh_rate=10 
    # 持仓数量
    g.stocknum=5
    
    # 设置参考基准为上证500
    set_benchmark('000905.XSHG') 
    # 使用真实价格交易
    set_option('use_real_price', True) 
    # 设置交易日志级别为error
    log.set_level('order', 'error')
    # 为全部交易品种设定固定值滑点
    set_slippage(FixedSlippage(0.02))
    
    # 指定周期交易函数
    run_daily(trade, 'every_bar')


def order_cost(dt):
    """交易费用设置"""
    
    # 2013年之后的交易费用
    if dt>datetime.datetime(2013,1, 1):
        set_order_cost(OrderCost(
            open_tax=0, 
            close_tax=0.001, 
            open_commission=0.0003, 
            close_commission=0.0003, 
            close_today_commission=0, 
            min_commission=5), 
            type='stock')
    # 2013年之前的交易费用
    else:
        set_order_cost(OrderCost(
            open_tax=0, 
            open_commission=0.003,
            close_commission=0.003, 
            close_tax=0.001,
            min_commission=5), 
            type='stock')
        
    
def trade(context):
    """交易函数"""
    
    '''
    +++++++++++++++++++++++++++++++++++++++++++++
    一、判断今天是否需要交易
    +++++++++++++++++++++++++++++++++++++++++++++
    '''
    g.days += 1
    if g.days % g.refresh_rate != 1:
        return
    
    '''
    +++++++++++++++++++++++++++++++++++++++++++++
    二、筛选股票
    +++++++++++++++++++++++++++++++++++++++++++++
    '''
    # 设定当前的交易费用
    order_cost(context.current_dt)
    
    # 从基本面筛选股票
    stock_to_choose = get_fundamentals(query(
        valuation.code, 
        valuation.pe_ratio, 
        valuation.market_cap, 
        indicator.eps, 
        indicator.inc_return, 
        indicator.inc_net_profit_annual
        ).filter(
            valuation.pe_ratio < 45,
            valuation.pe_ratio > 20,
            indicator.eps > 0.3,
            indicator.inc_net_profit_annual > 0.30,
            indicator.roe > 15
            ).order_by(
                indicator.roe.desc()
                ).limit(
                    50), date=None)

    # 剔除停牌股票
    stockset=list(stock_to_choose['code'])
    current_data = get_current_data()
    stockset = [stock for stock in stockset if not current_data[stock].paused]

    '''
    +++++++++++++++++++++++++++++++++++++++++++++
    三、股票买卖
    +++++++++++++++++++++++++++++++++++++++++++++
    '''
    # 获取持仓列表，不在买入列表中的持仓股票将被卖出
    sell_list = list(context.portfolio.positions.keys())
    for stock in sell_list:
        if stock not in stockset[:g.stocknum]:
            order_target_value(stock, 0) 

    # 分配资金
    if len(context.portfolio.positions) < g.stocknum :
        Num = g.stocknum - len(context.portfolio.positions)
        Cash = context.portfolio.cash/Num
    else: 
        return

    # 买入股票
    for stock in (set(stockset[:g.stocknum])-set(sell_list)):
        order_target_value(stock, Cash)