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)
        
        # Add OBI persistence tracking - store last N OBI values
        self.obi_window = RollingWindow[float](3)  # Track last 3 OBI values (15 minutes)
        
        # Add volume analysis - track total order book volume (bid + ask)
        self.total_volume_sma = SimpleMovingAverage(20)  # 20-period volume SMA
        
        # Position tracking flags
        self.is_long = False
        self.is_short = False
        
        # Entry/exit thresholds
        self.obi_entry_threshold = 0.7  # Align with project specification
        self.obi_exit_threshold = 0.7
        self.volume_multiplier = 1.2  # Require volume 20% above average
        
        # Schedule trading decisions every 5 minutes
        self.Schedule.On(self.DateRules.EveryDay(self.symbol), 
                         self.TimeRules.Every(timedelta(minutes=5)), 
                         self.ExecuteTrades)
        
        # Create custom charts
        self.Plot("Strategy", "OBI", 0)
        self.Plot("Strategy", "Entry Threshold", self.obi_entry_threshold)
        self.Plot("Strategy", "Exit Threshold", self.obi_exit_threshold)
        self.Plot("Volume", "Total Order Volume", 0)
        self.Plot("Volume", "Average Order Volume", 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)
        
        # Update total order book volume SMA
        total_volume = quote_bar.LastBidSize + quote_bar.LastAskSize
        self.total_volume_sma.Update(self.Time, total_volume)
    
    def ExecuteTrades(self):
        # Check if indicators are ready
        if not self.bid_volume_sma.IsReady or not self.ask_volume_sma.IsReady or not self.total_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)
        
        # Add current OBI to rolling window for persistence tracking
        self.obi_window.Add(obi)
        
        # Plot the current OBI value
        self.Plot("Strategy", "OBI", obi)
        
        # Log OBI value
        self.Debug(f"Time: {self.Time}, OBI: {obi}")
        
        # Check for significant order book volume
        total_current_volume = bid_vol + ask_vol
        avg_volume = self.total_volume_sma.Current.Value
        
        # Plot volume data
        self.Plot("Volume", "Total Order Volume", total_current_volume)
        self.Plot("Volume", "Average Order Volume", avg_volume)
        
        # Check if current total volume is above average (significant)
        has_significant_volume = total_current_volume > avg_volume * self.volume_multiplier
        
        if not has_significant_volume:
            self.Debug(f"Volume too low: Current={total_current_volume}, Avg={avg_volume}")
        
        # Check for persistent OBI (majority of values in window exceed threshold)
        has_persistent_bullish_obi = False
        has_persistent_bearish_obi = False
        
        if self.obi_window.IsReady:
            # Get values from the rolling window
            obi_values = [self.obi_window[i] for i in range(self.obi_window.Count)]
            
            # Check if at least 2 out of 3 recent OBI values are above threshold
            bullish_count = sum(1 for o in obi_values if o > self.obi_entry_threshold * 0.8)
            bearish_count = sum(1 for o in obi_values if o < -self.obi_entry_threshold * 0.8)
            
            has_persistent_bullish_obi = bullish_count >= 2
            has_persistent_bearish_obi = bearish_count >= 2
            
            self.Debug(f"OBI Window: [{', '.join([str(round(o, 4)) for o in obi_values])}]")
            self.Debug(f"Persistent Bullish: {has_persistent_bullish_obi}, Persistent Bearish: {has_persistent_bearish_obi}")
        
        # Scale position size based on OBI strength
        position_size = min(0.95, 0.5 + abs(obi) * 0.5)  # Scale from 0.5 to 0.95 based on OBI
        
        # Trading logic based on OBI thresholds, persistence, and volume
        # Long signal - OBI > threshold, persistent bullish, and significant volume
        if obi > self.obi_entry_threshold and has_persistent_bullish_obi and has_significant_volume and not self.is_long:
            if self.is_short:
                # Close short position first
                self.Liquidate()
                self.is_short = False
                
            # Enter long position with size based on OBI strength
            self.SetHoldings(self.symbol, position_size)
            self.is_long = True
            self.Debug(f"LONG signal at OBI: {obi} (Position Size: {position_size:.2f}, Persistent: {has_persistent_bullish_obi}, Volume: {has_significant_volume})")
            
        # Short signal - OBI < -threshold, persistent bearish, and significant volume
        elif obi < -self.obi_entry_threshold and has_persistent_bearish_obi and has_significant_volume and not self.is_short:
            if self.is_long:
                # Close long position first
                self.Liquidate()
                self.is_long = False
                
            # Enter short position with size based on OBI strength
            self.SetHoldings(self.symbol, -position_size)
            self.is_short = True
            self.Debug(f"SHORT signal at OBI: {obi} (Position Size: {position_size:.2f}, Persistent: {has_persistent_bearish_obi}, Volume: {has_significant_volume})")
            
        # Exit long position when OBI drops below threshold or volume dries up
        elif self.is_long and (obi < self.obi_exit_threshold or not has_significant_volume):
            self.Liquidate()
            self.is_long = False
            self.Debug(f"EXIT LONG at OBI: {obi}, Volume Significant: {has_significant_volume}")
            
        # Exit short position when OBI rises above -threshold or volume dries up
        elif self.is_short and (obi > -self.obi_exit_threshold or not has_significant_volume):
            self.Liquidate()
            self.is_short = False
            self.Debug(f"EXIT SHORT at OBI: {obi}, Volume Significant: {has_significant_volume}")