class RSIBasedUniverseSelection(QCAlgorithm):

    def Initialize(self):
        # Set Start Date, End Date, and Cash
        #-------------------------------------------------------
        self.SetTimeZone(TimeZones.NewYork)     #EDIT: Added Timezon
        self.SetStartDate(2011, 1, 1)   # Set Start Date
        self.SetEndDate(2012, 1, 1)      # Set End Date
        self.SetCash(10000)            # Set Strategy Cash
        #-------------------------------------------------------
        
        self.SetBenchmark("SPY")
        self.AddUniverse(self.CoarseSelectionFilter, self.FineSelectionFilter)
        self.UniverseSettings.Resolution = Resolution.Daily    #Needs to change to Resolution.Minute once code works, leaving Daily for now to minimize data
        self.UniverseSettings.SetDataNormalizationMode = DataNormalizationMode.SplitAdjusted
        self.UniverseSettings.FeeModel = ConstantFeeModel(0.0)
        self.UniverseSettings.Leverage = 1
        
        self.__numberOfSymbols = 100
        self.__numberOfSymbolsFine = 10
        self.indicators = {}
        self.rsi_period = 14
        self.percentagebuy = 0.1


    def CoarseSelectionFilter(self, coarse):
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)   # sort descending by daily dollar volume
        return [ x.Symbol for x in sortedByDollarVolume[:self.__numberOfSymbols] ]  # return the symbol objects of the top entries from our sorted collection
        
    def FineSelectionFilter(self, fine):  # sort the data by P/E ratio and take the top 'NumberOfSymbolsFine'
        sortedByPeRatio = sorted(fine, key=lambda x: x.OperationRatios.OperationMargin.Value, reverse=False)    # sort descending by P/E ratio
        self.universe = [ x.Symbol for x in sortedByPeRatio[:self.__numberOfSymbolsFine] ]  # take the top entries from our sorted collection
        return self.universe
        
    def OnSecuritiesChanged(self, changes):
        for security in changes.AddedSecurities:
            self.indicators[security.Symbol] = SymbolData(security.Symbol, self, self.rsi_period)
            
        # for security in changes.RemovedSecurities:
        #     if security.Invested:
        #         self.Liquidate(security.Symbol, "Universe Removed Security")
            
        #     if security in self.indicators:
        #         self.indicators.pop(security.Symbol, None)
        
        
    def OnData(self, data):
        for symbol in self.universe:
            if not data.ContainsKey(symbol):    #Tested and Valid/Necessary
                continue
            
            if data[symbol] is None:            #Tested and Valid/Necessary
                continue
            
            if not symbol in self.indicators:    #Tested and Valid/Necessary
                continue
            
            # Ensure indicators are ready to update rolling windows
            if not self.indicators[symbol].is_rsi_value_ready():
                continue
            
            self.indicators[symbol].rsi_window.Add(self.indicators[symbol].get_rsi_value())
            
            if not self.indicators[symbol].is_rsi_window_ready():
                continue
            
            if self.indicators[symbol].is_valid_buy_order():
                if self.Portfolio.MarginRemaining > 0.9 * self.percentagebuy * self.Portfolio.TotalPortfolioValue:
                    self.buyquantity = round((self.percentagebuy*self.Portfolio.TotalPortfolioValue)/data[symbol].Close)
                    self.MarketOrder(symbol, self.buyquantity)
            elif self.indicators[symbol].is_valid_sell_order():
                if self.Portfolio.MarginRemaining > 0.9 * self.percentagebuy * self.Portfolio.TotalPortfolioValue:
                    self.sellQuantity = round((self.percentagebuy*self.Portfolio.TotalPortfolioValue)/data[symbol].Close)
                    self.MarketOrder(symbol, -1*self.sellQuantity)
            elif self.Portfolio[symbol].Invested:
                if self.Portfolio[symbol].IsLong and self.indicators[symbol].is_valid_buy_order_liquidate():
                    self.Liquidate(symbol)
                if not self.Portfolio[symbol].IsLong and self.indicators[symbol].is_valid_sell_order_liquidate():
                    self.Liquidate(symbol)
                if self.Securities[symbol].Holdings.UnrealizedProfitPercent < 0.1:
                    self.Debug('Liquidate after 10% loss on trade' + str(symbol))
                    self.Liquidate(symbol)

class SymbolData(object):

    rolling_window_length = 3
    
    def __init__(self, symbol, context, rsi_period):
        
        self.symbol = symbol
        self.rsi_period = rsi_period
        self.rsi_value = context.RSI(self.symbol, rsi_period, MovingAverageType.Exponential)
        self.rsi_window = RollingWindow[float](self.rolling_window_length)
        
        # Warm up EMA indicators
        history = context.History([symbol], rsi_period + self.rolling_window_length, Resolution.Daily)
        for time, row in history.loc[symbol].iterrows():
            self.rsi_value.Update(time, row["close"])
            
            # Warm up rolling windows
            if self.rsi_value.IsReady:
                self.rsi_window.Add(self.rsi_value.Current.Value)
    
    def get_rsi_value(self):
        return self.rsi_value.Current.Value
        
        
    def is_rsi_value_ready(self):
        return self.rsi_value.IsReady
        
    def is_rsi_window_ready(self):
        return self.rsi_window.IsReady
    
    def is_valid_buy_order(self):
        return self.rsi_window[0] > 25 and self.rsi_window[1] < 25
        
    def is_valid_sell_order(self):
        return self.rsi_window[0] <70 and self.rsi_window[1] > 70
        
    def is_valid_buy_order_liquidate(self):
        return self.rsi_window[0] > 50
    
    def is_valid_sell_order_liquidate(self):
        return self.rsi_window[0] < 50