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 = 8 # 5 min timeframe
        self.sma_period = 55 # 5 min timeframe
        self.bb_period = 20 # 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:
            symbol_data = self.data.pop(security.Symbol, None)
            if symbol_data:
                symbol_data.dispose()
        
        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():
            if self.data[symbol].LastClose == 0:
                self.data[symbol].GapUp = False
                continue
            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)
        self.algorithm = algorithm
        
        self.consolidator_5min = TradeBarConsolidator(timedelta(minutes=5))
        self.consolidator_30min = TradeBarConsolidator(timedelta(minutes=30))
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator_5min)
        algorithm.SubscriptionManager.AddConsolidator(symbol, self.consolidator_30min)
        algorithm.RegisterIndicator(symbol, self.EMA, self.consolidator_5min)
        algorithm.RegisterIndicator(symbol, self.SMA, self.consolidator_5min)
        algorithm.RegisterIndicator(symbol, self.BB, self.consolidator_30min)
        
        # 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*5, sma*5, bb*30), Resolution.Minute)
        #if history.empty or 'close' not in history.columns:
        #    return
        #for index, row in history.loc[symbol].iterrows():
        #    tradeBar = TradeBar()
        #    tradeBar.Close = row['close']
        #    tradeBar.Open = row['open']
        #    tradeBar.High = row['high']
        #    tradeBar.Low = row['low']
        #    tradeBar.Volume = row['volume']
        #    tradeBar.Time = index
        #    self.consolidator_5min.Update(tradeBar)
        #    self.consolidator_30min.Update(tradeBar)
    
    def dispose(self):
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.Symbol, self.consolidator_5min)
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.Symbol, self.consolidator_30min)