from QuantConnect.Indicators import RelativeStrengthIndex, AverageTrueRange

'''
Universe: SPY and IEF
Timeframe: Daily (the only reason why it is on minute is because we need the OnOrderEvent)
Position size: 50%
Buy rules: After the market closes, buy on Market-On-Open order if the 3-day cumulative RSI(2) < 15.
    Use a stoploss with 2*ATR(1) below the open price (which is the same as fill price)
Sell rules: After the market closes, sell if RSI(2) < 70 using MOO order.

Needing almost 80 lines of code for this simple strategy seems a bit too much. Can the code be made more efficient/smaller?
Also: is there an easy way to 'attach' a stop order to a market order such that when the position gets closed,
    the stop order is automatically cancelled?
'''
class RSI_Strategy(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 1, 1) 
        self.SetCash(100000)
        
        tickers = ['SPY', 'IEF']
        self.symbol_data_by_symbol = {}
        for ticker in tickers:
            symbol = self.AddEquity(ticker, Resolution.Daily).Symbol
            self.symbol_data_by_symbol[symbol] = SymbolData(self, symbol)

    def OnData(self, data):
        for symbol, symbol_data in self.symbol_data_by_symbol.items():
            if not data.ContainsKey(symbol):
                continue
        
            if not self.Securities[symbol].Invested:
                if sum(list(symbol_data.rsi_window)) < 15:
                    quantity = self.CalculateOrderQuantity(symbol, 1 / len(self.symbol_data_by_symbol))
                    if quantity:
                        self.MarketOrder(symbol, quantity)
            
            elif symbol_data.rsi.Current.Value > 70:
                    self.Liquidate(symbol)

        
    #Attach a stop order to the market order. The reason why this cannot be done in the OnData code is 
    #that we cannot know the fill price before the order gets filled.
    def OnOrderEvent(self, orderEvent):
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        #If a market on open order gets filled
        if orderEvent.Status == OrderStatus.Filled and orderEvent.FillQuantity > 0: 
            fillPrice = orderEvent.FillPrice
            
            #Set SL order
            symbol_data = self.symbol_data_by_symbol[orderEvent.Symbol]
            stop_price = fillPrice - 2 * symbol_data.atr.Current.Value
            symbol_data.stopTicket = self.StopMarketOrder(orderEvent.Symbol, -symbol_data.amount, stop_price)
                
class SymbolData:
    amount = 0
    stopTicket = None
    
    def __init__(self, algorithm, symbol, rsi_indicator_length=2, rsi_window_length=3):
        
        # Create indicators
        self.atr = AverageTrueRange(1, MovingAverageType.Simple)
        self.rsi = RelativeStrengthIndex(rsi_indicator_length, MovingAverageType.Wilders)
        self.rsi_window = RollingWindow[float](rsi_window_length)
        
        # Warmup indicators
        history = algorithm.History(symbol, rsi_indicator_length + rsi_window_length, Resolution.Daily)
        for time, row in history.loc[symbol].iterrows():
            trade_bar = TradeBar(time, symbol, row.open, row.high, row.low, row.close, row.volume)
            self.atr.Update(trade_bar)
            self.rsi.Update(time, row.close)
            
            if self.rsi.IsReady:
                self.rsi_window.Add(self.rsi.Current.Value)
        
        # Setup consolidators to update indicators
        self.consolidator = TradeBarConsolidator(timedelta(1))
        self.consolidator.DataConsolidated += self.CustomHandler
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator)

    def CustomHandler(self, sender, consolidated):
        self.atr.Update(consolidated)
        self.rsi.Update(consolidated.Time, consolidated.Close)
        self.rsi_window.Add(self.rsi.Current.Value)