from AlgorithmImports import *

class DailyNQFutures(QCAlgorithm):

    def Initialize(self):
        # Set the start and end dates for the backtest
        self.SetStartDate(2020, 1, 1)
        #self.SetEndDate(2024, 10, 1)
        
        # Set the initial cash balance for the algorithm
        self.SetCash(100000)

        # Add the Nasdaq 100 E-mini futures contract to trade
        futureNQ = self.AddFuture(Futures.Indices.NASDAQ100EMini, Resolution.Minute)

        # Filter to select contracts that expire within 90 days from today
        futureNQ.SetFilter(timedelta(0), timedelta(90))

        # Initialize the contract symbol to None (will be updated later)
        self.contractSymbol = None

        # Set the parameters for the SMA periods
        self.sma_5min_period = 7  # 5-minute SMA period
        self.sma_30min_period = 21  # 30-minute SMA period

        # Define the indicators (will be set up later)
        self.sma_5min = None
        self.sma_30min = None

        # Schedule the daily buy/sell event (check for weekdays and market open inside methods)
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(15, 54), Action(self.TradeNQ))
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(15, 59), Action(self.SellNQ))

    def OnData(self, slice):
        # Iterate through future contracts in the slice (market data for the current time step)
        for chain in slice.FutureChains:
            contracts = list(chain.Value)
            if len(contracts) > 0:
                # Sort the contracts by expiration date
                nearest_contract = sorted(contracts, key=lambda x: x.Expiry)[0]
                
                # Check if the current contract is expiring in less than 4 days
                if nearest_contract.Expiry - self.Time < timedelta(days=4):
                    # Ensure there is a next contract available
                    if len(contracts) > 1:
                        next_contract = sorted(contracts, key=lambda x: x.Expiry)[1]
                        if self.contractSymbol != next_contract.Symbol:
                            self.contractSymbol = next_contract.Symbol
                            self.SetupIndicators(self.contractSymbol)
                            self.Debug(f"Rolled to new contract: {self.contractSymbol}, Expiration: {next_contract.Expiry}")
                else:
                    # If not rolling, just track the nearest contract
                    if self.contractSymbol != nearest_contract.Symbol:
                        self.contractSymbol = nearest_contract.Symbol
                        self.SetupIndicators(self.contractSymbol)

    def SetupIndicators(self, symbol):
        # Set up the 5-minute and 30-minute SMA indicators for the given contract symbol
        self.sma_5min = self.SMA(symbol, self.sma_5min_period, Resolution.Minute)
        self.sma_30min = self.SMA(symbol, self.sma_30min_period, Resolution.Minute)

    def TradeNQ(self):
        # Only trade if it's a weekday, the market is open, and we have a contract symbol
        if self.Time.weekday() < 5:  # 0 = Monday, 4 = Friday
            if self.contractSymbol and self.IsMarketOpen(self.contractSymbol):
                # Make sure both SMAs are ready (i.e., enough data points are available)
                if self.sma_5min.IsReady and self.sma_30min.IsReady:
                    # Check current invested position
                    current_position = self.Portfolio[self.contractSymbol].Quantity
                    
                    # Long if 5-minute SMA is above the 30-minute SMA (bullish signal)
                    if self.sma_5min.Current.Value > self.sma_30min.Current.Value:
                        if current_position <= 0:  # If not already long, buy
                            self.MarketOrder(self.contractSymbol, 1)  # Buy 1 contract
                            self.Debug(f"Bought 1 contract at {self.Time}, 5-min SMA: {self.sma_5min.Current.Value}, 30-min SMA: {self.sma_30min.Current.Value}")
                    
                    # Short if 5-minute SMA is below the 30-minute SMA (bearish signal)
                    elif self.sma_5min.Current.Value < self.sma_30min.Current.Value:
                        if current_position >= 0:  # If not already short, sell
                            self.MarketOrder(self.contractSymbol, -1)  # Short 1 contract
                            self.Debug(f"Shorted 1 contract at {self.Time}, 5-min SMA: {self.sma_5min.Current.Value}, 30-min SMA: {self.sma_30min.Current.Value}")
    
    def SellNQ(self):
        # Only liquidate if it's a weekday, the market is open, and we have a contract symbol
        if self.Time.weekday() < 5:  # 0 = Monday, 4 = Friday
            if self.contractSymbol and self.IsMarketOpen(self.contractSymbol):
                if self.Portfolio[self.contractSymbol].Invested:
                    self.Liquidate(self.contractSymbol)
                    self.Debug(f"Liquidated position at {self.Time}")

    def IsMarketOpen(self, symbol):
        # Check if the market for the given symbol is open at the current time
        exchange = self.Securities[symbol].Exchange
        return exchange.DateTimeIsOpen(self.Time)