from itertools import groupby

class UniverseRollingAlgorithm(QCAlgorithm):
   
   def Initialize(self):       #Initialize Dates, Cash, Equities, Fees, Allocation, Parameters, Indicators, Charts
       
       # Set Start Date, End Date, and Cash
       #-------------------------------------------------------
       self.SetTimeZone(TimeZones.NewYork)     #EDIT: Added Timezon
       self.SetStartDate(2021, 1, 1)   # Set Start Date
      # self.SetEndDate(2020, 4, 15)      # Set End Date
       self.SetCash(100000)            # Set Strategy Cash
       #-------------------------------------------------------
       
       # Set Custom Universe
       #-------------------------------------------------------
       self.AddUniverse(self.CoarseSelectionFilter, self.FineSelectionFilter)
       self.UniverseSettings.Resolution = Resolution.Minute    #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.universe = []
       #-------------------------------------------------------
       
     # self.SetBrokerageModel(BrokerageName.Alpaca, AccountType.Cash)      #EDIT: Added Brokerage, appears to have set fees to zero
       
       self.EMA_Period_Fast = 20
       self.EMA_Period_Slow = 50
       self.__numberOfSymbols     = 100
       self.__numberOfSymbolsFine = 10
       
       self.indicators = {}
       
       # Define Percentage Allocation
       #-------------------------------------------------------
       self.percentagebuy = 0.1
       #-------------------------------------------------------
       
       
       # Sector Selectior
       self.numberOfSymbolsCoarse = 500
       self.exposureToSector = 2
       self.lastMonth = -1
       self.SetSecurityInitializer(lambda x: x.SetMarketPrice(self.GetLastKnownPrice(x)))
       
           
   def CoarseSelectionFilter(self, coarse):
       
       if self.Time.month == self.lastMonth:
           return Universe.Unchanged
       self.lastMonth = self.Time.month
       
       allCoarse = [x for x in coarse if x.HasFundamentalData and x.Price > 1 and x.Volume > 1]
       finalCoarse = sorted(allCoarse, key = lambda x: x.DollarVolume, reverse = True)
       
       return [x.Symbol for x in finalCoarse][:self.numberOfSymbolsCoarse]
       
       #Old
       #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'
       filteredSymbols = []
       sortedBySector = [x for x in fine]
      # self.Log("ORDER NOTIFICATION >> {}".format(str(AssetClassification.MorningstarSectorCode)))
       for code, g in groupby(sortedBySector, lambda x: x.AssetClassification.MorningstarSectorCode):
           for x in sorted(g, key = lambda x: x.ValuationRatios.PERatio, reverse = True)[:self.exposureToSector]:
               filteredSymbols.append(x.Symbol)
               
       self.universe = filteredSymbols[:5]       
           
       return self.universe
       
       
       #Old
       #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):
       # Create indicator for each new security
       for security in changes.AddedSecurities:
           self.indicators[security.Symbol] = SymbolData(security.Symbol, self, self.EMA_Period_Fast, self.EMA_Period_Slow)
       
       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):         #Entry Point for Data and algorithm - Check Data, Define Buy Quantity, Process Volume, Check Portfolio, Check RSI, Execute Buy/Sell orders, Chart Plots
       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].slow_ema.IsReady:
               continue
           
           # Update EMA rolling windows
           self.indicators[symbol].fast_ema_window.Add(self.indicators[symbol].get_fast_EMA())
           self.indicators[symbol].slow_ema_window.Add(self.indicators[symbol].get_slow_EMA())
           
           # Check for Indicator Readiness within Rolling Window
           #-------------------------------------------------------
           if not (self.indicators[symbol].fast_ema_window.IsReady and self.indicators[symbol].slow_ema_window.IsReady): 
               continue    #return     #EDIT 
           
           #EXECUTE TRADING LOGIC HERE - 
           if self.Portfolio[symbol].Invested:
               # Sell condition
               if (self.indicators[symbol].fast_ema_window[1] >= self.indicators[symbol].slow_ema_window[1]) and (self.indicators[symbol].fast_ema_window[4] < self.indicators[symbol].slow_ema_window[4]):
                   self.Liquidate(symbol)
                   
               # Buy conditions
           elif self.Portfolio.MarginRemaining > 0.9 * self.percentagebuy * self.Portfolio.TotalPortfolioValue:
               if self.indicators[symbol].fast_ema_window[1] <= self.indicators[symbol].slow_ema_window[1] and \
                   (self.indicators[symbol].fast_ema_window[4] > self.indicators[symbol].slow_ema_window[4]):
                       
                   self.buyquantity = round((self.percentagebuy*self.Portfolio.TotalPortfolioValue)/data[symbol].Close)
                   self.MarketOrder(symbol, self.buyquantity)
           
class SymbolData(object):
   
   rolling_window_length = 5
   
   def __init__(self, symbol, context, fast_ema_period, slow_ema_period):
       
       self.symbol = symbol
       self.fast_ema_period = fast_ema_period
       self.slow_ema_period = slow_ema_period
       self.fast_ema = context.EMA(symbol, self.fast_ema_period, Resolution.Minute)    #, fillDataForward = True, leverage = 1, extendedMarketHours = False)
       self.slow_ema = context.EMA(symbol, self.slow_ema_period, Resolution.Minute)    #, fillDataForward = True, leverage = 1, extendedMarketHours = False)
       self.fast_ema_window = RollingWindow[float](self.rolling_window_length)
       self.slow_ema_window = RollingWindow[float](self.rolling_window_length)
       
       # Warm up EMA indicators
       history = context.History([symbol], slow_ema_period + self.rolling_window_length, Resolution.Minute)
       for time, row in history.loc[symbol].iterrows():
           self.fast_ema.Update(time, row["close"])
           self.slow_ema.Update(time, row["close"])
           
           # Warm up rolling windows
           if self.fast_ema.IsReady:
               self.fast_ema_window.Add(self.fast_ema.Current.Value)
           if self.slow_ema.IsReady:
               self.slow_ema_window.Add(self.slow_ema.Current.Value)
   
   def get_fast_EMA(self):
       return self.fast_ema.Current.Value
       
   def get_slow_EMA(self):
       return self.slow_ema.Current.Value