from AlgorithmImports import *
from math import *

class ESFuturesVWAP9EMACrossAlgorithm(QCAlgorithm):
    def Initialize(self):
        # Set start and end date for backtesting
        self.SetStartDate(2024, 1, 1)
        #self.SetEndDate(2024, 1, 1)
        
        # Set cash allocation
        self.SetCash(100000)

        resolution = Resolution.Second
        
        # Add ES futures contract (S&P 500 e-mini)
        self.es_future = self.AddFuture(Futures.Indices.SP500EMini, resolution, Market.CME)
        self.es_future.SetFilter(timedelta(0), timedelta(180))  # Set contract expiry filter
        
        # Variables to keep track of positions and entry state
        self.current_position = None
        self.position_entry_time = None  # Track entry time for position
        
        self.contract = {self.es_future.Symbol: None}
        
        # Schedule to liquidate positions outside regular market hours
        self.Schedule.On(
            self.DateRules.EveryDay(self.es_future.Symbol),
            self.TimeRules.BeforeMarketClose(self.es_future.Symbol, 1),
            self.ExitPositions
        )

        # Create indicators
        self.vwap = self.VWAP(self.es_future.Symbol, 14, resolution)
        self.ema9 = self.EMA(self.es_future.Symbol, 9, resolution)
        
        # Rolling windows to keep track of EMA, VWAP values, and recent bars
        self.ema_window = RollingWindow[float](2)    # Store the last 2 EMA values
        self.vwap_window = RollingWindow[float](2)    # Store the last 2 VWAP values
        self.price_window = RollingWindow[TradeBar](2)    # Store the last 2 bars
        
        # Request consolidated minute data
        self.consolidator = TradeBarConsolidator(timedelta(minutes=10))
        self.consolidator.DataConsolidated += self.OnDataConsolidated
        
        # Register the consolidator
        self.SubscriptionManager.AddConsolidator(self.es_future.Symbol, self.consolidator)

    def OnData(self, slice):
        for kvp in slice.FutureChains:
            symbol = kvp.Key
            chain = kvp.Value
            
            # Find the most liquid contract
            if symbol in self.contract:
                most_liquid_contract = sorted(chain, key=lambda contract: contract.OpenInterest, reverse=True)[0]
                self.contract[symbol] = most_liquid_contract

        # Check profit and stop loss targets for active positions
        if self.current_position is not None:
            contract = self.contract[self.es_future.Symbol]
            if contract is None: 
                return
            symbol = contract.Symbol

            # Retrieve current position details
            current_holdings = self.Portfolio[symbol]
            entry_price = current_holdings.AveragePrice if current_holdings.Invested else None

            if entry_price is not None:
                # Define profit and stop loss targets
                profit_target_points = 10
                stop_loss_points = 10

                if self.current_position == "long":
                    # Profit target and stop loss for long position
                    stop_loss_price = entry_price - stop_loss_points
                    profit_target_price = entry_price + profit_target_points

                    if slice[symbol].Close <= stop_loss_price:
                        self.Liquidate(symbol)
                        self.current_position = None
                        self.Debug(f"Long Stop Loss Hit: {self.Time}, Price: {slice[symbol].Close}")

                    elif slice[symbol].Close >= profit_target_price:
                        self.Liquidate(symbol)
                        self.current_position = None
                        self.Debug(f"Long Profit Target Hit: {self.Time}, Price: {slice[symbol].Close}")

                elif self.current_position == "short":
                    # Profit target and stop loss for short position
                    stop_loss_price = entry_price + stop_loss_points
                    profit_target_price = entry_price - profit_target_points

                    if slice[symbol].Close >= stop_loss_price:
                        self.Liquidate(symbol)
                        self.current_position = None
                        self.Debug(f"Short Stop Loss Hit: {self.Time}, Price: {slice[symbol].Close}")

                    elif slice[symbol].Close <= profit_target_price:
                        self.Liquidate(symbol)
                        self.current_position = None
                        self.Debug(f"Short Profit Target Hit: {self.Time}, Price: {slice[symbol].Close}")

    def OnDataConsolidated(self, sender, bar):
        contract = self.contract[self.es_future.Symbol]
        if contract is None: 
            return
        symbol = contract.Symbol

        # Add the latest bar to the rolling window
        self.price_window.Add(bar)

        # Ensure indicators are ready and rolling windows have enough data
        if not self.vwap.IsReady or not self.ema9.IsReady or self.price_window.Count < 2:
            return

        # Add the latest EMA and VWAP values to their respective rolling windows
        self.ema_window.Add(self.ema9.Current.Value)
        self.vwap_window.Add(self.vwap.Current.Value)

        # Ensure rolling windows have enough data
        if self.ema_window.Count < 2 or self.vwap_window.Count < 2:
            return

        # Get the current and previous values from rolling windows
        current_ema = self.ema_window[0]
        previous_ema = self.ema_window[1]
        current_vwap = self.vwap_window[0]
        previous_vwap = self.vwap_window[1]

        # Get the current and previous bars
        current_bar = self.price_window[0]
        previous_bar = self.price_window[1]

        # Long entry condition: 9EMA crosses above VWAP and the last 2 bars close above VWAP
        if previous_ema <= previous_vwap and current_ema > current_vwap and \
           current_bar.Close > previous_bar.Close and previous_bar.Close > previous_vwap:
            if self.current_position != "long" or not self.Portfolio.Invested:
                #self.Liquidate()
                limit_price = ceil(max(current_bar.High, previous_bar.High) + 1)  # Set limit price slightly above the high
                #self.LimitOrder(symbol, 1, limitPrice=limit_price)  # Use limit order for long entry
                self.MarketOrder(symbol, 1)
                self.current_position = "long"
                self.Debug(f"Entered Long: {self.Time}, EMA: {current_ema}, VWAP: {current_vwap}")

        # Short entry condition: 9EMA crosses below VWAP and the last 2 bars close below VWAP
        elif previous_ema >= previous_vwap and current_ema < current_vwap and \
             current_bar.Close < previous_bar.Close and previous_bar.Close < previous_vwap:
            if self.current_position != "short" or not self.Portfolio.Invested:
                #self.Liquidate()
                limit_price = ceil(min(current_bar.Low, previous_bar.Low) - 1)  # Set limit price slightly below the low
                #self.LimitOrder(symbol, -1, limitPrice=limit_price)  # Use limit order for short entry
                self.MarketOrder(symbol, -1)
                self.current_position = "short"
                self.Debug(f"Entered Short: {self.Time}, EMA: {current_ema}, VWAP: {current_vwap}")

    def ExitPositions(self):
        # Liquidate all positions if they exist
        if self.current_position is not None:
            self.Liquidate()
            self.current_position = None
            self.Debug(f"Position liquidated after the next candle: {self.Time}")