from AlgorithmImports import *

class OrderBookImbalanceStrategy(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 4)
        self.SetEndDate(2020, 9, 10)
        self.SetCash(100000)
        
        # Add SPY with quote data explicitly requested
        equity = self.AddEquity("SPY", Resolution.Minute)
        equity.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.symbol = equity.Symbol
        
        # Initialize SMA indicators for bid and ask volumes
        self.bid_volume_sma = SimpleMovingAverage(5)
        self.ask_volume_sma = SimpleMovingAverage(5)
        
        # Position tracking flags
        self.is_long = False
        self.is_short = False
        
        # Schedule trading decisions every 5 minutes
        self.Schedule.On(self.DateRules.EveryDay(self.symbol), 
                         self.TimeRules.Every(timedelta(minutes=5)), 
                         self.ExecuteTrades)
        
        # Create custom chart to track OBI values
        self.Plot("Strategy", "OBI", 0)

    def OnData(self, data):
        # Specifically check for QuoteBars data
        if not data.QuoteBars.ContainsKey(self.symbol):
            return
        
        # Get the quote bar
        quote_bar = data.QuoteBars[self.symbol]
        
        # Update SMAs with bid and ask volumes
        self.bid_volume_sma.Update(self.Time, quote_bar.LastBidSize)
        self.ask_volume_sma.Update(self.Time, quote_bar.LastAskSize)
    
    def ExecuteTrades(self):
        # Check if SMAs are ready
        if not self.bid_volume_sma.IsReady or not self.ask_volume_sma.IsReady:
            return
            
        # Get current values from SMAs
        bid_vol = self.bid_volume_sma.Current.Value
        ask_vol = self.ask_volume_sma.Current.Value
        
        # Avoid division by zero
        if bid_vol + ask_vol == 0:
            self.Debug("Skipping trade evaluation: Bid + Ask volume is zero")
            return
            
        # Calculate Order Book Imbalance using the formula: OBI(t) = (B(t) - A(t)) / (B(t) + A(t))
        obi = (bid_vol - ask_vol) / (bid_vol + ask_vol)
        
        # Plot the current OBI value
        self.Plot("Strategy", "OBI", obi)
        
        # Log OBI value
        self.Debug(f"Time: {self.Time}, OBI: {obi}")
        
        # Trading logic based on OBI thresholds
        # Long signal - OBI > 0.7
        if obi > 0.9 and not self.is_long:
            if self.is_short:
                # Close short position first
                self.Liquidate()
                self.is_short = False
                
            # Enter long position
            self.SetHoldings(self.symbol, 0.95)  # Using 95% of portfolio
            self.is_long = True
            self.Debug(f"LONG signal at OBI: {obi}")
            
        # Short signal - OBI < -0.7
        elif obi < -0.9 and not self.is_short:
            if self.is_long:
                # Close long position first
                self.Liquidate()
                self.is_long = False
                
            # Enter short position
            self.SetHoldings(self.symbol, -0.95)  # Using 95% of portfolio for short
            self.is_short = True
            self.Debug(f"SHORT signal at OBI: {obi}")
            
        # Exit long position when OBI drops below 0.7
        elif self.is_long and obi < 0.7:
            self.Liquidate()
            self.is_long = False
            self.Debug(f"EXIT LONG at OBI: {obi}")
            
        # Exit short position when OBI rises above -0.7
        elif self.is_short and obi > -0.7:
            self.Liquidate()
            self.is_short = False
            self.Debug(f"EXIT SHORT at OBI: {obi}")