from AlgorithmImports import *

class OpeningRangeBreakoutStocksInPlay(QCAlgorithm):
    
    def Initialize(self):
        # Configurable parameters
        self.START_YEAR = 2025
        self.START_MONTH = 1
        self.START_DAY = 1
        self.END_YEAR = 2025
        self.END_MONTH = 8
        self.END_DAY = 1
        self.INITIAL_CASH = 500_000
        
        self.MIN_PRICE = 5
        self.MIN_DOLLAR_VOLUME = 50_000_000
        self.TOP_COARSE_SYMBOLS = 100
        
        self.MIN_MARKET_CAP = 2_000_000_000
        self.TOP_FINE_SYMBOLS = 1
        
        self.WARMUP_DAYS = 20
        self.WARMUP_RESOLUTION = Resolution.Daily
        
        self.LIQUIDATION_MINUTES_BEFORE_CLOSE = 10
        self.MARKET_SYMBOL = "SPY"
        
        self.OR_START_HOUR = 9
        self.OR_START_MINUTE = 30
        self.OR_END_MINUTE = 35
        
        self.TRADE_START_HOUR = 9
        self.TRADE_START_MINUTE = 35
        
        self.POSITION_SIZE = 0.95
        self.STOP_LOSS_PCT = 0.02  # 2% stop-loss
        self.TAKE_PROFIT_PCT = 1.00  # 100% take-profit (note: original comment said 2%, but code used equivalent of 100%)

        self.SetStartDate(self.START_YEAR, self.START_MONTH, self.START_DAY)
        self.SetEndDate(self.END_YEAR, self.END_MONTH, self.END_DAY)
        self.SetCash(self.INITIAL_CASH)
        
        # Universe selection for Stocks in Play: Use coarse for high dollar volume, fine for market cap filter
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        # Add VIX for volatility check (data only)
        self.vix = self.AddIndex("VIX", Resolution.Minute).Symbol
        
        # Add tradable VIX ETF (UVXY)
        self.AddEquity("UVXY", Resolution.Minute)
        
        # Store dollar volume from coarse for sorting in fine
        self.dollar_vol = {}
        
        # Warm up for fundamentals
        self.SetWarmUp(self.WARMUP_DAYS, self.WARMUP_RESOLUTION)
        
        # Schedule end-of-day liquidation
        self.Schedule.On(self.DateRules.EveryDay(), 
                         self.TimeRules.BeforeMarketClose(self.MARKET_SYMBOL, self.LIQUIDATION_MINUTES_BEFORE_CLOSE), 
                         self.LiquidateAll)
        
        # Track OR high/low per symbol, entry prices for stops/profits
        self.or_high = {}
        self.or_low = {}
        self.entry_prices = {}
    
    def CoarseSelectionFunction(self, coarse):
        # Filter: Price > MIN_PRICE (avoid low-priced stocks), DollarVolume > MIN_DOLLAR_VOLUME (high activity/liquidity), HasFundamentalData
        filtered = [c for c in coarse if c.Price > self.MIN_PRICE and c.DollarVolume > self.MIN_DOLLAR_VOLUME and c.HasFundamentalData]
        
        # Sort by DollarVolume descending, take top TOP_COARSE_SYMBOLS for fine processing
        sorted_by_dollar_vol = sorted(filtered, key=lambda c: c.DollarVolume, reverse=True)[:self.TOP_COARSE_SYMBOLS]
        
        # Store dollar volume for fine sorting
        self.dollar_vol = {c.Symbol: c.DollarVolume for c in sorted_by_dollar_vol}
        
        return [c.Symbol for c in sorted_by_dollar_vol]
    
    def FineSelectionFunction(self, fine):
        # Filter: MarketCap > MIN_MARKET_CAP (avoid low cap)
        filtered_fine = [f for f in fine if f.MarketCap > self.MIN_MARKET_CAP]
        
        # Sort by stored DollarVolume descending, take top TOP_FINE_SYMBOLS
        sorted_fine = sorted(filtered_fine, key=lambda f: self.dollar_vol.get(f.Symbol, 0), reverse=True)[:self.TOP_FINE_SYMBOLS]
        
        return [f.Symbol for f in sorted_fine]
    
    def OnData(self, data):
        if not self.IsMarketOpen(self.MARKET_SYMBOL): return  # US equities hours
        
        # Get current VIX value for volatility check
        if self.vix in data.Bars:
            vix_close = data.Bars[self.vix].Close
            # Dynamically adjust position size: halve in high vol (VIX > 30)
            current_position_size = self.POSITION_SIZE / 2 if vix_close > 30 else self.POSITION_SIZE
        else:
            current_position_size = self.POSITION_SIZE  # Default if VIX data unavailable
        
        for symbol in list(self.UniverseManager.ActiveSecurities.Keys):
            if symbol not in data.Bars: continue
            if symbol == self.vix: continue  # Skip VIX, use only for data
            
            bar = data.Bars[symbol]
            time = self.Time.time()
            
            # Calculate 5-min OR: From open (9:30) to 9:35 ET
            if time.hour == self.OR_START_HOUR and time.minute >= self.OR_START_MINUTE and time.minute < self.OR_END_MINUTE:
                if symbol not in self.or_high:
                    self.or_high[symbol] = bar.High
                    self.or_low[symbol] = bar.Low
                else:
                    self.or_high[symbol] = max(self.or_high[symbol], bar.High)
                    self.or_low[symbol] = min(self.or_low[symbol], bar.Low)
            
            # After 9:35, check for upside breakout if not invested
            elif time.hour >= self.TRADE_START_HOUR and time.minute >= self.TRADE_START_MINUTE and not self.Portfolio[symbol].Invested:
                or_high = self.or_high.get(symbol, 0)
                or_low = self.or_low.get(symbol, 0)
                
                # Common margin check logic (using dynamic size)
                portfolio_value = self.Portfolio.TotalPortfolioValue
                approx_order_value = current_position_size * portfolio_value
                approx_required_margin = approx_order_value * 0.5  # Assuming 50% initial margin for 2x leverage
                
                if self.Portfolio.MarginRemaining >= approx_required_margin:
                    if bar.Close > or_high:
                        self.SetHoldings(symbol, current_position_size) 
                        self.entry_prices[symbol] = bar.Close
                        self.Debug(f"Buy {symbol} at {bar.Close} on breakout")
                    elif bar.Close < or_low:
                        self.SetHoldings(symbol, -current_position_size) 
                        self.entry_prices[symbol] = bar.Close
                        self.Debug(f"Short {symbol} at {bar.Close} on breakdown")
            
            # Manage exits: stop-loss and take-profit based on position direction
            if self.Portfolio[symbol].Invested:
                entry = self.entry_prices.get(symbol, 0)
                if self.Portfolio[symbol].IsLong:
                    if bar.Close <= entry * (1 - self.STOP_LOSS_PCT):  # Stop-loss
                        self.Liquidate(symbol)
                        self.Debug(f"Stop-loss {symbol} at {bar.Close}")
                    elif bar.Close >= entry * (1 + self.TAKE_PROFIT_PCT):  # Take-profit
                        self.Liquidate(symbol)
                        self.Debug(f"Take-profit {symbol} at {bar.Close}")
                elif self.Portfolio[symbol].IsShort:
                    if bar.Close >= entry * (1 + self.STOP_LOSS_PCT):  # Stop-loss
                        self.Liquidate(symbol)
                        self.Debug(f"Stop-loss {symbol} at {bar.Close}")
                    elif bar.Close <= entry * (1 - self.TAKE_PROFIT_PCT):  # Take-profit
                        self.Liquidate(symbol)
                        self.Debug(f"Take-profit {symbol} at {bar.Close}")
    
    def LiquidateAll(self):
        self.Liquidate()
        self.Debug("End-of-day liquidation")
        # Reset for next day
        self.or_high.clear()
        self.or_low.clear()
        self.entry_prices.clear()