class UniverseSelection(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 4, 1)  # Set Start Date
        self.SetEndDate(2020, 5, 7) # 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_stock_price = 10
        
        self.min_days_after_earnings = 10
        self.max_days_after_earnings = 80
        self.ema_period =  8  #  8 period on  5 Min TF is  40 on 1 Min TF 
        self.sma_period = 55  # 55 period on  5 Min TF is 275 on 1 Min TF 
        self.bb_period  = 20  # 20 period on 30 min TF is 600 on 1 Min TF
        self.bb_k = 2
        self.gap_distance = 0.02 # 2%
        self.data = {}
        self.selectedStocks = []
        
        
        self.Schedule.On(self.DateRules.EveryDay(),
                        self.TimeRules.AfterMarketOpen(self.SPY, 30),
                        self.SelectStocks)
            
        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) > self.min_stock_price)]
        
        # 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 select    
    def SelectStocks(self):
        
        for symbol in self.data.keys():
            
            sd = self.data[symbol]
            isReady = sd.EMA.IsReady and sd.SMA.IsReady and sd.BB.IsReady
            
            if isReady and 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.selectedStocks.append(symbol)
                self.Debug(f'{symbol.ToString()} {self.Time.ctime()} GapUp:{sd.GapUp} EMA:{sd.EMA} SMA:{sd.SMA} BB:{sd.BB.UpperBand.Current.Value} Close:{self.Securities[symbol].Close}')
    
    
    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].Open - 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, emaPeriod, smaPeriod, bbPeriod, bbKvalue, algorithm):
        self.Symbol = symbol
        self.LastClose = 0
        self.GapUp = False
        self.EMA = ExponentialMovingAverage(emaPeriod)
        self.SMA = SimpleMovingAverage(smaPeriod)
        self.BB = BollingerBands(bbPeriod, bbKvalue, MovingAverageType.Exponential)
        self.algorithm = algorithm
        
        self.consolidator_5min = TradeBarConsolidator(timedelta(minutes=5))
        self.consolidator_30min = TradeBarConsolidator(timedelta(minutes=30))
        
        # Exit if no daily data. We need 1 bar of data (one day)
        # ---------------------------------------------------------
        history = algorithm.History(symbol, 1, Resolution.Daily)
        if history.empty or 'close' not in history.columns:
            return
        
        # iterate over historical data, get/store last close (of prev bar)
        # ------------------------------------------------------------------
        else:
            for index, row in history.loc[symbol].iterrows():
                self.LastClose = row['close']
        
        history = algorithm.History(symbol, max(emaPeriod*5, smaPeriod*5, bbPeriod*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
            tradeBar.Symbol = symbol
            self.consolidator_5min.Update(tradeBar)
            self.consolidator_30min.Update(tradeBar)
        
        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)

    def dispose(self):
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.Symbol, self.consolidator_5min)
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.Symbol, self.consolidator_30min)