# region imports
from AlgorithmImports import *
from datetime import timezone
# endregion

class CustomTimeRuleAlgorithm(QCAlgorithm):
    def initialize(self):
        self.set_start_date(2021, 1, 4)
        self.set_end_date(2021, 1, 5)
        self.set_cash(100000)
        self.add_equity("SPY", Resolution.MINUTE)
        self.use_extended_hrs = False  # Set to True if you want to use extended market hours

        # Schedule an event using the custom time rule
        self.schedule.on(
            self.date_rules.every_day("SPY"),
            self.custom_time_rule("SPY", timedelta(minutes=5)),
            self.my_scheduled_function
        )

    def custom_time_rule(self, symbol, interval):
        # Define the custom time rule using FuncTimeRule with a lambda function
        return FuncTimeRule(
            name="EveryIntervalDuringMarketHours",
            createUtcEventTimesFunction=lambda dates: [
                dt for date in dates
                for dt in self.get_custom_times(symbol, date, interval)
            ]
        )
    
    def get_custom_times(self, symbol, date, interval):
        # Get market open and close times for the specified date
        exchange = self.securities[symbol].exchange
        market_open = exchange.hours.get_next_market_open(date, self.use_extended_hrs)
        market_close = exchange.hours.get_next_market_close(market_open, self.use_extended_hrs)

        # Convert market times from EST to UTC
        market_open = market_open.replace(tzinfo=timezone(timedelta(hours=-5))).astimezone(timezone.utc)
        market_close = market_close.replace(tzinfo=timezone(timedelta(hours=-5))).astimezone(timezone.utc)
        
        # Generate times every 'interval' minutes between market open and close
        times = []
        current_time = market_open
        while current_time < market_close:
            times.append(current_time)
            current_time += interval
        return times

    def my_scheduled_function(self):
        self.log(f"Executed at: {self.time}")
        self.debug(f"Executed at: {self.time}")