# region imports
from AlgorithmImports import *
# endregion

class Test(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2022, 4, 11)  # Set Start Date
        self.SetEndDate(2022, 4, 13)  # Set End Date
        self.SetCash(100000)  # Set Strategy Cash

        # Raw data normalization
        self.SetSecurityInitializer(lambda x: x.SetDataNormalizationMode(DataNormalizationMode.Adjusted))

        # Add SPY and define symbols
        self.AddEquity("SPY", resolution=Resolution.Minute, extendedMarketHours=True)
        self.selectedSymbols = []

        # Add universe
        self.AddUniverse(self.CoarseUniverseSelection)

        # Scheduled event
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.At(9, 25), self.CheckStocks)
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.At(21, 00), self.ResetUniverse)

    def OnData(self, data: Slice):
        pass

    def CoarseUniverseSelection(self, coarse):
        self.Debug(f"Universe selection time: {self.Time}")
        for stock in coarse:
            if (stock.Price > 10):
                # Subscribe to equity
                self.AddEquity(stock.Symbol, resolution=Resolution.Minute, extendedMarketHours=True)
                self.selectedSymbols.append(stock.Symbol)
        
        return []

    
    def CheckStocks(self):
        self.Debug(f"Check stocks time: {self.Time}")
        number_not_in_active_securities = 0
        # Remove security
        for symbol in self.selectedSymbols:
            try:
                security = self.ActiveSecurities[symbol] # Can't find security even though it was added at the start of the day by CoarseUniverseSelection
                price = security.Price
            except:
                number_not_in_active_securities += 1
        
        self.Debug(f"{number_not_in_active_securities} of {len(self.selectedSymbols)} not in active securities")

    def ResetUniverse(self):
        self.Debug(f"Resetting universe: {self.Time}")
        self.selectedSymbols = []
        self.yesterdayClose = {}

        for kvp in self.ActiveSecurities:
            symbol = kvp.Key
            self.RemoveSecurity(symbol)
    
        self.AddEquity("SPY", resolution=Resolution.Minute, extendedMarketHours=True)