### <summary>
### Simple RSI Strategy intended to provide a minimal algorithm example using
### one indicator with the most basic plotting
### </summary>
from datetime import timedelta
class RSIAlgorithm(QCAlgorithm):
    # 1 - Add the FANG stocks (Facebook, Amazon, , Netflix, Google)
    # 2 - Cycle through stocks
    # 3 - Cycle through list adding each equity
    # 3 - Create an indicator dict like backtrader
    def Initialize(self):
        '''Initialise the data and resolution required, as well as the cash and start-end dates for your algorithm. All algorithms must initialized.'''
        # Set our main strategy parameters
        self.SetStartDate(2020,8,1)    # Set Start Date
        #self.SetEndDate(2020,8,6)      # Set End Date
        self.SetCash(400)            # Set Strategy Cash
        self.SetWarmUp(timedelta(30))
        RSI_Period    = 14                # RSI Look back period
        self.RSI_OB   = 75                # RSI Overbought level
        self.RSI_OS   = 50                # RSI Oversold level
        self.Allocate = 1              # Percentage of captital to allocate
        self.Equities = ["TQQQ"]
        self.Indicators = dict()
        self.Charts = dict()
        self.Consolidators = dict()
        self.slow = dict()
        self.fast = dict()
        # Find more symbols here: http://quantconnect.com/data
        
        for Symbol in self.Equities:
            self.AddSecurity(SecurityType.Equity, Symbol, Resolution.Minute)
            self.Consolidators[Symbol] = dict()
            self.AddEquity(Symbol, Resolution.Daily)
            # Each Equity requires its own consoilidator! See:
            # https://www.quantconnect.com/forum/discussion/1936/multiple-consolidators/p1
            # https://www.quantconnect.com/forum/discussion/1587/multiple-symbol-indicator-values-in-consolidated-bar-handler/p1
            # ------------------------
            # Create our consolidators
            self.Consolidators[Symbol]['W1 Con'] = TradeBarConsolidator(timedelta(days=5))
            # Register our Handlers
            self.Consolidators[Symbol]['W1 Con'].DataConsolidated += self.On_W1
            self.Indicators[Symbol] = dict()
            self.Indicators[Symbol]['RSI'] = dict()
            self.Indicators[Symbol]['RSI']['D'] = self.RSI(Symbol, RSI_Period)
            self.Indicators[Symbol]['RSI']['W'] = RelativeStrengthIndex(Symbol, RSI_Period)
            # Register the indicaors with our stock and consolidator
            self.RegisterIndicator(Symbol, self.Indicators[Symbol]['RSI']['W'], self.Consolidators[Symbol]['W1 Con'])
            # Finally add our consolidators to the subscription
            # manager in order to receive updates from the engine
            self.SubscriptionManager.AddConsolidator(Symbol, self.Consolidators[Symbol]['W1 Con'])
            self.Charts[Symbol] = dict()
            # Plot the RSI
            RSIChartName = Symbol+" RSI"
            self.Charts[Symbol]['RSI'] = Chart(RSIChartName, ChartType.Stacked)
            self.Charts[Symbol]['RSI'].AddSeries(Series("D1", SeriesType.Line))
            self.Charts[Symbol]['RSI'].AddSeries(Series("W1", SeriesType.Line))
            self.AddChart(self.Charts[Symbol]['RSI'])
            # Create a custom volume chart
            VolChartName = Symbol+" Volume"
            self.Charts[Symbol]['VOL'] = Chart(VolChartName, ChartType.Stacked)
            self.Charts[Symbol]['VOL'].AddSeries(Series('Buying Volume', SeriesType.Bar))
            self.Charts[Symbol]['VOL'].AddSeries(Series('Selling Volume', SeriesType.Bar))
            self.AddChart(self.Charts[Symbol]['VOL'])
            EMAChartName = Symbol+" EMA"
            self.Charts[Symbol]['EMA'] = Chart(EMAChartName, ChartType.Stacked)
            self.Charts[Symbol]['EMA'].AddSeries(Series('FAST EMA', SeriesType.Line))
            self.Charts[Symbol]['EMA'].AddSeries(Series('SLOW EMA', SeriesType.Line))
            self.AddChart(self.Charts[Symbol]['EMA'])
            self.fast[Symbol] = self.EMA(Symbol, 5, Resolution.Daily);
            # create a 30 day exponential moving average
            self.slow[Symbol] = self.EMA(Symbol, 10, Resolution.Daily);
        # Ensure that the Indicator has enough data before trading,.
        # x5 to get enough weekly data
        self.SetWarmUp(RSI_Period*5)
    
    def On_W1(self,sender,bar):
        '''
        This method will be called every time a new 30 minute bar is ready.
        '''
        # Make sure we are not warming up
        if self.IsWarmingUp: return
        Symbol = str(bar.get_Symbol())
        self.Plot(Symbol+' RSI', 'W1', self.Indicators[Symbol]['RSI']['W'].Current.Value)
    
    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
        Arguments:
            data: Slice object keyed by symbol containing the stock data
        '''
        # Make sure we are not warming up
        if self.IsWarmingUp: return
        # Loop through our equities
        for Symbol in self.Equities:
            if not self.fast[Symbol].IsReady:
                return
            # wait for our slow ema to fully initialize
            if not self.slow[Symbol].IsReady:
                return
            # Add some alias for reading ease
            Close = data[Symbol].Close
            Volume = data[Symbol].Volume
            D1_RSI = self.Indicators[Symbol]['RSI']['D'].Current.Value
            W1_RSI = self.Indicators[Symbol]['RSI']['W'].Current.Value
            #self.Debug("{}: Close: {} RSI: {}".format(Symbol, Close, D1_RSI))
            if data[Symbol].Close >= data[Symbol].Open:
                self.Plot(Symbol+" Volume", 'Buying Volume', Volume)
            else:
                self.Plot(Symbol+" Volume", 'Selling Volume', Volume)
            self.Plot(Symbol +' RSI', 'D1', D1_RSI)
            self.Plot(Symbol+" EMA", 'Fast EMA', self.fast[Symbol].Current.Value)
            self.Plot(Symbol+" EMA", 'Slow EMA', self.slow[Symbol].Current.Value)
            # Determine our entry and exit conditions
            # Do it here to avoid long lines later
            Long_Cond1 = D1_RSI < self.RSI_OS
            Long_Cond2 = W1_RSI < self.RSI_OS
            Long_Cond3 = self.fast[Symbol].Current.Value > self.slow[Symbol].Current.Value
            Exit_Cond1 = D1_RSI > self.RSI_OB
            Exit_Cond2 = W1_RSI > self.RSI_OB
            Exit_Cond3 = self.fast[Symbol].Current.Value < self.slow[Symbol].Current.Value
            if not self.Securities[Symbol].Invested:
                # If not, the long conditions
                if all([Long_Cond1, Long_Cond3]):
                    # Buy!
                    self.SetHoldings(Symbol, self.Allocate)
                    self.Debug("{}: Buy: {} RSI: {}".format(Symbol, Close, D1_RSI))
            else:
                if all([Exit_Cond1, Exit_Cond3]):
                    # Sell!
                    #self.Liquidate(Symbol)
                    self.Debug("{}: Sell: {} RSI: {}".format(Symbol, Close, D1_RSI))