# coding: utf-8

# In[ ]:

# 导入函数库
import jqdata
import pandas as pd
import numpy as np

# 初始化函数，设定基准等等
def initialize(context):
    # 设定沪深300作为基准
    set_benchmark('000300.XSHG')
    # 开启动态复权模式(真实价格)
    set_option('use_real_price', True)
    # 股票类每笔交易时的手续费是：买入时佣金万分之三，卖出时佣金万分之三加千分之一印花税, 每笔交易佣金最低扣5块钱
    set_order_cost(OrderCost(close_tax=0.001, open_commission=0.0003, close_commission=0.0003, min_commission=5), type='stock')

#获取沪深300股票池
stock_set=get_index_stocks('000300.XSHG')
#此处可增加选股条件
q = query(
    valuation.code, # 股票代码
).filter(
    valuation.code.in_(stock_set),#只对设定股票池执行
)
fdf = get_fundamentals(q)
#取前50只股
fdf=fdf.head(50)
#获取股票列表
stock_list=list(fdf['code'])
trend={}
for security in stock_list:
    trend[security]='None'
    
def handle_data(context, data):
    N = 20   # 计算TR时的N
    M = 49   # 计算MATRIX时的M
    num=3
    length_of_data = N+M+num  # 取closeprice的天数，为了足够计算MATRIX、TRIX
    
    cash=context.portfolio.available_cash/50
    for security in stock_list:
        close_price=attribute_history(security,length_of_data,'1d',('close'),skip_paused=True)
        where_are_nan = np.isnan(close_price)  
        where_are_inf = np.isinf(close_price)  
        close_price[where_are_nan] = 0  
        close_price[where_are_inf] = 0
        
        MA5=close_price['close'][-5:].mean()
        MA10=close_price['close'][-10:].mean()
        ma5=close_price['close'][-6:-1].mean()
        ma10=close_price['close'][-11:-1].mean()
        
        price_array={}
        for i in range(0,M+num):
            price_array[i]=close_price['close'][i:i+N]
        #TR=收盘价的N日指数移动平均
        TR={}
        for i in range(0,M+num):
            TR[i]=np.mean(price_array[i])  
        #TRIX=(TR-昨日TR)/昨日TR*100
        TRIX={}
        for i in range(1,M+num):
            if TR[i]==0:
                TRIX[i]=0
                continue
            TRIX[i]=(TR[i]-TR[i-1])/TR[i]*100  
        #MATRIX=TRIX的M日简单移动平均
        MATRIX={}
        for i in range(0,num):
            TRIX_sum=0
            for j in range(1,M):
                TRIX_sum=TRIX_sum+TRIX[i+j]
            MATRIX[i]=TRIX_sum/M
        current_price=data[security].price
        
        length=0
        for i in range(0,num-1):
            if TRIX[M+i]>MATRIX[i]:
                length=length+1
            else:
                length=length-1
        if length>0 and MATRIX[num-1]>MATRIX[0]:
            trend[security]='up'
        elif length<0 and MATRIX[num-1]<MATRIX[0]:
            trend[security]='down'
        
        if trend[security]=='up':
            order_value(security,cash)
        elif trend[security]=='down' and security in context.portfolio.positions:
            if MA5<MA10:
                order_target(security,0)
            elif MA5>MA10 and ma5<ma10 :
                order_value(security,cash)