"""

SPX Opening Range Breakout (ORB) Options Strategy
----------------------------------
Quantish.io | Quantish.substack.com


This algorithm implements a 60-minute opening range breakout strategy on SPX options.
The strategy identifies the price range during the first hour of trading (9:30-10:30 AM)
and trades breakouts above or below this range using credit spreads.

Strategy Logic:
- Upside Breakout: Sell bull put spread with short strike at opening range low
- Downside Breakout: Sell bear call spread with short strike at opening range high
- Uses $15 wide spreads, holds to expiration
- Filters: Opening range must be >= 0.2% of SPX price), no FOMC days, no trades after 12pm
"""


from AlgorithmImports import *

class SPXORBCreditSpreads(QCAlgorithm):

    def initialize(self):

        # Initialize backtest                
        self.set_start_date(2022, 5, 1)                
        self.set_cash(10000)
        self.set_security_initializer(lambda security: security.set_market_price(self.get_last_known_price(security)))
        
        # Setup SPX index and options
        index = self.add_index("SPX", Resolution.MINUTE).symbol
        option = self.add_index_option(index, "SPXW", Resolution.MINUTE)
        option.set_filter(lambda x: x.include_weeklys().strikes(-50, 10).expiration(0, 0))
                
        self.spxw = option.symbol
        self.spx = index
        self.tickets = []
        self.spread_map = {}  # symbol -> spread_id mapping
                
        # Entry Parameters, with default values
        self.spread_width     = int(self.get_parameter("spread_width", 15))
        self.trade_bearish    = 1 == int(self.get_parameter("trade_bearish", 1))
        self.trade_bullish    = 1 == int(self.get_parameter("trade_bullish", 1))        
        self.skip_fridays     = 1 == int(self.get_parameter("skip_fridays", 1))
        self.skip_wednesdays  = 1 == int(self.get_parameter("skip_wednesdays", 1))
        self.mon_tue_only     = 1 == int(self.get_parameter("mon_tue_only", 0))

        # Additional Risk management parameters
        self.fixed_order_size    = int(self.get_parameter("fixed_order_size", 0)) 
        self.pct_risk            = float(self.get_parameter("pct_risk", .2)) # used if order size is 0
        self.min_range_threshold = 0.002  # Minimum range threshold (0.2% of underlying price)

        # Opening range variables
        self.opening_range_minutes = 60  # Use 60-minute opening range
        self.opening_range_high = None
        self.opening_range_low = None
        self.range_established = False
        self.breakout_traded = False
        self.daily_reset_done = False
        
        # Get FOMC dates
        self.fomc_dates = self.get_fomc_dates()
        
        self.log("SPX ORB Options Strategy Initialized")

    # Main data handler that orchestrates the entire trading logic flow
    # Processes each data tick to update opening range, check for breakouts, and execute trades
    # Coordinates daily resets, range establishment, and breakout detection
    # -----------------------------------------------------------------------------------------
    def on_data(self, slice: Slice) -> None:
        
        # Early exit conditions
        if not self._should_process_data(slice):
            return
            
        current_time = self.time
        current_price = slice.bars[self.spx].close
        
        # Handle daily resets
        self._handle_daily_resets(current_time)
        
        # Update opening range during establishment period
        self._update_opening_range(current_time, current_price)
        
        # Check if we should trade breakouts
        if not self._can_trade_breakout(current_time):
            return
            
        # Check for and execute breakout trades
        self._check_and_execute_breakouts(slice, current_price)

    
    # Validate whether current data tick should be processed for trading
    # Checks algorithm state, day filters, FOMC exclusions, and data availability
    # Returns false to skip processing if any exclusion criteria are met
    # ---------------------------------------------------------------------------
    def _should_process_data(self, slice: Slice) -> bool:
        
        if self.is_warming_up:
            self.log("Algorithm is warming up, skipping data")
            return False

        ## TEMP CHECK 
        if self.mon_tue_only and (self.time.weekday() not in [0,1]):
            return False


        if self.skip_fridays and self.time.weekday() == 4:
            return False

        if self.skip_wednesdays and self.time.weekday() == 2:
            return False

            
        # Skip FOMC meeting days
        if self.time.strftime("%Y-%m-%d") in self.fomc_dates:
            # self.log(f"FOMC meeting day {self.time.strftime('%Y-%m-%d')}, skipping trading")
            return False
            
        # Check if SPX data is available
        if self.spx not in slice or slice[self.spx] is None:
            # self.log("SPX data not available in slice")
            return False
            
        return True

    # Reset daily trading variables at market open and end of day
    # Ensures fresh state for each trading session by clearing opening range data
    # Manages daily_reset_done flag to prevent multiple resets on same day
    # ---------------------------------------------------------------------------
    def _handle_daily_resets(self, current_time) -> None:
        
        # Reset daily variables at market open
        if current_time.hour == 9 and current_time.minute == 31 and not self.daily_reset_done:
            self.opening_range_high = None
            self.opening_range_low = None
            self.range_established = False
            self.range_is_valid = False
            self.breakout_traded = False
            self.daily_reset_done = True
            # self.log("Daily variables reset for new trading day")
            
        # Reset flag at end of day
        if current_time.hour == 16:
            self.daily_reset_done = False

    # Track and update opening range high/low during establishment period (9:30-10:30 AM)
    # Continuously monitors price action to identify the true range bounds
    # Automatically finalizes range at 10:30 AM and validates minimum size requirements
    # -----------------------------------------------------------------------------------
    def _update_opening_range(self, current_time, current_price: float) -> None:
        
        # Establish opening range during first 60 minutes
        if (current_time.hour == 9 and current_time.minute >= 30) or \
           (current_time.hour == 10 and current_time.minute < 31):
            
            if self.opening_range_high is None:
                self.opening_range_high = current_price
                self.opening_range_low = current_price
                # self.log(f"Starting opening range: {current_price:.2f}")
            else:
                self.opening_range_high = max(self.opening_range_high, current_price)
                self.opening_range_low = min(self.opening_range_low, current_price)
                
        # Mark range as established at 10:30 AM
        if current_time.hour == 10 and current_time.minute == 31 and not self.range_established:
            self._finalize_opening_range(current_price)

    # Complete opening range establishment and validate it meets minimum threshold
    # Ensures range is large enough (0.2% of SPX price) to constitute a valid trading signal
    # Prevents trading on days with insufficient volatility during opening hour
    # --------------------------------------------------------------------------------------
    def _finalize_opening_range(self, current_price: float) -> None:
        """Finalize the opening range and validate it meets minimum requirements"""
        
        self.range_established = True
        
        # Check if range meets minimum threshold
        if self.opening_range_high and self.opening_range_low:
            range_size = self.opening_range_high - self.opening_range_low
            min_range_required = current_price * self.min_range_threshold
            
            if range_size < min_range_required:
                self.log(f"Range too small: {range_size:.2f} < {min_range_required:.2f}")
                self.range_is_valid = False
            else:
                self.range_is_valid = True

        return

    # Determine if conditions allow for breakout trading execution
    # Validates range establishment, time window (before noon), and position status
    # Prevents multiple trades per day and ensures proper timing constraints
    # -----------------------------------------------------------------------------
    def _can_trade_breakout(self, current_time) -> bool:
        
        # Only trade after range is established and before noon
        if not self.range_established or not self.range_is_valid or self.breakout_traded or current_time.hour >= 12:
            return False

        # Return if open position exists
        if any([self.portfolio[x.symbol].invested for x in self.tickets]):
            self.log("Existing position found, skipping new trades")
            return False
            
        return True

    # Monitor for breakout signals and trigger appropriate trade execution
    # Detects when price breaks above opening range high or below opening range low
    # Initiates corresponding credit spread strategy based on breakout direction
    # -----------------------------------------------------------------------------
    def _check_and_execute_breakouts(self, slice: Slice, current_price: float) -> None:
        
        # Check for breakout signals
        upside_breakout = current_price > self.opening_range_high
        downside_breakout = current_price < self.opening_range_low
        
        if upside_breakout and self.trade_bullish:
            # self.log(f"Upside breakout detected: {current_price:.2f} > {self.opening_range_high:.2f}")
            self.current_spread_id = f"upside_{self.time.strftime('%Y%m%d_%H%M%S')}"
            self._execute_upside_breakout_trade(slice)

        elif downside_breakout and self.trade_bearish:
            # self.log(f"Downside breakout detected: {current_price:.2f} < {self.opening_range_low:.2f}")
            self.current_spread_id = f"downside_{self.time.strftime('%Y%m%d_%H%M%S')}"
            self._execute_downside_breakout_trade(slice)
        

    # Execute bull put spread when price breaks above opening range high
    # Sells put spread with short strike at opening range low (expecting support)
    # Uses opening range low as anchor point, assuming it will act as support level
    # -----------------------------------------------------------------------------
    def _execute_upside_breakout_trade(self, slice: Slice) -> None:
        
        current_price = slice.bars[self.spx].close

        # Get option chain
        chain = slice.option_chains.get(self.spxw)
        if not chain:
            self.log("No option chain available for upside breakout")
            return

        # Get the nearest expiry date of the contracts
        expiry = min([x.expiry for x in chain])
        # self.log(f"Using expiry: {expiry}, Chain has {len(chain)} contracts")
        
        # Sell short put spread (bullish breakout)
        puts = sorted([i for i in chain if i.expiry == expiry and i.right == OptionRight.PUT], 
                     key=lambda x: x.strike)
        if len(puts) < 2:
            self.log(f"Insufficient puts available: {len(puts)}")
            return
            
        # self.log(f"Available put strikes: {[p.strike for p in puts[:5]]}...{[p.strike for p in puts[-5:]]}")
        
        # Find puts for short strike - range edge
        target_short_strike = self.opening_range_low
        strike_description = f"Range low (${target_short_strike:.2f})|{current_price:.2f}"
        valid_short_puts = puts

        if not valid_short_puts:
            self.log("No puts available below current price for short strike")
            return
        
        short_put = min(valid_short_puts, key=lambda x: abs(x.strike - target_short_strike))
        # self.log(f"Selected short put strike: {short_put.strike} (target: {strike_description})")

        # Find put closest to $x below short strike
        target_long_strike = short_put.strike - self.spread_width
        valid_puts = [put for put in puts if put.strike <= target_long_strike]
        long_put = max(valid_puts, key=lambda put: put.strike) if valid_puts else None

        if long_put:
            put_spread = OptionStrategies.bull_put_spread(self.spxw, short_put.strike, long_put.strike, expiry)
            order_size = self.calc_position_size(short_put, long_put)
                        
            tag = (f"+P{long_put.strike} -P{short_put.strike} | "
                f"orb_low={self.opening_range_low:.2f} | orb_high={self.opening_range_high:.2f} | "
                f"spx={current_price:.2f}")
            

            self.tickets = self.buy(put_spread, order_size, tag=tag)

            # self.log(f"Upside breakout: Sold put spread {long_put.strike}/{short_put.strike}, size: {order_size}")
            self.breakout_traded = True
        else:
            self.log(f"No long put found ${self.spread_width} below short strike {short_put.strike}")

    # Execute bear call spread when price breaks below opening range low
    # Sells call spread with short strike at opening range high (expecting resistance)
    # Uses opening range high as anchor point, assuming it will act as resistance level
    # ---------------------------------------------------------------------------------
    def _execute_downside_breakout_trade(self, slice: Slice) -> None:
        
        current_price = slice.bars[self.spx].close

        # Get option chain
        chain = slice.option_chains.get(self.spxw)
        if not chain:
            self.log("No option chain available for downside breakout")
            return

        # Get the nearest expiry date of the contracts
        expiry = min([x.expiry for x in chain])
        
        # Sell short call spread (bearish breakout)
        calls = sorted([i for i in chain if i.expiry == expiry and i.right == OptionRight.CALL], 
                      key=lambda x: x.strike)
        if len(calls) < 2:
            self.log(f"Insufficient calls available: {len(calls)}")
            return
            
        # Find calls for short strike - range edge
        target_short_strike = self.opening_range_high
        strike_description = f"Range high ${target_short_strike:.2f}"
        # valid_short_calls = calls
        valid_short_calls = [call for call in calls if call.strike >= target_short_strike]

        # For calls, ensure short strike is above current price (OTM for bearish position)
        if not valid_short_calls:
            self.log("No calls available above current price for short strike")
            return
        
        short_call = min(valid_short_calls, key=lambda x: abs(x.strike - target_short_strike))

        # Find call closest to $x above short strike
        target_long_strike = short_call.strike + self.spread_width
        valid_calls = [call for call in calls if call.strike >= target_long_strike]
        long_call = min(valid_calls, key=lambda call: call.strike) if valid_calls else None
                
        if long_call:
            call_spread = OptionStrategies.bear_call_spread(self.spxw, short_call.strike, long_call.strike, expiry)
            order_size = self.calc_position_size(short_call, long_call)


            # Enhanced tag creation            
            tag = (f"+C{long_call.strike} -C{short_call.strike} | "
                f"orb_low={self.opening_range_low:.2f} | orb_high={self.opening_range_high:.2f} | "
                f"spx={current_price:.2f}")

            self.tickets = self.buy(call_spread, order_size, tag=tag)
            self.breakout_traded = True
        else:
            # self.log(f"No long call found ${self.spread_width} above short strike {short_call.strike}")
            pass

    # Calculate appropriate position size based on risk management parameters
    # Uses either fixed order size or percentage-based sizing depending on configuration
    # Ensures position sizing aligns with available capital and risk tolerance
    # -----------------------------------------------------------------------------------
    def calc_position_size(self, short_contract, long_contract):
        if self.fixed_order_size > 0:
            return self.fixed_order_size
        else:
            buying_power = self.portfolio.total_portfolio_value - self.portfolio.total_holdings_value
            alloc_power = self.pct_risk * buying_power 

            strike_diff = abs(short_contract.strike - long_contract.strike)
            contract_unit = self.securities[short_contract.symbol].contract_unit_of_trade
            margin_req = Math.max((strike_diff * contract_unit * 1), 0)  # required for one
            
            order_size = math.floor(alloc_power / margin_req)
            
            
            return order_size

    # Return comprehensive set of FOMC meeting dates to avoid trading
    # Federal Reserve meetings create unpredictable volatility that disrupts opening range patterns
    # Strategy skips these days to avoid adverse market reactions to monetary policy announcements
    # --------------------------------------------------------------------------------------------
    def get_fomc_dates(self):
        return {
            "2025-06-18", "2025-05-07", "2025-03-19", "2025-01-29",
            "2024-12-18", "2024-11-07", "2024-09-18", "2024-07-31",
            "2024-06-12", "2024-05-01", "2024-03-20", "2024-01-31",
            "2023-12-13", "2023-11-01", "2023-09-20", "2023-07-26",
            "2023-06-14", "2023-05-03", "2023-03-22", "2023-02-01",        
            "2022-12-14", "2022-11-02", "2022-09-21", "2022-07-27",
            "2022-06-15", "2022-05-04", "2022-03-16", "2021-12-15",
            "2021-11-03", "2021-09-22", "2021-07-28", "2021-06-16",
            "2021-04-28", "2021-03-17", "2021-01-27", "2020-12-16",
            "2020-11-05", "2020-09-16", "2020-07-29", "2020-06-10",
            "2020-04-29", "2020-03-31", "2020-03-23", "2020-03-19",
            "2020-03-15", "2020-03-03", "2019-10-30", "2019-09-18",
            "2019-07-31", "2018-12-19", "2018-09-26", "2018-06-13",
            "2018-03-21", "2017-12-13", "2017-06-14", "2017-03-15",
            "2016-12-14", "2015-12-16", "2011-06-22", "2008-12-16",
            "2008-10-29", "2008-10-08", "2008-09-16", "2008-08-05",
            "2008-04-30", "2008-03-18", "2008-03-16", "2008-01-30",
            "2008-01-22", "2007-12-11", "2007-10-31", "2007-09-18",
            "2007-08-17", "2007-08-07", "2007-06-28", "2007-05-09",
            "2007-03-21", "2007-01-31", "2006-12-12", "2006-10-25",
            "2006-09-20", "2006-08-08", "2006-06-29", "2006-05-10",
            "2006-03-28", "2006-01-31", "2005-12-13", "2005-08-09",
            "2005-06-30", "2005-05-03", "2005-03-22", "2005-02-02",
            "2004-12-14", "2004-11-10", "2004-09-21", "2004-08-10",
            "2004-06-30", "2004-05-04", "2004-03-16", "2004-01-28",
            "2003-12-09", "2003-10-28", "2003-09-16", "2003-08-12",
            "2003-06-25", "2003-05-06", "2003-03-18", "2003-01-29",
            "2002-12-10", "2002-11-06", "2002-09-24", "2002-08-13",
            "2002-06-26", "2002-05-07", "2002-03-19", "2002-01-30",
            "2001-12-11", "2001-11-06", "2001-10-02", "2001-09-17",
            "2001-08-21", "2001-06-27", "2001-05-15", "2001-04-18",
            "2001-03-20", "2001-01-31", "2001-01-03", "2000-12-19",
            "2000-11-15", "2000-10-03", "2000-08-22", "2000-06-28",
            "2000-05-16", "2000-03-21", "2000-02-02"
        }

class CustomSlippageModel:
    def get_slippage_approximation(self, asset: Security, order: Order) -> float:
        # We set a maximum 2% slippage that is linearly linked to the ratio of the order size versus the previous bar's volume.
        return asset.price * 0.02 * min(1, order.absolute_quantity / asset.volume)

class CustomSecurityInitializer(BrokerageModelSecurityInitializer):
    def __init__(self, brokerage_model, security_seeder):
        super().__init__(brokerage_model, security_seeder)
    def initialize(self, security):
        super().initialize(security)
        # To set the slippage model to the custom one for all assets entering the universe.
        security.set_slippage_model(CustomSlippageModel())