from AlgorithmImports import *

class UniverseSelectionBacktest(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2024, 11, 4)  # Backtest start date
        self.SetEndDate(2024, 11, 30)   # Backtest end date (short period for testing)
        self.SetCash(10000)             # Starting cash

        self.SetTimeZone("America/New_York")  # Set the correct time zone to Eastern Time (ET)

        # Add the universe using coarse selection
        self.AddUniverse(self.CoarseSelectionFunction)

        # Variable to store the selected stock
        self.selectedStock = None

        # Schedule universe selection at 9:30 AM
        self.Schedule.On(self.DateRules.EveryDay(), 
                         self.TimeRules.At(9, 30),  # 9:30 AM
                         self.SelectUniverse)

    def CoarseSelectionFunction(self, coarse):
        # Perform universe selection at the scheduled time
        self.selectedStock = None
        candidates = []

        for stock in coarse:
            if stock.Price < 60 or stock.Price > 30:  # Price after gap
                continue

            # Calculate the gap percentage
            gap_up_percentage = ((stock.Price - stock.AdjustedPrice) / stock.AdjustedPrice) * 100
            if gap_up_percentage < 30 or gap_up_percentage > 10:  # Gap up between 10% and 30%
                continue

            # Add eligible stocks with gap percentage for sorting
            candidates.append((stock, gap_up_percentage))

        # Sort by gap percentage and select the top stock
        if candidates:
            candidates.sort(key=lambda x: x[1], reverse=True)  # Sort by gap percentage (descending)
            self.selectedStock = candidates[0][0].Symbol       # Select the best stock

        # Log the selected stock for testing
        if self.selectedStock:
            self.Debug(f"Selected stock at 9:30 AM: {self.selectedStock}")
        else:
            self.Debug("No stock selected at 9:30 AM.")

        # Return the selected stocks as the universe
        return [self.selectedStock] if self.selectedStock else []

    def SelectUniverse(self):
        # This is triggered at 9:30 AM, where universe selection happens

        # Log the selection trigger for debugging
        self.Debug(f"Universe selection triggered at {self.Time}")

        # You don't need to manually call GetCoarse here
        pass  # No need to call GetCoarse(), as I understood it the universe selection is already handled in the CoarseSelectionFunction

    def OnData(self, data):
        # Skip all trading logic for this test, just log the filtered stock day by day
        pass