#region imports
from AlgorithmImports import *
from collections import deque
#endregion

class LeapOptionsScreenerStrategy(QCAlgorithm):
    """
    Strategy that screens for stocks based on volatility and options data,
    then purchases long-term call options (LEAPs) on them.
    12/4/2025 - minor tweaks to mx stock price, imp vol, stop loss and vix threshold to resemble the 2 year backtest that did 30%. 
    also added ETFs. also rememeber that every 7 days is a weekend on screening and no trades will trigger.
    SPX- YTD up 12.5% as of 12/4/2025
    spx - 2024 perf - 24.01 SPX 2023 Annaul Perf - 24.73%
    so if im starting in 1//1/2023 i need to beat ...70% return and 20.5% drawdown in spring 2025... so i can hit 50% return and 10% drawdown or 
    other combination.
    """

    def Initialize(self):
        """Initialize the algorithm."""
        self.SetStartDate(2022, 11, 1)
        self.SetEndDate(2025, 12, 1)
        self.SetCash(100000)
        self.SetWarmup(timedelta(days=30))

        # === Strategy Parameters ===
        self.MIN_STOCK_PRICE = 2.50
        self.MAX_STOCK_PRICE = 350.0
        self.MIN_IMPLIED_VOLATILITY = 0.50   # 100%
        self.STOP_LOSS_PERCENT = -0.10  # Stop Loss at 200% loss
        self.TAKE_PROFIT_PERCENT = 0.33  # Take Profit at 300% gain
        self.POSITION_ALLOCATION = 0.02 # 2% per new position
        self.SPY_ALLOCATION = 0.80
        self.XLU_ALLOCATION = 0.20
        self.MAX_PORTFOLIO_ALLOCATION = 1.6 # Max 85% of portfolio in stocks but my PM always under 40%
        self.VIX_THRESHOLD = 20.50 # VIX level to define a BEAR market regime
        
        
        # === Internal State ===
        self.trade_dates = {}
        self.last_week_iv = {} # To store the previous week's IV for comparison
        self.option_symbols = {} # To track equity -> option symbol mapping
        self.last_screening_date = None
        self.screening_frequency_days = 14  # Every 2 weeks (change to 21 for 3 weeks)
        self.vix = None # To hold the VIX symbol
        self.market_regime = 'BULL' # Start with a BULL assumption

        # Set universe settings
        self.UniverseSettings.Asynchronous = True
        self.UniverseSettings.Resolution = Resolution.Daily
        # Add the universe selection function. The timing of the screening is handled by a scheduled event.
        self.AddUniverse(self.UniverseSelectionFunction)
        
        # Add VIX data to use as a market regime filter
        # TEMPLATE NOTE: VIX must be added as an Index, not an Equity, to ensure data is loaded correctly.
        self.vix = self.AddIndex("VIX", Resolution.Daily).Symbol
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.xlu = self.AddEquity("XLU", Resolution.Daily).Symbol
    def UniverseSelectionFunction(self, coarse: List[CoarseFundamental]) -> List[Symbol]:
        """
        Select top 50 most liquid stocks by dollar volume.
        This runs once a week on Monday morning.
        """
        self.Log(f"Coarse selection running at {self.Time}...")
        # Filter for price, fundamental data, and minimum volume
        filtered = [
            c for c in coarse # Remove 'HasFundamentalData' to include ETFs
            if self.MIN_STOCK_PRICE < c.Price < self.MAX_STOCK_PRICE and # Price filter
               c.Volume > 500000 # Liquidity filter
        ]

        # Sort by dollar volume and take the top 50
        top_50 = sorted(filtered, key=lambda c: c.DollarVolume, reverse=True)[:50]
        
        self.Log(f"Selected top {len(top_50)} liquid securities for screening.")
        return [c.Symbol for c in top_50]

    def OnData(self, data):
        """Main trading logic - executed on data events."""
        # TEMPLATE NOTE: The main logic is called from OnData, not a scheduled event.
        # This is CRITICAL because it provides the 'data' slice, which is the only
        # way to access fresh OptionChains data for IV calculation.
        if self.IsWarmingUp:
            return
        
        self.WeeklyScreeningAndTrading(data)
        #Maintain base spy and xlu pos    
        self.SetHoldings(self.spy, self.SPY_ALLOCATION)
        self.SetHoldings(self.xlu, self.XLU_ALLOCATION)
    def WeeklyScreeningAndTrading(self, data):
        """
        Runs periodically (every `screening_frequency_days`) to rank universe stocks
        and manage a portfolio of the top 10.
        - First run: Ranks by absolute IV (or price as fallback).
        - Subsequent runs: Ranks by IV % increase from the previous screening.
        """
        # First, manage existing positions with the current data slice
        self.ManagePositions(data)

        # Only run if enough days have passed since last screening
        if self.last_screening_date is not None:
            days_since = (self.Time.date() - self.last_screening_date).days
            if days_since < self.screening_frequency_days:
                return  # Skip this screening
        
        # --- VIX Market Regime Filter ---
        vix_security = self.Securities[self.vix]
        if not vix_security.HasData:
            self.Log("VIX data not ready yet. Skipping screening cycle.")
            return

        vix_price = vix_security.Price
        if vix_price > self.VIX_THRESHOLD:
            self.market_regime = 'BEAR'
        else:
            self.market_regime = 'BULL'

        if self.market_regime == 'BEAR':
            self.Log(f"VIX is at {vix_price:.2f} (>{self.VIX_THRESHOLD}). Market regime is BEAR. Pausing new trades.")
            return # Stop looking for new trades in a bear market
        '''
        self.last_screening_date = self.Time.date()
        self.Log(f"--- Starting Weekly Screening Trading on {self.Time.date()} ---")
        self.Log("=" * 80)
        self.Log("IV DIAGNOSTIC LOG")
        self.Log("=" * 80)
        '''

        # Step 1: Get current IV for all active securities and filter for a high IV pool
        # CORRECT: Use the helper function to get true Implied Volatility from option chains
        current_iv_data = {}
        for sec in self.ActiveSecurities.Values:
            # Skip securities that don't have data or are the VIX index itself
            # TEMPLATE NOTE: Must filter for Equity or ETF type. self.ActiveSecurities contains stocks, ETFs, and their options.
            if sec.Type != SecurityType.Equity:
                continue
                
            if not sec.HasData or sec.Symbol == self.vix:
                continue
                
            # Use our new helper function to get true IV
            iv = self.GetImpliedVolatility(sec.Symbol, data)
            
            # Only consider valid IVs
            if iv > 0:
                current_iv_data[sec.Symbol] = iv

        high_iv_pool = {s: iv for s, iv in current_iv_data.items() if iv > self.MIN_IMPLIED_VOLATILITY}
        self.Log(f"Found {len(high_iv_pool)} stocks with IV > {self.MIN_IMPLIED_VOLATILITY:.0%}. Now checking for bullish sentiment.")

        # Step 2: Filter the high IV pool for stocks with a bullish call/put volume ratio
        bullish_sentiment_pool = {}
        for symbol, iv in high_iv_pool.items():
            # TEMPLATE NOTE: Always check for the chain's existence in the current data slice.
            if not data.OptionChains or symbol not in data.OptionChains:
                continue
            try:
                option_chain = data.OptionChains[symbol] # Access with bracket notation.
            except:
                continue
            
            if not option_chain or len(option_chain) == 0:
                continue
            
            total_call_volume = sum(c.Volume for c in option_chain if c.Right == OptionRight.Call)
            total_put_volume = sum(c.Volume for c in option_chain if c.Right == OptionRight.Put)

            if total_put_volume == 0:
                # If put volume is zero, any call volume makes the ratio effectively infinite (very bullish)
                call_put_ratio = float('inf') if total_call_volume > 0 else 0
            else:
                call_put_ratio = total_call_volume / total_put_volume

            if call_put_ratio >= 1.10:
                bullish_sentiment_pool[symbol] = iv
                self.Log(f"  -> {symbol.Value} passed sentiment check. C/P Ratio: {call_put_ratio:.2f} ({total_call_volume}/{total_put_volume})")

        self.Log(f"Found {len(bullish_sentiment_pool)} stocks with bullish sentiment. Now checking for IV increase.")

        # If no stocks passed the sentiment filter, fall back to the original high IV pool
        if not bullish_sentiment_pool:
            self.Log("No stocks passed sentiment filter. Falling back to high IV pool.")
            bullish_sentiment_pool = high_iv_pool

        # Step 3: Rank candidates. If it's the first run, rank by absolute IV. Otherwise, rank by IV increase.
        top_sorted = []
        if not self.last_week_iv:
            self.Log("First screening run. Ranking candidates by absolute Implied Volatility.")
            top_sorted = sorted(bullish_sentiment_pool.items(), key=lambda item: item[1], reverse=True)
        else:
            # Calculate the IV increase since the last screening
            iv_increase_candidates = {}
            for symbol, current_iv in bullish_sentiment_pool.items():
                previous_iv = self.last_week_iv.get(symbol)
                if previous_iv and previous_iv > 0:
                    iv_increase = (current_iv - previous_iv) / previous_iv
                    iv_increase_candidates[symbol] = iv_increase
            
            if iv_increase_candidates:
                self.Log("Ranking candidates by percentage increase in Implied Volatility.")
                top_sorted = sorted(iv_increase_candidates.items(), key=lambda item: item[1], reverse=True)

        # Log the top 5 for debugging
        log_metric = "IV" if not self.last_week_iv else "IV Inc"
        top_5_log = [f"{s.Value}: {val:.2%}" for s, val in top_sorted[:5]]
        #self.Log(f"Top 5 candidates by {log_metric}: {', '.join(top_5_log)}")

        # Add a check for an empty list of candidates
        if not top_sorted:
            self.Log("No candidates passed all filters. Skipping this screening cycle.")
            # Still save IV data for the next cycle
            self.last_week_iv = current_iv_data
            return

        # Select the top 25 symbols to trade
        top_25_symbols = [symbol for symbol, _ in top_sorted[:25]]
        
        # --- "Add-Only" Portfolio Management ---

        # Iterate through ranked candidates and add new positions if conditions are met
        for symbol in top_25_symbols:
            # Recalculate current allocation before each trade decision for accuracy
            current_stock_allocation = sum(
                h.HoldingsValue for h in self.Portfolio.Values if h.Invested and h.Type == SecurityType.Equity
            ) / self.Portfolio.TotalPortfolioValue

            # Check if we have room for a new position under the 85% cap
            if current_stock_allocation + self.POSITION_ALLOCATION > self.MAX_PORTFOLIO_ALLOCATION:
                self.Log(f"Max portfolio allocation of {self.MAX_PORTFOLIO_ALLOCATION:.0%} reached. No more positions will be added this cycle.")
                break

            if self.Portfolio[symbol].Invested:
                self.Log(f"[SKIP] {symbol.Value}: Already invested.")
                continue
            
            self.Log(f"[BUY] {symbol.Value}: Current allocation {current_stock_allocation:.2%}. Adding new {self.POSITION_ALLOCATION:.0%} position.")
            self.SetHoldings(symbol, self.POSITION_ALLOCATION)
            self.trade_dates[symbol] = self.Time.date()
        
        
        # Save current IV data for next week's comparison
        self.last_week_iv = current_iv_data

        invested = len([h for h in self.Portfolio.Values if h.Invested])
        self.Log(f"[POSITIONS] Currently invested in: {invested} positions.")

    def GetImpliedVolatility(self, symbol: Symbol, data) -> float:
        """
        Retrieves the Implied Volatility (IV) for a given stock symbol by looking at
        At-The-Money (ATM) options expiring in roughly 30-45 days.
        """
        
        # TEMPLATE NOTE: Must use the 'data' slice to access fresh option chains.
        if not data.OptionChains or symbol not in data.OptionChains:
            self.Log(f"  - {symbol.Value}: No option chain in data slice.")
            return 0.0
        
        option_chain = data.OptionChains[symbol]
        
        if not option_chain or len(option_chain) == 0:
            self.Log(f"  - {symbol.Value}: Option chain is empty in data slice.")
            return 0.0
        
        # 2. Filter for Call options only
        calls = [x for x in option_chain if x.Right == OptionRight.Call]
        if not calls:
            self.Log(f"  - {symbol.Value}: No call contracts found in the chain.")
            return 0.0
        
        # 3. Select contracts closest to 30-45 days out
        target_date = self.Time + timedelta(days=30)
        near_term_contracts = [
            x for x in calls
            if 25 <= (x.Expiry - self.Time).days <= 45
        ]
        
        if not near_term_contracts:
            # Fallback: grab closest to 30 days
            near_term_contracts = sorted(calls, key=lambda x: abs((x.Expiry - target_date).days))[:5]
        
        if not near_term_contracts:
            self.Log(f"  - {symbol.Value}: No near-term contracts found.")
            return 0.0
        
        # 4. Select ATM contract
        underlying_price = self.Securities[symbol].Price
        if underlying_price <= 0:
            self.Log(f"  - {symbol.Value}: Underlying price is zero or invalid.")
            return 0.0
        
        atm_contract = min(near_term_contracts, key=lambda x: abs(x.Strike - underlying_price))
        
        # 5. Return IV
        iv = atm_contract.ImpliedVolatility

        # Add this line BEFORE the return
        if iv == 0.0:
            self.Log(f"⚠️ {symbol.Value}: IV is 0.0 - Greeks not calculating")
        else:
            self.Log(f"✓ {symbol.Value}: IV={iv:.2%}")

        return max(0.0, iv)

    def OnSecuritiesChanged(self, changes: SecurityChanges):
        """Minimal - only clean up data for removed securities."""
        # Add option subscriptions for new equities
        for security in changes.AddedSecurities:
            if security.Type == SecurityType.Equity:
                self.Log(f"New stock in universe: {security.Symbol.Value}. Subscribing to options.")
                option = self.AddOption(security.Symbol, resolution=Resolution.Hour)
                option.SetFilter(lambda u: u.IncludeWeeklys().Strikes(-3, 3).Expiration(0, 45))
                # TEMPLATE NOTE: Both SetPricingModel() and EnableGreekApproximation are required
                # For American-style options (equities), BlackScholes is not suitable. Use a model that supports early exercise.
                option.PriceModel = OptionPriceModels.CrankNicolsonFD()
                option.EnableGreekApproximation = True
                self.option_symbols[security.Symbol] = option.Symbol

        # Remove subscriptions for equities that have left the universe
        for security in changes.RemovedSecurities:
            if security.Type == SecurityType.Equity and security.Symbol in self.option_symbols:
                self.Log(f"Stock left universe: {security.Symbol.Value}. Unsubscribing from options.")
                option_symbol_to_remove = self.option_symbols.pop(security.Symbol)
                self.RemoveSecurity(option_symbol_to_remove)

            # Clean up any historical data
            if security.Symbol in self.last_week_iv:
                del self.last_week_iv[security.Symbol]
                
    def ManagePositions(self, data):
        """Scheduled function to manage open positions (e.g., exit)."""
        for holding in self.Portfolio.Values:
            
            if not holding.Invested or holding.Type != SecurityType.Equity or holding.Symbol in [self.spy, self.xlu]:
                continue

            symbol = holding.Symbol
            unrealized_profit_percent = holding.UnrealizedProfitPercent
            '''
            # 3. IV Stop Loss: Exit if the reason for entry (high IV) is gone
            current_iv = self.GetImpliedVolatility(symbol, data)
            if 0 < current_iv < 45:
                self.Log(f"IV STOP LOSS for {symbol.Value}. Current IV {current_iv:.2%} < Threshold {self.MIN_IMPLIED_VOLATILITY:.2%}. Liquidating.")
                self.Liquidate(symbol)
                if symbol in self.trade_dates: del self.trade_dates[symbol]
                continue # Move to next holding
            '''        
            # 1. Stop Loss
            if unrealized_profit_percent <= self.STOP_LOSS_PERCENT:
                self.Log(f"STOP LOSS TRIGGERED for {symbol.Value}. Unrealized P&L: {unrealized_profit_percent:.2%}. Liquidating.")
                self.Liquidate(symbol)
                if symbol in self.trade_dates: del self.trade_dates[symbol]
                continue # Move to next holding

            # 2. Take Profit
            if unrealized_profit_percent >= self.TAKE_PROFIT_PERCENT:
                self.Log(f"TAKE PROFIT TRIGGERED for {symbol.Value}. Unrealized P&L: {unrealized_profit_percent:.2%}. Liquidating.")
                self.Liquidate(symbol)
                if symbol in self.trade_dates: del self.trade_dates[symbol]
                continue # Move to next holding\

    def OnOrderEvent(self, orderEvent: OrderEvent):
        """Handle order events for logging."""
        if orderEvent.Status != OrderStatus.Filled:
            return

        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        self.Log(f"FILLED: {order.Direction} {order.Quantity} of {order.Symbol.Value} at ${orderEvent.FillPrice:.2f}")

    def OnEndOfAlgorithm(self):
        """Final logging at the end of the algorithm."""
        self.Log("LEAP Options Screener Strategy Ended")
        self.Log(f"Final Portfolio Value: ${self.Portfolio.TotalPortfolioValue:,.2f}")