from AlgorithmImports import *
from datetime import timedelta
import numpy as np

from volatility_utils import VolatilityCalculator
from options_utils import OptionsManager
from trade_manager import TradeManager

"""
This strategy exploits earnings-related IV crush through ATM calendar 
spreads, entering positions when both IV term structure is abnormally 
steep and IV/RV ratio indicates overpricing before announcements. 

By selling front-month options against back-month options at identical 
strikes, it capitalizes on the asymmetric volatility compression between 
expirations, generating profits regardless of moderate price movement 
in the underlying. Primary risk includes gap moves exceeding expected 
range and abnormal post-earnings IV behavior.

Backtest Author  : u/shock_and_awful (reddit)
Strategy Creator : @VolatilityVibes (youtube)
"""

class EarningsVolatilityCalendarSpread(QCAlgorithm):
    def initialize(self):
        """
        Initialize the algorithm with settings, schedules, and variables.
        Sets up the earnings data source, universe settings, and scheduling for precise entry/exit timing.
        """
        self.set_start_date(2024, 6, 1)
        self.set_end_date(2024, 11, 3)
        self.set_cash(1000000)
        self.set_security_initializer(BrokerageModelSecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))

        # Initialize utility classes
        self.volatility = VolatilityCalculator(self)
        self.options = OptionsManager(self)
        self.trade_manager = TradeManager(self)

        # Local Enums for Earning report time. 
        # TODO: Will change later and replace with official LEAN ENUMS
        # Important because we dont know if LEAN expects 1 for BMO 
        self.IS_BMO = 0  # Before Market Open
        self.IS_AMC = 1  # After Market Close

        # Add earnings data source to track upcoming corporate earnings
        self.add_universe(EODHDUpcomingEarnings, self.selection)
        
        # Set universe settings for data resolution and normalization
        self.universe_settings.resolution = Resolution.MINUTE
        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        
        # Initialize variables to track earnings
        self.earnings_schedule = {}  # Will store: {symbol: (date, is_amc)}
        
        # Define threshold values from the technical specification
        self.threshold_volume = 1500000  # Minimum average daily volume requirement
        self.threshold_iv_rv_ratio = 1.25  # Minimum IV30/RV30 ratio requirement
        self.threshold_ts_slope = -0.00406  # Maximum term structure slope requirement
        self.risk_per_trade = 0.01  # 1% of portfolio per trade for position sizing
        
        # Define a symbol for scheduling (needed for market hours reference)
        self.schedule_symbol = Symbol.Create("SPY", SecurityType.EQUITY, Market.USA)
        self.add_equity("SPY", Resolution.MINUTE)
        
        # Schedule entry evaluation 15 minutes before market close
        self.schedule.on(
            self.date_rules.every_day(self.schedule_symbol), 
            self.time_rules.before_market_close(self.schedule_symbol, 15), 
            self.evaluate_entry_conditions
        )
        
        # Schedule position closing 15 minutes after market open
        self.schedule.on(
            self.date_rules.every_day(self.schedule_symbol), 
            self.time_rules.after_market_open(self.schedule_symbol, 15), 
            self.close_positions
        )
    
    def selection(self, earnings: List[EODHDUpcomingEarnings]) -> List[Symbol]:
        """
        Filter stocks with upcoming earnings within the next 4 days.
        
        Args:
            earnings: List of upcoming earnings announcements
            
        Returns:
            List of symbols that meet the selection criteria
        """
        # Look ahead 4 days to handle weekends properly
        selected_symbols = []
        # List of allowed symbols
        allowed_symbols = ["AAPL", "NVDA", "MSFT", "AMZN", "META", "GOOGL", "AVGO",\
                            "GOOG", "TSLA", "JPM", "LLY", "V", "XOM", "UNH", "MA", \
                            "NFLX", "COST", "JNJ", "PG", "WMT", "ABBV", "ADBE"]
        for earning in earnings:
            # Only process allowed symbols
            if earning.symbol.value not in allowed_symbols:
                continue
                
            if earning.report_date <= self.time + timedelta(days=4):
                # Store both the report date and whether it's AMC (1) or BMO (0)
                is_amc = earning.report_time == self.IS_AMC  
                self.earnings_schedule[earning.symbol] = (earning.report_date, is_amc)
                selected_symbols.append(earning.symbol)
        return selected_symbols
    

    def on_securities_changed(self, changes: SecurityChanges) -> None:
        """
        Handle changes to the securities universe.
        Adds options for new securities without liquidating positions for removed securities.
        
        Args:
            changes: Securities added or removed from the universe
        """
        # Add option chains for newly added securities
        for added in [security for security in changes.added_securities if security.type == SecurityType.EQUITY]:
            self.log(f"Added security {added.symbol} to universe")
            
            # Add option contracts for the equity
            option = self.add_option(added.symbol)
            
            # Set filter for option contracts - wider range to ensure finding suitable contracts
            option.set_filter(lambda u: (u.strikes(-2, +2).expiration(0, 60)))

    def evaluate_entry_conditions(self):
        """
        Evaluates entry conditions for calendar spreads 15 minutes before market close.
        This is the core strategy logic that checks all predictor variables against thresholds.
        """
        self.log(f"[3:45 PM] Evaluating entry conditions at {self.time}")
        
        # Track symbols that we've processed and can remove from earnings schedule
        symbols_to_remove = []
        
        # Check securities with earnings tomorrow
        for underlying_symbol, earnings_info in list(self.earnings_schedule.items()):
            earnings_date, is_amc = earnings_info
            
            # Only trade if earnings are tomorrow
            if earnings_date.date() != (self.time + timedelta(days=1)).date():
                continue
                
            self.log(f"[{underlying_symbol} Entry Check] - Earnings on {earnings_date} {'AMC' if is_amc else 'BMO'}")
            
            # Check if underlying is tradable and has valid price
            if not self.options.is_underlying_valid(underlying_symbol):
                # If we've decided not to enter, we can remove this from tracking
                symbols_to_remove.append(underlying_symbol)
                continue
            
            # Get option contracts for possible calendar spread
            atm_strike, expiry_dict = self.options.get_atm_options(underlying_symbol)
            if atm_strike is None or expiry_dict is None:
                symbols_to_remove.append(underlying_symbol)
                continue
                
            # Find suitable expiration pairs for calendar spread
            expiry_pair = self.options.find_expiry_pair(expiry_dict)
            if expiry_pair is None:
                symbols_to_remove.append(underlying_symbol)
                continue
                
            near_term_days, longer_term_days = expiry_pair

            # Check average volume (least expensive calculation first)
            avg_volume = self.volatility.calculate_average_volume(underlying_symbol)
            if avg_volume is None or avg_volume <= self.threshold_volume:
                symbols_to_remove.append(underlying_symbol)
                continue
        
            # Get term structure data
            days_array, iv_array = self.volatility.calculate_term_structure(expiry_dict)
            
            # Calculate term structure slope
            ts_slope_result = self.volatility.calculate_term_structure_slope(days_array, iv_array)
            if ts_slope_result is None or ts_slope_result[0] >= self.threshold_ts_slope:
                symbols_to_remove.append(underlying_symbol)
                continue

            ts_slope, term_structure = ts_slope_result
                            
            # Calculate realized volatility and IV/RV ratio using the already computed term structure
            iv_rv_ratio = self.volatility.calculate_iv_rv_ratio(underlying_symbol, term_structure)
            if iv_rv_ratio is None or iv_rv_ratio <= self.threshold_iv_rv_ratio:
                symbols_to_remove.append(underlying_symbol)
                continue
                
            # Log predictor variables
            self.log(
                f"....Term Slope: {round(ts_slope,5)}, "
                f"IV/RV: {round(iv_rv_ratio,2)}, "
                f"Avg Vol: {round(avg_volume,0)}"
            )

            # All conditions met, select contracts for calendar spread
            self.log(f"....All entry conditions met for {underlying_symbol}")
            
            # Select appropriate option contracts
            near_term_contract, longer_term_contract = self.options.select_option_contracts(
                expiry_dict, near_term_days, longer_term_days, OptionRight.CALL)
                
            if not near_term_contract or not longer_term_contract:
                symbols_to_remove.append(underlying_symbol)
                continue
            
            # Calculate position size
            position_size = self.options.calculate_position_size(
                near_term_contract, longer_term_contract, self.risk_per_trade)
            
            # Create calendar spread trade
            trade_success = self.trade_manager.enter_calendar_spread(
                underlying_symbol, near_term_contract, longer_term_contract, position_size)
            
            # If we didn't enter a trade, add to removal list
            if not trade_success:
                symbols_to_remove.append(underlying_symbol)
        
        # Clean up earnings schedule for symbols we've processed and decided not to trade
        for symbol in symbols_to_remove:
            if symbol in self.earnings_schedule:
                self.log(f"Removing {symbol} from earnings schedule - no trade entered")
                del self.earnings_schedule[symbol]

    def close_positions(self):
        """
        Close positions 15 minutes after market open for stocks that had earnings.
        Handles different timing for BMO (Before Market Open) and AMC (After Market Close) reports.
        """
        self.log(f"[9:45 AM] Checking for open positions to close at {self.time}")
        
        # Process exit logic based on earnings timing
        closed_symbols = self.trade_manager.process_earnings_exit(self.earnings_schedule)
        
        # Remove closed symbols from earnings schedule
        for symbol in closed_symbols:
            if symbol in self.earnings_schedule:
                del self.earnings_schedule[symbol]
                self.log(f"....Removed {symbol} from earnings schedule after position closed")
    
    def on_order_event(self, order_event):
        # Order event handling (if needed)
        pass 

from AlgorithmImports import *

class OptionsManager:
    """
    Handles all options-related functionality including finding ATM options,
    selecting expiration pairs, and building option positions.
    """
    
    def __init__(self, algorithm):
        """
        Initialize with a reference to the main algorithm.
        
        Args:
            algorithm: The main algorithm instance for accessing options chains and other methods
        """
        self.algorithm = algorithm
    
    def get_atm_options(self, symbol):
        """
        Get the ATM option contracts for the symbol.
        
        Args:
            symbol: The underlying security symbol
            
        Returns:
            tuple: (atm_strike, expiry_dict) or (None, None) if not found
        """
        # Get option chain
        option_chain = self.algorithm.option_chain(symbol)
        if option_chain is None or not list(option_chain):
            self.algorithm.log(f"....No option chain available for {symbol}")
            return None, None
                
        underlying_price = self.algorithm.securities[symbol].price
                
        # Find ATM strike - get the strike closest to current price
        strikes = [contract.strike for contract in option_chain]
        atm_strike = min(strikes, key=lambda x: abs(x - underlying_price))
        
        # Get all contracts at ATM strike
        atm_contracts = [contract for contract in option_chain if contract.strike == atm_strike]
        
        # Separate by expiration dates
        expiry_dict = {}
        for contract in atm_contracts:
            days_to_expiry = (contract.expiry - self.algorithm.time).days
            if days_to_expiry not in expiry_dict:
                expiry_dict[days_to_expiry] = []
            expiry_dict[days_to_expiry].append(contract)
        
        # Need at least two different expiration dates
        if len(expiry_dict) < 2:
            # self.algorithm.log(f"Not enough expiration dates for {symbol}")
            return None, None
            
        return atm_strike, expiry_dict
    
    def find_expiry_pair(self, expiry_dict):
        """
        Find suitable pair of expirations for a calendar spread.
        
        Args:
            expiry_dict: Dictionary of option contracts by days to expiry
            
        Returns:
            tuple: (near_term_days, longer_term_days) or None if not found
        """
        # Find near-term and longer-term expirations
        expiry_days = sorted(expiry_dict.keys())
        near_term_days = expiry_days[0]
        
        # Find expiration ~30 days later than nearest (as specified in technical doc)
        target_days = near_term_days + 30
        longer_term_days = None
        for d in expiry_days:
            if d > near_term_days:
                if longer_term_days is None or abs(d - target_days) < abs(longer_term_days - target_days):
                    longer_term_days = d
        
        if longer_term_days is None:
            # self.algorithm.log(f"No suitable longer-term expiration")
            return None
            
        return near_term_days, longer_term_days
    
    def select_option_contracts(self, expiry_dict, near_term_days, longer_term_days, right=OptionRight.CALL):
        """
        Select option contracts for the given expirations and option right.
        
        Args:
            expiry_dict: Dictionary of option contracts by days to expiry
            near_term_days: Days to expiry for near-term option
            longer_term_days: Days to expiry for longer-term option
            right: Option right (CALL or PUT)
            
        Returns:
            tuple: (near_term_contract, longer_term_contract) or (None, None) if not found
        """
        # Select ATM contracts for the two expirations with the specified right
        near_term_options = [c for c in expiry_dict[near_term_days] if c.right == right]
        longer_term_options = [c for c in expiry_dict[longer_term_days] if c.right == right]
        
        if not near_term_options or not longer_term_options:
            return None, None
            
        # Select the contracts
        near_term_contract = near_term_options[0]
        longer_term_contract = longer_term_options[0]
        
        return near_term_contract, longer_term_contract
    
    def calculate_position_size(self, near_term_contract, longer_term_contract, risk_per_trade):        
        """
        Calculate position size based on risk percentage of portfolio.
        Implements the fixed percentage position sizing from the technical specification.
        
        Args:
            near_term_contract: The near-term option contract
            longer_term_contract: The longer-term option contract
            risk_per_trade: Percentage of portfolio to risk per trade
            
        Returns:
            Number of contracts to trade
        """
        # Check if we should just return 1 contract
        if (self.algorithm.get_parameter("justOneContract", 0)) == 1:
            return 1

        # Get prices
        near_term_price = near_term_contract.bid_price
        longer_term_price = longer_term_contract.ask_price
        
        # Calculate net debit for the calendar spread
        net_debit = longer_term_price - near_term_price
        
        # Calculate maximum risk based on portfolio value
        max_risk = self.algorithm.portfolio.cash * risk_per_trade
        
        # Calculate position size
        if net_debit <= 0:
            # If net credit, use minimum size
            return 1
            
        position_size = max(1, int(max_risk / (net_debit * 100)))  # multiply by 100 for option contract multiplier
        
        return position_size
    
    def is_underlying_valid(self, symbol):
        """
        Check if the underlying security is valid and has a price.
        
        Args:
            symbol: The underlying security symbol
            
        Returns:
            bool: True if valid, False otherwise
        """
        if symbol not in self.algorithm.securities:
            self.algorithm.log(f"....{symbol} not in securities collection")
            return False
            
        underlying_price = self.algorithm.securities[symbol].price
        if underlying_price == 0:
            self.algorithm.log(f"....Invalid price for {symbol}")
            return False
            
        return True 

from AlgorithmImports import *
from datetime import timedelta

class TradeManager:
    """
    Handles trade execution, position tracking, and order management
    for the earnings volatility calendar spread strategy.
    """
    
    def __init__(self, algorithm):
        """
        Initialize with a reference to the main algorithm.
        
        Args:
            algorithm: The main algorithm instance for accessing trading methods
        """
        self.algorithm = algorithm
        self.trades = {}  # Will store: {symbol: (near_term_symbol, long_term_symbol, position_size)}
    
    def enter_calendar_spread(self, symbol, near_term_contract, longer_term_contract, position_size):
        """
        Enter a calendar spread trade for the given symbol.
        
        Args:
            symbol: The underlying security symbol
            near_term_contract: The near-term option contract
            longer_term_contract: The longer-term option contract
            position_size: Number of contracts to trade
            
        Returns:
            bool: True if the trade was entered successfully
        """
        if not near_term_contract or not longer_term_contract:
            self.algorithm.log(f"....Could not find suitable option contracts for {symbol}")
            return False
            
        # Add option contracts before trading them
        near_term_contract_symbol = self.algorithm.add_option_contract(near_term_contract).symbol
        longer_term_contract_symbol = self.algorithm.add_option_contract(longer_term_contract).symbol

        # Place orders for calendar spread
        # Sell near-term, buy longer-term (long calendar spread)
        near_term_order = self.algorithm.sell(near_term_contract_symbol, position_size)
        longer_term_order = self.algorithm.buy(longer_term_contract_symbol, position_size)
        
        # Store trade information for later reference and exit
        self.trades[symbol] = (near_term_contract.symbol, longer_term_contract.symbol, position_size)
        
        self.algorithm.log(f"....Entered calendar spread for {symbol}: ")
        self.algorithm.log(f"....Sold {position_size} of {near_term_contract.symbol}")
        self.algorithm.log(f"....Bought {position_size} of {longer_term_contract.symbol}")
        
        return True
    
    def close_positions_for_symbol(self, symbol):
        """
        Close positions for a specific symbol.
        
        Args:
            symbol: The underlying symbol to close positions for
            
        Returns:
            bool: True if positions were closed, False if no positions found
        """
        if symbol not in self.trades:
            return False
            
        near_term_symbol, longer_term_symbol, position_size = self.trades[symbol]
        
        # Close positions - buy back what we sold, sell what we bought
        self.algorithm.buy(near_term_symbol, position_size)
        self.algorithm.sell(longer_term_symbol, position_size)
        
        self.algorithm.log(f"....Closed position for {symbol} after earnings")
        
        # Remove from active trades
        del self.trades[symbol]
        
        return True
    
    def close_all_positions(self):
        """
        Close all open positions.
        
        Returns:
            int: Number of positions closed
        """
        positions_closed = 0
        for symbol in list(self.trades.keys()):
            if self.close_positions_for_symbol(symbol):
                positions_closed += 1
                
        return positions_closed
    
    def process_earnings_exit(self, earnings_schedule):
        """
        Process exit logic for positions based on earnings reports.
        
        Args:
            earnings_schedule: Dictionary mapping symbols to earnings info (date, is_amc)
            
        Returns:
            list: Symbols for which positions were closed
        """
        closed_symbols = []
        
        for underlying_symbol, trade_info in list(self.trades.items()):
            if underlying_symbol not in earnings_schedule:
                continue

            earnings_info = earnings_schedule.get(underlying_symbol)
            if earnings_info is None:
                continue
                
            earnings_date, is_amc = earnings_info
            should_exit = False
            
            # Check if we should exit based on earnings timing (BMO/AMC)
            if is_amc:
                # For AMC earnings, exit the morning after the earnings date
                if earnings_date.date() < self.algorithm.time.date():
                    should_exit = True
                    self.algorithm.log(f"....{underlying_symbol} Exited - AMC earnings on {earnings_date.date()} (previous day or earlier)")
            else:
                # For BMO earnings, exit the same morning
                if earnings_date.date() == self.algorithm.time.date():
                    should_exit = True
                    self.algorithm.log(f"....{underlying_symbol} Exited - BMO earnings today {earnings_date.date()}")
            
            if not should_exit:
                continue
                
            if self.close_positions_for_symbol(underlying_symbol):
                closed_symbols.append(underlying_symbol)
        
        return closed_symbols
    
    def get_active_trades(self):
        """
        Get list of active trades.
        
        Returns:
            dict: Dictionary of active trades
        """
        return self.trades 

import numpy as np
from scipy.interpolate import interp1d
from AlgorithmImports import *

class VolatilityCalculator:
    """
    Handles all volatility and term structure related calculations for the strategy.
    """
    
    def __init__(self, algorithm):
        """
        Initialize with a reference to the main algorithm.
        
        Args:
            algorithm: The main algorithm instance for accessing history and other methods
        """
        self.algorithm = algorithm
    
    def build_term_structure(self, days, ivs):
        """
        Build term structure using linear interpolation to estimate IV at any DTE.
        
        Args:
            days: Array of days to expiration
            ivs: Array of implied volatilities corresponding to each expiration
            
        Returns:
            Function that interpolates IV for any given days to expiration
        """
        # Sort by days to ensure proper ordering
        sort_idx = days.argsort()
        days_sorted = days[sort_idx]
        ivs_sorted = ivs[sort_idx]
        
        # Create interpolation function
        spline = interp1d(days_sorted, ivs_sorted, kind='linear', fill_value="extrapolate")
        
        def term_spline(dte):
            if dte < days_sorted[0]:
                return ivs_sorted[0]
            elif dte > days_sorted[-1]:
                return ivs_sorted[-1]
            else:
                return float(spline(dte))
                
        return term_spline
    
    def calculate_term_structure(self, expiry_dict):
        """
        Calculate term structure data from option expirations.
        
        Args:
            expiry_dict: Dictionary of option contracts by days to expiry
            
        Returns:
            tuple: (days_array, iv_array)
        """
        days_array = np.array(sorted(expiry_dict.keys()))
        
        # Calculate average IV for each expiration (combining calls and puts)
        iv_array = np.array([
            np.mean([c.implied_volatility for c in expiry_dict[days]])
            for days in days_array
        ])
        
        return days_array, iv_array
    
    def calculate_term_structure_slope(self, days_array, iv_array):
        """
        Calculate the slope of the IV term structure.
        
        Args:
            days_array: Array of days to expiration
            iv_array: Array of implied volatilities
            
        Returns:
            tuple: (slope_value, term_structure_function)
        """
        # Create interpolation function for IV term structure
        term_structure = self.build_term_structure(days_array, iv_array)
        
        # Get IV at near term and at 45 days
        near_term_days = days_array[0]
        near_term_iv = term_structure(near_term_days)
        iv45 = term_structure(45)
        
        # Calculate term structure slope
        ts_slope = (iv45 - near_term_iv) / (45 - near_term_days)
        
        return ts_slope, term_structure
    
    def calculate_realized_volatility(self, symbol):
        """
        Calculate realized volatility using Yang-Zhang method.
        This is a more accurate volatility measurement that accounts for overnight jumps.
        
        Args:
            symbol: Stock symbol to calculate volatility for
            
        Returns:
            Annualized realized volatility or None if calculation fails
        """
        # Get price history
        history = self.algorithm.history(symbol, 30, Resolution.DAILY)
        if history.empty or len(history) < 30:
            return None
            
        # Extract OHLC prices
        open_prices = history['open'].values
        high_prices = history['high'].values
        low_prices = history['low'].values
        close_prices = history['close'].values
        
        # Calculate components for Yang-Zhang volatility
        n = len(close_prices)
        close_to_close = np.log(close_prices[1:] / close_prices[:-1])
        open_to_open = np.log(open_prices[1:] / open_prices[:-1])
        log_hl = np.log(high_prices[1:] / low_prices[1:])
        
        # Calculate overnight volatility (close to open)
        overnight_vol = np.sum(np.square(np.log(open_prices[1:] / close_prices[:-1]))) / (n - 1)
        
        # Calculate open to close volatility
        open_close_vol = np.sum(np.square(np.log(close_prices[1:] / open_prices[1:]))) / (n - 1)
        
        # Calculate Rogers-Satchell volatility
        log_ho = np.log(high_prices[1:] / open_prices[1:])
        log_lo = np.log(low_prices[1:] / open_prices[1:])
        log_co = np.log(close_prices[1:] / open_prices[1:])
        rs_vol = np.sum(log_ho * (log_ho - log_co) + log_lo * (log_lo - log_co)) / (n - 1)
        
        # Calculate k factor
        k = 0.34 / (1.34 + (n + 1) / (n - 1))
        
        # Combine the components for Yang-Zhang volatility
        yang_zhang_vol = overnight_vol + k * open_close_vol + (1 - k) * rs_vol
        
        # Convert to annualized volatility (assuming 252 trading days per year)
        annualized_vol = np.sqrt(yang_zhang_vol * 252)
        
        return annualized_vol
    
    def calculate_average_volume(self, symbol):
        """
        Calculate 30-day average trading volume for a symbol.
        Used as one of the three predictor variables.
        
        Args:
            symbol: Stock symbol to calculate average volume for
            
        Returns:
            30-day average volume or None if calculation fails
        """
        history = self.algorithm.history(symbol, 30, Resolution.DAILY)
        if history.empty or len(history) < 30:
            return None
            
        return history['volume'].mean()
    
    def calculate_iv_rv_ratio(self, symbol, term_structure):
        """
        Calculate the IV/RV ratio for a given symbol
        
        Args:
            symbol: The underlying security symbol
            term_structure: The precomputed term structure function
            
        Returns:
            float: IV/RV ratio or None if calculation failed
        """
        # Calculate 30-day realized volatility using Yang-Zhang method
        rv30 = self.calculate_realized_volatility(symbol)
        if rv30 is None or rv30 == 0:
            self.algorithm.log(f"....Could not calculate realized volatility for {symbol}")
            return None
        
        # Get implied volatility for 30 days from the provided term structure
        iv30 = term_structure(30)
        iv_rv_ratio = iv30 / rv30
        
        return iv_rv_ratio