class UniverseSelection(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 1)  # Set Start Date
        #self.SetEndDate(2021, 1, 5) # Set End Date
        self.SetCash(100000)  # Set Strategy Cash
        
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        self.SPY = self.AddEquity('SPY', Resolution.Minute).Symbol
        
        self.day = -1
        self.num_coarse = 20
        self.min_days_after_earnings = 10
        self.max_days_after_earnings = 80
        self.ema_period = 40 # 8 on 5 min timeframe
        self.sma_period = 275 # 55 on 5 min timeframe
        self.bb_period = 600 # 20 on 30 min timeframe
        self.bb_k = 2
        self.gap_distance = 0.02 # 2%
        self.data = {}
        self.selected = []
        
        
        self.Schedule.On(self.DateRules.EveryDay(),
                        self.TimeRules.AfterMarketOpen(self.SPY, 30),
                        self.Selection)
            
        self.Schedule.On(self.DateRules.EveryDay(),
                        self.TimeRules.AfterMarketOpen(self.SPY, 1),
                        self.AtMarketOpen)
        
        self.Schedule.On(self.DateRules.EveryDay(),
                        self.TimeRules.BeforeMarketClose(self.SPY, 1),
                        self.BeforeMarketCloses)
                        
    def CoarseSelectionFunction(self, coarse):
        
        if self.day == self.Time.day:
            return Universe.Unchanged

        self.day = self.Time.day

        # drop stocks which have no fundamental data or have too low prices
        selected = [x for x in coarse if (x.HasFundamentalData) and (float(x.Price) > 5)]
        # rank the stocks by dollar volume 
        filtered = sorted(selected, key=lambda x: x.DollarVolume, reverse=True)
        return [ x.Symbol for x in filtered[:self.num_coarse]]
    
    def FineSelectionFunction(self, fine):
        
        filtered = [x for x in fine if x.SecurityReference.IsPrimaryShare
                    and x.SecurityReference.SecurityType == "ST00000001"
                    and x.SecurityReference.IsDepositaryReceipt == 0
                    and x.CompanyReference.IsLimitedPartnership == 0
                    and x.EarningReports.FileDate < self.Time - timedelta(days=self.min_days_after_earnings)
                    and x.EarningReports.FileDate > self.Time - timedelta(days=self.max_days_after_earnings)]

        return [x.Symbol for x in filtered]
        
        
    def Selection(self):
        
        for symbol in self.data.keys():
            if self.data[symbol].GapUp and self.data[symbol].EMA > self.data[symbol].SMA and self.data[symbol].BB.UpperBand.Current.Value < self.Securities[symbol].Close:
                self.selected.append(symbol)
                self.Debug(str(symbol))
    
    
    def OnSecuritiesChanged(self, changes):
        
        for security in changes.RemovedSecurities:
            if security.Symbol in self.data:
                del self.data[security.Symbol]
        
        for security in changes.AddedSecurities:
            if security.Symbol not in self.data:
                self.data[security.Symbol] = SymbolData(security.Symbol, self.ema_period, self.sma_period, self.bb_period, self.bb_k, self)
    
    
    def AtMarketOpen(self):
        
        for symbol in self.data.keys():
            gap = (self.Securities[symbol].Close - self.data[symbol].LastClose) / self.data[symbol].LastClose
            if gap > self.gap_distance:
                self.data[symbol].GapUp = True
            else:
                self.data[symbol].GapUp = False
            
    
    def BeforeMarketCloses(self):
        
        for symbol in self.data.keys():
            self.data[symbol].LastClose = self.Securities[symbol].Close
            

class SymbolData(object):
    def __init__(self, symbol, ema, sma, bb, k, algorithm):
        self.Symbol = symbol
        self.LastClose = 0
        self.GapUp = False
        self.EMA = ExponentialMovingAverage(ema)
        self.SMA = SimpleMovingAverage(sma)
        self.BB = BollingerBands(bb, k, MovingAverageType.Exponential) 
        
        algorithm.RegisterIndicator(symbol, self.EMA, Resolution.Minute, Field.Close)
        algorithm.RegisterIndicator(symbol, self.SMA, Resolution.Minute, Field.Close)
        algorithm.RegisterIndicator(symbol, self.BB, Resolution.Minute, Field.Close)
        
        # Warm up indicators
        history = algorithm.History(symbol, 1, Resolution.Daily)
        if history.empty or 'close' not in history.columns:
            return
        for index, row in history.loc[symbol].iterrows():
            self.LastClose = row['close']
            
        history = algorithm.History(symbol, max(ema, sma, bb), Resolution.Minute)
        if history.empty or 'close' not in history.columns:
            return
        for index, row in history.loc[symbol].iterrows():
            self.EMA.Update(index, row['close'])
            self.SMA.Update(index, row['close'])
            self.BB.Update(index, row['close'])