class SPY_SMA(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 4, 25)  # 设置回测起始日期
        self.SetEndDate(2021, 4, 27) # 设置回测结束日期
        self.SetCash(100000000)  # 设置现金100000000
        
        self.spy = self.AddEquity("SPY" , Resolution.Minute , Market.USA , True, 1, True).Symbol #添加SPY数据
        
        fiveCountConsolidator = TradeBarConsolidator(TimeSpan.FromMinutes(5))  # 创造一个5分钟tradebar数据整合器            
        fiveCountConsolidator.DataConsolidated += self.FiveBarHandler #给5分钟tradebar数据整合器添加事件
        self.SubscriptionManager.AddConsolidator(self.spy, fiveCountConsolidator) #注册整合器来自动更新5分钟tradebar数据
        self.sma_5min_window = RollingWindow[float](2) #创建存储2个sma5数据的移动窗口
        self.sma_30min_window = RollingWindow[float](2) #创建存储2个sma30数据的移动窗口
        self.close_window = RollingWindow[float](2) #创建存储2个5分钟k线收盘价数据的移动窗口
        
        self.sma_5min_5 = SimpleMovingAverage(5)    # 计算基于5min 的MA5  
        self.sma_5min_30 = SimpleMovingAverage(30)  # 计算基于30min 的MA30 
        self.RegisterIndicator(self.spy, self.sma_5min_5, fiveCountConsolidator)  # 注册SPY的每日数据来自动更新sma5指标
        self.RegisterIndicator(self.spy, self.sma_5min_30, fiveCountConsolidator)  #注册SPY的每日数据来自动更新sma30指标
        
        self.quantity = 0 #设置初始股票持有量为0
        self.stop_trade = False #设置停止交易标志
        self.current_day = 0  #记录日期
        self.day_loss = 0 #设置初始每日亏损金额为0
        self.day_values = 0 #每天起始资产金额
        self.SetWarmUp(30) #  为所有订阅的数据预热30笔数据

    def OnData(self, data):
        pass
        
    def FiveBarHandler(self, sender, bar):
        if bar.Time.day != self.current_day:
            self.day_values = self.Portfolio.Cash + self.Portfolio.UnsettledCash
            self.current_day = bar.Time.day
            self.stop_trade = False
            
        if bar.Time.day == self.current_day and not self.stop_trade:
            self.close_window.Add(bar.Close) #获取5分钟k线收盘价数值，将其添加到移动窗口中
            self.sma_5min_window.Add(self.sma_5min_5.Current.Value) #获取sma5的数值，将其添加到移动窗口中
            self.sma_30min_window.Add(self.sma_5min_30.Current.Value) #获取sma30的数值，将其添加到移动窗口中
            if self.sma_5min_window.IsReady and self.sma_30min_window.IsReady: #
                if (self.sma_5min_window[0] > self.sma_30min_window[0]) \
                    and (self.close_window[0] > self.sma_5min_window[0]) and self.quantity <= 0:
                    self.MarketOrder(self.spy, -self.quantity + 100) #满足上述条件，提交市价订单，持多仓
                    self.quantity = 100 #更新股票持有量为100
                elif (self.sma_5min_window[0] < self.sma_30min_window[0]) \
                    and (self.close_window[0] < self.sma_5min_window[0]) and self.quantity > 0:
                    self.MarketOrder(self.spy, -self.quantity - 100) #满足上述条件，提交市价订单，持空仓
                    self.quantity = -100 #更新股票持有量为-100
        
            _current_values = self.Portfolio.Cash + self.Portfolio.UnsettledCash
            if _current_values < self.day_values:
                loss = self.day_values - _current_values #计算当日损失
                if (loss/self.day_values) >= 0.05 and not self.stop_trade: #当日亏损超过5%,清算资产
                    self.Liquidate()
                    self.stop_trade = True #设置当日策略暂停标志