from AlgorithmImports import *

"""
Opening Range Breakout Universe 

An implementation of the ORB strategy 
Shared by: Quantconnect https://www.quantconnect.com/research/18444/opening-range-breakout-for-stocks-in-play/p1
From research by: Zarattini, Barbon and Aziz.  https://papers.ssrn.com/sol3/papers.cfm?abstract_id=4729284

---

This algorithm implements a classic opening range breakout strategy across a universe of liquid US equities.
The strategy identifies stocks with high relative volume and ATR, then places directional trades based on
breakouts from the opening range (first N minutes of trading).

Key Features:
- Dynamic universe selection based on dollar volume
- Opening range calculation using custom consolidators
- Risk management with stop-loss orders
- Daily liquidation before market close
- Zero commission fee model for backtesting

Strategy Logic:
1. Select top liquid stocks each month
2. Calculate opening range for each stock (default 5 minutes)
3. Identify stocks with high relative volume and ATR
4. Place breakout trades above/below opening range
5. Use stop-loss orders for risk management
6. Liquidate all positions before market close
"""


class OpeningRangeBreakoutUniverseAlgorithm(QCAlgorithm):

    def initialize(self):
        """
        Initialize the algorithm with parameters, universe, and scheduling.
        
        Sets up the trading environment including:
        - Backtest period and capital allocation
        - Zero commission fee model
        - Risk management parameters
        - Universe selection criteria
        - Daily liquidation schedule
        - Indicator warm-up period
        """
        # Backtest configuration
        self.set_start_date(2016, 1, 1)
        # self.set_end_date(2017, 1, 1)  # Uncomment to set end date
        self.set_cash(10_000_000)  # Starting capital: $10M

        # Apply zero commission fee model for accurate backtesting
        self.set_security_initializer(lambda security: security.set_fee_model(ConstantFeeModel(0)))

        # Core strategy parameters
        self.max_positions = 20         # Maximum number of concurrent positions
        self.risk = 0.01               # Risk per position as fraction of portfolio (1%)
        self.entry_gap = 0.1           # Entry buffer as fraction of ATR from breakout level
        
        # Universe and indicator parameters (configurable via algorithm parameters)
        self._universe_size = self.get_parameter("universeSize", 1000)  # Number of stocks in universe
        self._atr_threshold = 0.5      # Minimum ATR value for stock selection
        self._indicator_period = 14    # Period for ATR and volume SMA calculations (days)
        self._opening_range_minutes = self.get_parameter("openingRangeMinutes", 5)  # Opening range duration
        self._leverage = 4             # Maximum leverage for universe securities
        
        # Dictionary to store SymbolData objects for each security
        self._symbol_data_by_symbol = {}

        # Add SPY as benchmark and for scheduling reference
        self._spy = self.add_equity("SPY").symbol

        # Configure universe selection for liquid US equities
        self.universe_settings.leverage = self._leverage
        # Note: Asynchronous universe updates disabled to maintain correct time synchronization
        # self.universe_settings.asynchronous = True  # won't work correct since self.time will be wrong
        self.last_month = None  # Track month changes for universe updates
        self._universe = self.add_universe(self.filter_universe)

        # Schedule daily position liquidation 1 minute before market close
        # This ensures all positions are closed daily as per strategy design
        self.schedule.on(
            self.date_rules.every_day(self._spy),
            self.time_rules.before_market_close(self._spy, 1),
            self.liquidate
        )

        # Warm-up indicators for 2x the indicator period to ensure stability
        self.set_warm_up(timedelta(days=2 * self._indicator_period))

    def filter_universe(self, fundamentals):
        """
        Filter universe to select the most liquid stocks for trading.
        
        Universe selection criteria:
        - Price > $5 (excludes penny stocks)
        - Excludes SPY benchmark
        - Sorted by dollar volume (liquidity)
        - Limited to top N stocks (configurable via universe_size parameter)
        
        Updates monthly to reduce computational overhead while maintaining
        exposure to the most liquid securities.
        
        Args:
            fundamentals: List of fundamental data for all available securities
            
        Returns:
            List of Symbol objects for selected securities, or Universe.UNCHANGED
        """
        # Only update universe selection on the first day of each month
        # This reduces computational load while maintaining exposure to liquid stocks
        if self.time.month == self.last_month:
            return Universe.UNCHANGED

        self.last_month = self.time.month
        
        # Filter and sort securities by dollar volume (liquidity measure)
        # Exclude penny stocks (price < $5) and SPY benchmark
        return [
            f.symbol for f in sorted(
                [f for f in fundamentals if f.price > 5 and f.symbol != self._spy],
                key=lambda f: f.dollar_volume,
                reverse=True  # Highest dollar volume first
            )[:self._universe_size]  # Limit to top N securities
        ]

    def on_securities_changed(self, changes):
        """
        Handle universe changes by initializing SymbolData for new securities.
        
        Creates SymbolData objects for each new security that enters the universe.
        Each SymbolData object manages:
        - Opening range calculation via consolidators
        - ATR and volume indicators
        - Trade entry and exit logic
        - Order management
        
        Args:
            changes: SecurityChanges object containing added/removed securities
        """
        # Initialize SymbolData objects for new securities entering the universe
        for security in changes.added_securities:
            self._symbol_data_by_symbol[security.symbol] = SymbolData(
                self, 
                security, 
                self._opening_range_minutes,
                self._indicator_period
            )

    def on_data(self, slice):
        """
        Main trading logic executed on each data slice.
        
        Strategy execution flow:
        1. Check if opening range period has completed (9:30 AM + opening_range_minutes)
        2. Filter stocks based on selection criteria:
           - Active price > 0
           - In current universe
           - Relative volume > 1 (above average)
           - ATR > threshold (sufficient volatility)
        3. Rank by relative volume and limit to max_positions
        4. Scan each selected stock for breakout opportunities
        
        Args:
            slice: Data slice containing market data for all securities
        """
        # Skip processing during warm-up period or outside trading window
        # Only process after opening range completes (9:30 AM + opening_range_minutes)
        if self.is_warming_up or not (self.time.hour == 9 and self.time.minute == 30 + self._opening_range_minutes):
            return

        # Filter stocks meeting selection criteria and rank by relative volume
        filtered = sorted(
            [
                self._symbol_data_by_symbol[s] for s in self.active_securities.keys
                if self.active_securities[s].price > 0  # Active price check
                and s in self._universe.selected  # In current universe (Note: condition may be redundant)
                and self._symbol_data_by_symbol[s].relative_volume > 1  # Above average volume
                and self._symbol_data_by_symbol[s].ATR.current.value > self._atr_threshold  # Sufficient volatility
            ],
            key=lambda x: x.relative_volume,  # Sort by relative volume (liquidity/interest)
            reverse=True  # Highest relative volume first
        )[:self.max_positions]  # Limit to maximum number of positions

        # Scan each selected stock for breakout trading opportunities
        for symbolData in filtered:
            symbolData.scan()

    def on_order_event(self, orderEvent):
        """
        Handle order events for trade management.
        
        Processes filled orders to trigger follow-up actions:
        - Entry orders trigger stop-loss order placement
        - Maintains order state in corresponding SymbolData objects
        
        Args:
            orderEvent: OrderEvent containing order status and details
        """
        # Only process filled orders
        if orderEvent.status != OrderStatus.FILLED:
            return

        # Forward order events to corresponding SymbolData for trade management
        if orderEvent.symbol in self._symbol_data_by_symbol:
            self._symbol_data_by_symbol[orderEvent.symbol].on_order_event(orderEvent)


"""
SymbolData Class

Manages individual security data and trading logic for the opening range breakout strategy.

Key responsibilities:
- Calculate opening range using time-based consolidators
- Track relative volume using simple moving average
- Maintain ATR indicator for volatility measurement
- Execute breakout trading logic with risk management
- Handle order lifecycle and stop-loss placement

The class encapsulates all security-specific state and logic, allowing the main algorithm
to focus on universe selection and coordination.
"""


class SymbolData:
    def __init__(self, algorithm: QCAlgorithm, security, openingRangeMinutes, indicatorPeriod):
        """
        Initialize SymbolData with indicators and consolidators.
        
        Sets up the infrastructure for tracking:
        - Opening range bars via time-based consolidation
        - ATR indicator for volatility measurement
        - Volume SMA for relative volume calculation
        - Order tickets for trade management
        
        Args:
            algorithm: Reference to main QCAlgorithm instance
            security: Security object for this symbol
            openingRangeMinutes: Duration of opening range in minutes
            indicatorPeriod: Period for ATR and volume SMA calculations
        """
        # Algorithm and security references
        self.algorithm = algorithm
        self.security = security
        
        # Opening range and volume tracking
        self.opening_bar = None      # Current day's opening range bar
        self.relative_volume = 0     # Current relative volume ratio
        
        # Technical indicators
        self.ATR = algorithm.ATR(security.symbol, indicatorPeriod, resolution=Resolution.DAILY)
        self.volumeSMA = SimpleMovingAverage(indicatorPeriod)  # For relative volume calculation
        
        # Order management
        self.stop_loss_price = None   # Calculated stop loss price
        self.entry_ticket = None      # Entry order ticket
        self.stop_loss_ticket = None  # Stop loss order ticket

        # Time-based consolidator for opening range calculation
        # Consolidates bars for the specified opening range duration
        self.consolidator = algorithm.consolidate(
            security.symbol, 
            TimeSpan.from_minutes(openingRangeMinutes), 
            self.consolidation_handler
        )

    def consolidation_handler(self, bar):
        """
        Handle consolidated opening range bars.
        
        Processes each opening range bar to:
        1. Update relative volume calculation
        2. Store the opening bar for breakout analysis
        3. Maintain volume SMA for relative volume tracking
        
        Only processes one opening bar per trading day to avoid multiple
        entries on the same day.
        
        Args:
            bar: Consolidated bar representing the opening range
        """
        # Skip if we already have an opening bar for this trading day
        # This prevents multiple entries on the same day
        if self.opening_bar and self.opening_bar.time.date() == bar.time.date():
            return

        # Calculate relative volume (current volume / average volume)
        # Relative volume > 1 indicates above-average trading interest
        self.relative_volume = (
            bar.volume / self.volumeSMA.current.value 
            if self.volumeSMA.is_ready and self.volumeSMA.current.value > 0 
            else 0
        )
        
        # Update volume SMA with current bar's volume
        self.volumeSMA.update(bar.end_time, bar.volume)
        
        # Store opening bar for breakout analysis
        self.opening_bar = bar

    def scan(self):
        """
        Scan for opening range breakout opportunities.
        
        Implements the core breakout logic:
        1. Analyze opening bar direction (bullish/bearish)
        2. Place breakout trades based on opening range high/low
        3. Set appropriate stop-loss levels using ATR-based buffer
        
        Bullish breakout: Enter long above opening range high
        Bearish breakout: Enter short below opening range low
        
        Only executes if opening bar data is available.
        """
        # Skip if no opening bar data available
        if not self.opening_bar:
            return

        # Determine opening bar direction and place appropriate breakout trade
        if self.opening_bar.close > self.opening_bar.open:
            # Bullish opening bar - place long breakout above opening range high
            # Stop loss set below entry with ATR-based buffer
            self.place_trade(
                self.opening_bar.high,  # Entry price (breakout above high)
                self.opening_bar.high - self.algorithm.entry_gap * self.ATR.current.value  # Stop loss
            )
        elif self.opening_bar.close < self.opening_bar.open:
            # Bearish opening bar - place short breakout below opening range low
            # Stop loss set above entry with ATR-based buffer
            self.place_trade(
                self.opening_bar.low,   # Entry price (breakout below low)
                self.opening_bar.low + self.algorithm.entry_gap * self.ATR.current.value   # Stop loss
            )

    def place_trade(self, entryPrice, stopPrice):
        """
        Calculate position size and place breakout trade with risk management.
        
        Position sizing methodology:
        1. Calculate risk per position based on portfolio allocation
        2. Determine quantity based on entry-to-stop distance
        3. Apply portfolio allocation limits
        4. Place stop market entry order
        
        Risk management:
        - Fixed risk per position (algorithm.risk / max_positions)
        - Stop loss calculated using ATR-based buffer
        - Position size limited by portfolio allocation constraints
        
        Args:
            entryPrice: Price level for trade entry (breakout level)
            stopPrice: Stop loss price level
        """
        # Calculate risk allocation per position
        # Total portfolio risk distributed equally across max positions
        risk_per_position = (
            self.algorithm.portfolio.total_portfolio_value * self.algorithm.risk
        ) / self.algorithm.max_positions
        
        # Calculate position quantity based on risk and stop distance
        # Quantity = Risk Amount / (Entry Price - Stop Price)
        quantity = int(risk_per_position / (entryPrice - stopPrice))

        # Apply portfolio allocation limits using QuantConnect's built-in calculator
        # This ensures we don't exceed reasonable position sizes
        quantity_limit = self.algorithm.calculate_order_quantity(
            self.security.symbol, 
            1 / self.algorithm.max_positions  # Equal allocation across max positions
        )
        
        # Determine position direction and apply quantity limits
        if quantity > 0:
            sign = 1      # Long position
        elif quantity < 0:
            sign = -1     # Short position
        else:
            sign = 0      # No position (should not occur with valid inputs)
            
        # Apply the smaller of calculated quantity or allocation limit
        quantity = int(min(abs(quantity), quantity_limit) * sign)

        # Place trade if quantity is valid
        if quantity != 0:
            # Store stop loss price for later use
            self.stop_loss_price = stopPrice
            
            # Place stop market entry order
            # This will trigger when price breaks through the entry level
            self.entry_ticket = self.algorithm.stop_market_order(
                self.security.symbol, 
                quantity, 
                entryPrice, 
                "Entry"  # Order tag for identification
            )

    def on_order_event(self, orderEvent):
        """
        Handle order events for trade management.
        
        Processes filled entry orders to automatically place corresponding
        stop-loss orders. This ensures every position has risk management
        protection immediately upon entry.
        
        Trade management flow:
        1. Entry order fills
        2. Stop-loss order placed immediately
        3. Position is now protected with automatic exit
        
        Args:
            orderEvent: OrderEvent containing order details and status
        """
        # Check if this is our entry order being filled
        if self.entry_ticket and orderEvent.order_id == self.entry_ticket.order_id:
            # Place stop loss order immediately upon entry fill
            # Quantity is negative of entry (opposite direction for exit)
            self.stop_loss_ticket = self.algorithm.stop_market_order(
                self.security.symbol, 
                -self.entry_ticket.quantity,  # Opposite direction of entry
                self.stop_loss_price,         # Pre-calculated stop price
                "Stop Loss"                   # Order tag for identification
            )