# region imports
from AlgorithmImports import *
import time
# endregion

class TestRemoveUniverse(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 5, 9)  # Set Start Date
        self.SetEndDate(2021, 5, 15)
        self.SetCash(100000)  # Set Strategy Cash

        self.AddEquity("SPY", resolution=Resolution.Minute, extendedMarketHours=True)        

        # Add universe
        self.UniverseSettings.Resolution = Resolution.Minute
        self.UniverseSettings.ExtendedMarketHours = True
        self.UniverseSettings.FillForward = True
        self.coarse_universe = self.AddUniverse(self.CoarseUniverseSelection)

        self.AddUniverseSelection(ScheduledUniverseSelectionModel(
            self.DateRules.EveryDay("SPY"), 
            self.TimeRules.AfterMarketOpen("SPY", -5),
            self.ScheduledSymbolSelect))

        self.cur_time = time.time()
        self.all_daily_symbols = []
        self.num_sub_to = 10

        self.UniverseSettings.MinimumTimeInUniverse = timedelta(minutes=1)

    def OnData(self, data: Slice):
        if (self.Time.hour == 8 and self.Time.minute == 31):
            self.Debug(f"Time: {self.Time} | Length of slice: {len(data.Bars.Values)}")
        
        if (self.Time.hour == 9 and self.Time.minute == 31):
            self.Debug(f"Time: {self.Time} | Length of slice: {len(data.Bars.Values)}")

    def CoarseUniverseSelection(self, coarse):
        self.all_daily_symbols = []
        self.cur_time = time.time()
        symbols = []
        for stock in coarse:
            if (stock.AdjustedPrice > 1):
                self.all_daily_symbols.append(stock.Symbol)
        
        # Return temporary symbols
        return self.all_daily_symbols[:self.num_sub_to]
    
    def ScheduledSymbolSelect(self, date):
        # Overwrite current universe and only subscribe to the following symbols
        return self.all_daily_symbols[100:102]