# 可以自己import我们平台支持的第三方python模块，比如pandas、numpy等。
import talib
import numpy as np
import math
import pandas
# 在这个方法中编写任何的初始化逻辑。context对象将会在你的算法策略的任何方法之间做传递。
def init(context):

    #选择我们感兴趣的股票
    context.s1 = "000001.XSHE"
    context.s2 = "601988.XSHG"
    context.s3 = "000068.XSHE"
    context.stocks = [context.s1,context.s2,context.s3]

    context.TIME_PERIOD = 14
    context.HIGH_RSI = 85
    context.LOW_RSI = 30
    context.ORDER_PERCENT = 0.3

# 你选择的证券的数据更新将会触发此段逻辑，例如日或分钟历史数据切片或者是实时数据切片更新
def handle_bar(context, bar_dict):
    # 开始编写你的主要的算法逻辑

    # bar_dict[order_book_id] 可以拿到某个证券的bar信息
    # context.portfolio 可以拿到现在的投资组合状态信息

    # 使用order_shares(id_or_ins, amount)方法进行落单

    # TODO: 开始编写你的算法吧！

    # 对我们选中的股票集合进行loop，运算每一只股票的RSI数值
   for stock in context.stocks:
        # 读取历史数据
        prices = history(context.TIME_PERIOD+1,'1d','close')[stock].values

        # 用Talib计算RSI值
        rsi_data = talib.RSI(prices,timeperiod=context.TIME_PERIOD)[-1]

        curPosition = context.portfolio.positions[stock].quantity
        #用剩余现金的30%来购买新的股票
        target_available_cash = context.portfolio.cash * context.ORDER_PERCENT
        print("target_available_cash", stock, target_available_cash)

        #当RSI大于设置的上限阀值，清仓该股票
        if rsi_data > context.HIGH_RSI and curPosition > 0:
            print("order_target_value(%s,0)" % stock)
            order_target_value(stock,0)

        #当RSI小于设置的下限阀值，用剩余cash的一定比例补仓该股
        if rsi_data < context.LOW_RSI:
            logger.info("target available cash caled: " + str(target_available_cash))
            #如果剩余的现金不够一手 - 100shares，那么会被ricequant 的order management system reject掉
            order_value(stock, target_available_cash)
            print("order_value(%s, %s)" % (stock, target_available_cash))
