# region imports
from AlgorithmImports import *
# endregion
# region imports
from datetime import timedelta
from model import LogisticRegressionModel
# endregion

class TradeDecisionEngine:
    """
    Implements the trade decision logic for SPX Long Straddle Strategy.
    Uses logistic regression model to determine trade entry conditions.
    """
    
    def __init__(self, algorithm):
        """
        Initialize the decision engine with reference to the main algorithm.
        
        Parameters:
        - algorithm: QCAlgorithm instance with access to market data
        """
        self.algorithm = algorithm
        self.model = LogisticRegressionModel()
        self.probability_threshold = 0.40  # Exact threshold value
    
    def should_enter_trade(self):
        """
        Determine if a trade should be entered based on model probability
        and additional filters.
        
        Returns:
        - boolean: True if all trade conditions are met, False otherwise
        """
        # Get indicator values using required calculations
        current_term_slope = self.algorithm.CalculateTermSlope()
        current_iv_zscore = self.algorithm.CalculateIVZScore(252)  # Must use 252 days
        
        # Get VIX value
        current_vix = self.algorithm.Securities[self.algorithm.vix_index_symbol].Price
        
        # Calculate trade probability
        z, trade_probability = self.model.calculate_probability(current_term_slope, current_iv_zscore)
        
        # Check conditions with exact thresholds
        probability_condition = trade_probability > self.probability_threshold
        vix_condition = current_vix >= 21
        
        # Log all conditions with Phase 3 prefix
        self.algorithm.CustomLog(f"[PHASE3] Probability calculation: z={z:.6f}, probability={trade_probability:.6f}")
        self.algorithm.CustomLog(f"[PHASE3] Conditions: Term Slope={current_term_slope:.6f}, IV Z-Score={current_iv_zscore:.6f}, VIX={current_vix:.2f}")
        self.algorithm.CustomLog(f"[PHASE3] VIX filter: {vix_condition} (VIX = {current_vix:.2f})")
        self.algorithm.CustomLog(f"[PHASE3] Trade probability: {trade_probability:.6f}, Threshold: {self.probability_threshold}, Decision: {probability_condition and vix_condition}")
        
        # Decision must check both conditions
        return probability_condition and vix_condition
        
    def verify_dte(self, option_chain):
        """
        Verify that options in the chain have exactly 1 DTE.
        
        Parameters:
        - option_chain: Collection of option contracts
        
        Returns:
        - boolean: True if valid 1 DTE options exist, False otherwise
        """
        if not option_chain:
            self.algorithm.CustomLog("[PHASE3] No option chain available to verify DTE")
            return False
            
        # Target expiry date should be exactly 1 day from now
        target_expiry_date = self.algorithm.Time.date() + timedelta(days=1)
        
        # Filter for contracts expiring tomorrow (1 DTE)
        contracts_1dte = [contract for contract in option_chain 
                          if contract.Expiry.date() == target_expiry_date]
        
        has_valid_dte = len(contracts_1dte) > 0
        
        if not has_valid_dte:
            self.algorithm.CustomLog(f"[PHASE3] No 1 DTE contracts found, target expiry: {target_expiry_date}")
        else:
            self.algorithm.CustomLog(f"[PHASE3] Verified 1 DTE options: {len(contracts_1dte)} contracts available")
            
        return has_valid_dte 

# region imports
from datetime import datetime
from AlgorithmImports import *
# endregion

class TradeExecutionEngine:
    """
    Handles position sizing and order execution for SPX Long Straddle Strategy.
    Implements exact specifications from Phase 4 of the system design.
    """
    
    def __init__(self, algorithm):
        """Initialize with reference to main algorithm."""
        self.algorithm = algorithm
        self.active_orders = {}  # OrderId -> OrderDetails tracking
        
    def execute_straddle_position(self, atm_options):
        """
        Execute a straddle position with exact position sizing rules.
        
        Parameters:
        - atm_options: ATM call and put options to trade
        """
        # Calculate position size (exactly 20% of portfolio)
        portfolio_value = self.algorithm.Portfolio.TotalPortfolioValue
        position_value = portfolio_value * 0.20
        
        # Calculate option costs
        call_price = atm_options.Call.LastPrice
        put_price = atm_options.Put.LastPrice
        total_cost = call_price + put_price
        
        # Calculate contracts with hard limit of 1
        contracts_by_value = math.floor(position_value / total_cost)
        contracts = min(contracts_by_value, 1)
        
        # Log position sizing details
        self.algorithm.CustomLog(f"[PHASE4] Position sizing: Portfolio=${portfolio_value:.2f}, "
                               f"Allocation=${position_value:.2f} (20%), "
                               f"Option Cost=${total_cost:.2f}, "
                               f"Contracts calculated={contracts_by_value}, "
                               f"Final contracts={contracts}")
        
        # If zero contracts calculated but we have enough for 1, proceed with 1
        if contracts == 0 and portfolio_value >= total_cost:
            self.algorithm.CustomLog("[PHASE4] Warning: Calculated contracts is zero. Attempting to trade 1 contract.")
            contracts = 1
            
        # Place market orders for both options
        if contracts > 0:
            try:
                # Place call order
                call_order = self.algorithm.MarketOrder(atm_options.Call.Symbol, contracts)
                self.active_orders[call_order.OrderId] = OrderDetails(call_order, "Call")
                
                # Place put order
                put_order = self.algorithm.MarketOrder(atm_options.Put.Symbol, contracts)
                self.active_orders[put_order.OrderId] = OrderDetails(put_order, "Put")
                
                # Log order placement
                self.algorithm.CustomLog(
                    f"[PHASE4] Straddle executed: "
                    f"Call={atm_options.Call.Symbol}, "
                    f"Put={atm_options.Put.Symbol}, "
                    f"Total Cost=${total_cost * contracts:.2f}"
                )
                
            except Exception as e:
                self.algorithm.CustomLog(f"[PHASE4] Error executing orders: {str(e)}")
        else:
            self.algorithm.CustomLog("[PHASE4] Insufficient buying power for minimum position")
            
    def on_order_event(self, order_event):
        """Handle order status updates and partial fills."""
        if order_event.OrderId in self.active_orders:
            order_details = self.active_orders[order_event.OrderId]
            
            # Log order status
            self.algorithm.CustomLog(
                f"[PHASE4] Order {order_event.OrderId} status: {order_event.Status} "
                f"({order_details.option_type})"
            )
            
            # Handle partial fills
            if order_event.Status == OrderStatus.PartiallyFilled:
                self.algorithm.CustomLog(
                    f"[PHASE4] Partial fill for Order {order_event.OrderId}: "
                    f"{order_event.FillQuantity}/{order_event.Quantity}"
                )

class OrderDetails:
    """Simple struct to track order information."""
    def __init__(self, order, option_type):
        self.order = order
        self.option_type = option_type  # "Call" or "Put" 

# region imports
from AlgorithmImports import *
import numpy as np
import math
from model import LogisticRegressionModel
from decision import TradeDecisionEngine
from execution import TradeExecutionEngine
from trade_management import TradeManagementEngine
# endregion

# Renamed class to SPXLongStraddleStrategy
class SPXLongStraddleStrategy(QCAlgorithm):

    def Initialize(self):
        """
        Phase 1: Initialize strategy parameters, data subscriptions, and scheduling.
        Phase 2: Add warm-up period and initialize storage for historical IVs
        Phase 3: Initialize logistic regression model and trade decision engine
        Phase 4: Initialize trade execution engine
        Phase 5: Initialize trade management engine
        """
        self.SetStartDate(2019, 3, 31)  # Example Start Date
        # self.SetEndDate(2025, 4, 14)    # Example End Date
        self.SetCash(100000)           # Example Cash

        # --- Set warm-up period for historical IV calculation ---
        self.SetWarmUp(365)  # 365 calendar days to ensure at least 252 trading days
        
        # --- Model Coefficients (Phase 3) ---
        self.model_intercept = -0.60
        self.model_term_slope_coeff = -0.57
        self.model_iv_zscore_coeff = 0.36
        # self.CustomLog(f"[PHASE3] Model coefficients initialized: Intercept={self.model_intercept}, TermSlope={self.model_term_slope_coeff}, IVZScore={self.model_iv_zscore_coeff}")

        # --- Symbol Setup ---
        self.spx_index_symbol = self.AddIndex("SPX", Resolution.Minute).Symbol
        self.vix_index_symbol = self.AddIndex("VIX", Resolution.Minute).Symbol
        
        # --- Modified Option Setup with Better Filtering ---
        self.spx_option = self.AddIndexOption(self.spx_index_symbol, "SPXW", Resolution.Minute)
        self.option_symbol = self.spx_option.Symbol
        # Set filter for trading options (1 DTE)
        self.spx_option.SetFilter(lambda universe: universe
                                  .IncludeWeeklys() # Ensure we are using weeklys (SPXW)
                                  .Expiration(0, 1) # Only 0 and 1 DTE (Contracts expire T+1, so 1 DTE is needed on trade day)
                                  .Strikes(-10, 10)) # Strikes around ATM
        
        # --- Separate Option Chain for IV Calculation (30 DTE) ---
        self.iv_option = self.AddIndexOption(self.spx_index_symbol, Resolution.Minute)
        self.iv_option_symbol = self.iv_option.Symbol
        # Set filter targeting ~30 day options for IV calculation
        self.iv_option.SetFilter(lambda universe: universe
                               .IncludeWeeklys() # Include both standard and weekly options
                               .Expiration(25, 35) # Target options ~30 days out (25-35 day window)
                               .Strikes(-1, 1)) # Only closest strikes to ATM for more accurate IV
        
        # --- Phase 2: Historical IV Storage ---
        self.historical_iv_values = []  # Store historical IV values
        self.last_iv_date = None  # Track the last date we calculated IV for
        
        # --- Phase 3: Initialize trade decision engine ---
        self.decision_engine = TradeDecisionEngine(self)
        self.CustomLog("[PHASE3] Trade decision engine initialized with 40% probability threshold")
        
        # --- Phase 4: Initialize trade execution engine ---
        self.execution_engine = TradeExecutionEngine(self)
        self.CustomLog("[PHASE4] Trade execution engine initialized")
        
        # --- Scheduling (Phase 1.1) ---
        # Schedule trade evaluation before market close
        self.Schedule.On(self.DateRules.EveryDay(self.spx_index_symbol),
                         self.TimeRules.At(15, 47), # 3:47 PM ET
                         self.EvaluateSignal)
        self.CustomLog("[PHASE3] Scheduled EvaluateSignal daily at 3:47 PM ET for trade evaluation")

        # Schedule expiration check 2 minutes before market close
        self.Schedule.On(self.DateRules.EveryDay(self.spx_index_symbol),
                         self.TimeRules.At(15, 58), # 3:58 PM ET
                         self.CheckExpiringPositions)
        self.CustomLog("[PHASE5] Scheduled expiration check daily at 3:58 PM ET")

        # --- Phase 2: Schedule IV calculation once per day ---
        self.Schedule.On(self.DateRules.EveryDay(self.spx_index_symbol),
                         self.TimeRules.AfterMarketOpen(self.spx_index_symbol, 30),
                         self.CalculateDailyIV)

        # --- Phase 5: Initialize trade management engine ---
        self.trade_manager = TradeManagementEngine(self)
        self.CustomLog("[PHASE5] Trade management engine initialized with 90% stop loss")

        # self.CustomLog("[PHASE1] Algorithm initialization complete.")


    def OnData(self, data: Slice):
        """
        Phase 1: Basic option chain processing to verify ATM selection.
        Phase 5: Monitor stop loss conditions.
        """
        # Check if option chain data is available for trading
        chain = data.OptionChains.get(self.option_symbol)
        if chain is None or not chain:
            return # Wait for chain data

        # Only log ATM selection verification once per day around market open
        if self.Time.hour == 9 and self.Time.minute == 31 and self.Time.second == 0:
            # self.CustomLog(f"[PHASE1] Received SPX option chain for {self.Time.date()}.")

            underlying_price = chain.Underlying.Price
            if underlying_price <= 0:
                 # self.CustomLog(f"[PHASE1] Underlying price is not valid: {underlying_price}")
                 return

            # Filter for contracts expiring tomorrow (1 DTE relative to trade entry)
            target_expiry_date = self.Time.date() + timedelta(days=1)
            contracts_1dte = [contract for contract in chain
                              if contract.Expiry.date() == target_expiry_date]

            # Debug logging for DTE issue
            # if not contracts_1dte:
            #     # self.CustomLog(f"[DEBUG-DTE] Target expiry date: {target_expiry_date} (Day: {target_expiry_date.strftime('%A')})")
            #     available_expiries = sorted(set(contract.Expiry.date() for contract in chain))
            #     if available_expiries:
            #         # self.CustomLog(f"[DEBUG-DTE] Available expiries ({len(available_expiries)}): {available_expiries[:5]} {'' if len(available_expiries) <= 5 else '...'}")
            #         # Log closest expiry to target to help diagnose
            #         closest_expiry = min(available_expiries, key=lambda d: abs((d - target_expiry_date).days))
            #         self.CustomLog(f"[DEBUG-DTE] Closest available expiry: {closest_expiry} ({(closest_expiry - self.Time.date()).days} DTE)")
            #     else:
            #         self.CustomLog(f"[DEBUG-DTE] No expiry dates available in chain")

            if not contracts_1dte:
                # self.CustomLog(f"[PHASE1] No contracts found for target expiry {target_expiry_date} in the chain.")
                # self.CustomLog(f"[PHASE1] No 1 DTE contracts found for ATM selection check.")
                return
            else:
                # self.CustomLog(f"[DEBUG-DTE] Got 1 DTE contract!")
                contracts_to_use = contracts_1dte

            # Find ATM strike (closest to underlying price)
            atm_strike = min(contracts_to_use, key=lambda x: abs(x.Strike - underlying_price)).Strike
            # self.CustomLog(f"[PHASE1] Underlying Price: {underlying_price:.2f}, Identified ATM Strike: {atm_strike}")

            # Select ATM call and put
            atm_call = next((c for c in contracts_to_use if c.Strike == atm_strike and c.Right == OptionRight.Call), None)
            atm_put = next((c for c in contracts_to_use if c.Strike == atm_strike and c.Right == OptionRight.Put), None)

            # if atm_call and atm_put:
            #     # self.CustomLog(f"[PHASE1] ATM Call found: {atm_call.Symbol.Value}")
            #     # self.CustomLog(f"[PHASE1] ATM Put found: {atm_put.Symbol.Value}")
            # else:
            #     # self.CustomLog(f"[PHASE1] Could not find both ATM Call and Put for strike {atm_strike}.")

        # Check IV calculation chain (30-day options)
        iv_chain = data.OptionChains.get(self.iv_option_symbol)
        if iv_chain is None or not iv_chain:
            return
            
        # # Log IV option chain data once per day (for verification)
        # if self.Time.hour == 9 and self.Time.minute == 32 and self.Time.second == 0:
        #     # Get closest expiration to 30 days
        #     today = self.Time.date()
        #     target_date = today + timedelta(days=30)
        #     options_by_expiry = {}
            
        #     # Group options by expiry date
        #     for contract in iv_chain:
        #         expiry = contract.Expiry.date()
        #         if expiry not in options_by_expiry:
        #             options_by_expiry[expiry] = []
        #         options_by_expiry[expiry].append(contract)
            
        #     # Find expiry closest to 30 days
        #     if options_by_expiry:
        #         closest_expiry = min(options_by_expiry.keys(), key=lambda d: abs((d - today).days - 30))
        #         days_to_expiry = (closest_expiry - today).days
        #         # self.CustomLog(f"[PHASE2] 30-day IV chain: Found options expiring in {days_to_expiry} days ({closest_expiry})")
        #         # self.CustomLog(f"[PHASE2] Number of contracts in 30-day chain: {len(options_by_expiry[closest_expiry])}")
        #     else:
        #         # self.CustomLog("[PHASE2] No options found in the 30-day IV chain")

        # Phase 5: Check stop loss if we have open positions
        if self.trade_manager.check_stop_loss():
            # Get current positions for tagging
            open_positions = [pos for pos in self.Portfolio.Values if pos.Invested]
            # Create liquidation orders
            orders = [self.Liquidate(pos.Symbol, tag=f"STOP LOSS") for pos in open_positions]



    def EvaluateSignal(self):
        """
        Phase 2: Calculate indicators and log values
        Phase 3: Evaluate trade signal using logistic regression model
        Phase 4: Execute trades when conditions are met
        """
        if self.IsWarmingUp:
            return
            
        # self.CustomLog(f"[PHASE3] EvaluateSignal triggered at {self.Time}")
        
        # Verify we have option chain data with 1 DTE options
        chain = self.OptionChain(self.option_symbol)
        if not chain or not self.decision_engine.verify_dte(chain):
            self.CustomLog("[PHASE3] Cannot evaluate signal: missing option chain or invalid DTE")
            return
        
        # Calculate and log the term structure slope
        term_slope = self.CalculateTermSlope()
        
        # Calculate and log the IV Z-Score
        iv_zscore = self.CalculateIVZScore()
        
        # Phase 3: Calculate probability and evaluate trade conditions
        should_trade = self.ShouldEnterTrade()
        
        # self.CustomLog(f"[PHASE3] Signal evaluation complete: Term Slope={term_slope:.6f}, IV Z-Score={iv_zscore:.6f}, Should Trade={should_trade}")
        
        # Phase 4: If trade conditions are met, execute the trade
        if should_trade:
            self.CustomLog("[PHASE4] Trade conditions met, preparing to execute straddle position")
            atm_options = self.GetATMOptions()
            if atm_options.Call and atm_options.Put:
                self.execution_engine.execute_straddle_position(atm_options)
            else:
                self.CustomLog("[PHASE4] Could not find valid ATM options for straddle execution")


    def CalculateDailyIV(self):
        """
        Calculate and store the daily ATM IV value using 30-day options.
        Used for historical IV tracking.
        """
        # if self.IsWarmingUp:
        #     # self.CustomLog(f"[PHASE2] Still in warm-up period, collecting historical IV data.")
        
        # Only calculate once per day
        if self.last_iv_date == self.Time.date():
            return
            
        current_iv = self.GetCurrentATMIV()
        
        # Store the IV value if valid
        if current_iv > 0:
            self.historical_iv_values.append(current_iv)
            self.last_iv_date = self.Time.date()
            # self.CustomLog(f"[PHASE2] Daily IV calculated: {current_iv:.6f}, Historical IV count: {len(self.historical_iv_values)}")
        else:
            self.CustomLog(f"[PHASE2] Could not calculate valid IV for today.")


    def yeo_johnson_transform(self, x, lambda_param=0.5):
        """
        Apply Yeo-Johnson transformation to a single value.
        Using fixed lambda=0.5 as per strategy specification.
        """
        if x >= 0:
            numerator = np.power((x + 1), lambda_param) - 1
            return numerator / lambda_param if lambda_param != 0 else np.log(x + 1)
        else:
            numerator = np.power((-x + 1), (2 - lambda_param)) - 1
            return -numerator / (2 - lambda_param) if lambda_param != 2 else -np.log(-x + 1)

    def CalculateTermSlope(self):
        """
        Calculate the term structure slope by:
        1. Get SPX option chains for different expirations (1, 7, 14, 30, 60, 90 days)
        2. For each expiration, find ATM options and calculate average IV
        3. Perform linear regression on [days, IV] data points
        4. Apply Yeo-Johnson transformation to the slope
        """
        # Target DTE values for the term structure
        target_dte_values = [1, 7, 14, 30, 60, 90]
        term_structure_points = []
        
        # Get SPX underlying price
        spx_price = self.Securities[self.spx_index_symbol].Price
        
        # Get full option chain for current time
        option_chain = self.OptionChain(self.spx_index_symbol)
        if not option_chain:
            return 0
            
        # Extract all available expiration dates from the option chain
        today = self.Time.date()
        available_expiries = []
        expiry_options = {}
        
        for contract in option_chain:
            expiry = contract.Expiry.date()
            days_to_expiry = (expiry - today).days
            
            # Only consider expirations in the future
            if days_to_expiry > 0:
                if expiry not in expiry_options:
                    available_expiries.append((days_to_expiry, expiry))
                    expiry_options[expiry] = []
                expiry_options[expiry].append(contract)
        
        # Sort available expiries by DTE
        available_expiries.sort(key=lambda x: x[0])
        
        if not available_expiries:
            return 0
            
        # Find closest available expiration for each target DTE using a greedy approach
        used_expiries = set()
        dte_to_expiry = {}
        
        for target_dte in target_dte_values:
            # Find the closest available expiration that hasn't been used
            closest_expiry = None
            min_diff = float('inf')
            
            for actual_dte, expiry in available_expiries:
                if expiry in used_expiries:
                    continue
                    
                diff = abs(actual_dte - target_dte)
                if diff < min_diff:
                    min_diff = diff
                    closest_expiry = (actual_dte, expiry)
            
            # If found, mark as used and add to mapping
            if closest_expiry:
                actual_dte, expiry = closest_expiry
                used_expiries.add(expiry)
                dte_to_expiry[target_dte] = (actual_dte, expiry)
        
        # Calculate term structure points using the mapped expiration dates
        for target_dte, (actual_dte, expiry) in dte_to_expiry.items():
            # Get the options for this expiration
            options_for_expiry = expiry_options[expiry]
            
            # Find the strike price closest to current SPX price
            closest_strike = min(set(contract.Strike for contract in options_for_expiry), 
                               key=lambda strike: abs(strike - spx_price))
            
            # Get all options with that strike price
            atm_options = [contract for contract in options_for_expiry 
                         if contract.Strike == closest_strike]
            
            # Average the IVs of these options
            if atm_options:
                total_iv = sum(option.ImpliedVolatility for option in atm_options)
                avg_iv = total_iv / len(atm_options)
                term_structure_points.append([actual_dte, avg_iv])
        
        # Need at least 3 points for valid regression
        if len(term_structure_points) < 3:
            return 0
        
        # Extract days and ivs from points
        days = [point[0] for point in term_structure_points]
        ivs = [point[1] for point in term_structure_points]
        
        # Perform linear regression using numpy
        slope_coefficient = np.polyfit(days, ivs, 1)[0]
        
        # Apply Yeo-Johnson transformation to the slope
        transformed_slope = self.yeo_johnson_transform(slope_coefficient)
        
        self.CustomLog(f"[PHASE2] Term structure slope: raw={slope_coefficient:.6f}, transformed={transformed_slope:.6f}")
        
        return transformed_slope

    def CustomLog(self, message):
        if not self.IsWarmingUp:
            self.Log(f"{message}")

    def GetCurrentATMIV(self):
        """
        Get the current ATM implied volatility by:
        1. Get the 30-day SPX option chain (from the specially filtered chain)
        2. Find ATM options (within 0.5% of current price)
        3. Calculate the average IV of ATM call and put
        
        Modified to specifically use the 30-day option chain
        """
        # Get current 30-day option chain
        option_chain = self.OptionChain(self.iv_option_symbol)
        
        # Check if the option chain exists and has contracts
        if option_chain is None or not option_chain:
            # self.CustomLog("[PHASE2] No options available in the 30-day chain")
            return 0
            
        # Get SPX price
        spx_price = self.Securities[self.spx_index_symbol].Price
        
        # Find options closest to 30 days to expiration
        today = self.Time.date()
        target_date = today + timedelta(days=30)
        
        # Group options by expiry
        options_by_expiry = {}
        for contract in option_chain:
            expiry = contract.Expiry.date()
            if expiry not in options_by_expiry:
                options_by_expiry[expiry] = []
            options_by_expiry[expiry].append(contract)
        
        # If no options available, return 0
        if options_by_expiry is None or len(options_by_expiry) == 0:
            # self.CustomLog("[PHASE2] No options grouped by expiry in the 30-day chain")
            return 0
            
        # Find expiry closest to 30 days
        closest_expiry = min(options_by_expiry.keys(), key=lambda d: abs((d - today).days - 30))
        thirty_day_options = options_by_expiry[closest_expiry]
        days_to_expiry = (closest_expiry - today).days
        
        # self.CustomLog(f"[PHASE2] Using options expiring in {days_to_expiry} days for IV calculation")
        
        # Find the strike price closest to current SPX price
        closest_strike = min(set(contract.Strike for contract in thirty_day_options), 
                           key=lambda strike: abs(strike - spx_price))
        
        # Get ATM options with the closest strike
        atm_options = [contract for contract in thirty_day_options 
                     if contract.Strike == closest_strike]
        
        # Get ATM call and put for straddle
        atm_call = next((contract for contract in atm_options 
                      if contract.Right == OptionRight.Call), None)
        atm_put = next((contract for contract in atm_options 
                     if contract.Right == OptionRight.Put), None)
        
        # If both call and put are found, return the average IV of the straddle
        if atm_call is not None and atm_put is not None:
            iv = (atm_call.ImpliedVolatility + atm_put.ImpliedVolatility) / 2
            # self.CustomLog(f"[PHASE2] 30-day ATM straddle IV: {iv:.6f} (Call IV: {atm_call.ImpliedVolatility:.6f}, Put IV: {atm_put.ImpliedVolatility:.6f})")
            return iv
        
        # If only one is found, return that IV
        elif atm_call is not None:
            # self.CustomLog(f"[PHASE2] Using only 30-day ATM call IV: {atm_call.ImpliedVolatility:.6f}")
            return atm_call.ImpliedVolatility
        elif atm_put is not None:
            # self.CustomLog(f"[PHASE2] Using only 30-day ATM put IV: {atm_put.ImpliedVolatility:.6f}")
            return atm_put.ImpliedVolatility
        
        # If no ATM options found, return 0
        else:
            # self.CustomLog("[PHASE2] No ATM options found for IV calculation")
            return 0


    def GetHistoricalATMStraddleIV(self, historical_date):
        """
        Helper method that retrieves or estimates the historical ATM straddle IV
        In a live implementation, this would use a database of historical IVs
        For this implementation, we'll use our accumulated values from the warm-up period
        """
        # For simplicity, we'll use our accumulated values during warm-up
        # In a real implementation, additional steps would be needed to get 
        # historical option data for specific dates
        
        # Not used in this implementation - we're accumulating data as we go in CalculateDailyIV
        return 0


    def CalculateIVZScore(self, lookback_period=252):
        """
        Calculate the Z-Score of the current IV relative to historical IVs:
        1. Get historical IV values (252 trading days)
        2. Calculate mean and standard deviation
        3. Get current IV
        4. Calculate Z-Score as (current_iv - mean_iv) / std_iv
        """
        # Check if we have enough historical data
        if len(self.historical_iv_values) < lookback_period:
            # self.CustomLog(f"[PHASE2] Insufficient historical IV data: {len(self.historical_iv_values)}/{lookback_period} days")
            return 0
        
        # Use the most recent lookback_period values
        historical_iv_values = self.historical_iv_values[-lookback_period:]
        
        # Handle missing data with linear interpolation
        # Not needed here as we're ensuring valid values in CalculateDailyIV
        
        # Calculate statistics
        mean_iv = np.mean(historical_iv_values)
        std_iv = np.std(historical_iv_values, ddof=1)  # Use sample standard deviation
        
        # Get current IV (average of ATM call and put - straddle IV)
        current_iv = self.GetCurrentATMIV()
        
        # Calculate Z-score with exact formula
        z_score = (current_iv - mean_iv) / std_iv
        
        # Log results with high precision
        # self.CustomLog(f"[PHASE2] IV History: Mean={mean_iv:.6f}, StdDev={std_iv:.6f}")
        # self.CustomLog(f"[PHASE2] Current IV: {current_iv:.6f}, Z-Score: {z_score:.6f}")
        
        return z_score

    def CalculateTradeProbability(self, term_slope, iv_zscore):
        """
        Calculate the probability of a successful trade using the logistic regression model.
        
        Parameters:
        - term_slope: The term structure slope
        - iv_zscore: The implied volatility Z-score
        
        Returns:
        - probability: Value between 0 and 1 representing trade probability
        """
        # Use the model to calculate probability
        z, probability = LogisticRegressionModel().calculate_probability(term_slope, iv_zscore)
        
        # Log calculation details
        # self.CustomLog(f"[PHASE3] Probability calculation: z={z:.6f}, probability={probability:.6f}")
        
        return probability
        
    def ShouldEnterTrade(self):
        """
        Determine if a trade should be entered based on model and market conditions.
        
        Returns:
        - boolean: True if all trade conditions are met, False otherwise
        """
        # Use decision engine to evaluate trade conditions
        return self.decision_engine.should_enter_trade()

    def GetATMOptions(self):
        """
        Get ATM call and put options for straddle execution.
        Returns object with Call and Put properties.
        """
        # Get current option chain
        chain = self.OptionChain(self.option_symbol)
        if not chain:
            self.CustomLog("[PHASE4] No option chain available for ATM selection")
            return None
        
        # Get SPX price
        spx_price = self.Securities[self.spx_index_symbol].Price
        
        # Get target expiration (1 DTE)
        target_expiry = self.Time.date() + timedelta(days=1)
        
        # Filter for options with target expiration
        chain_for_expiry = [contract for contract in chain 
                          if contract.Expiry.date() == target_expiry]
        
        if not chain_for_expiry:
            self.CustomLog("[PHASE4] No options found for target expiration")
            return None
        
        # Find strike price closest to current SPX price
        closest_strike = min(set(contract.Strike for contract in chain_for_expiry), 
                           key=lambda strike: abs(strike - spx_price))
        
        # Get ATM call and put
        atm_call = next((contract for contract in chain_for_expiry 
                       if contract.Strike == closest_strike 
                       and contract.Right == OptionRight.Call), None)
        atm_put = next((contract for contract in chain_for_expiry 
                      if contract.Strike == closest_strike 
                      and contract.Right == OptionRight.Put), None)
        
        if atm_call and atm_put:
            self.CustomLog(f"[PHASE4] Selected ATM options: Strike={closest_strike}, "
                          f"Call={atm_call.Symbol}, Put={atm_put.Symbol}")
            
            # Create result object
            result = type('ATMOptions', (), {})()
            result.Call = atm_call
            result.Put = atm_put
            return result
        else:
            self.CustomLog("[PHASE4] Could not find both ATM call and put options")
            return None

    def OnOrderEvent(self, orderEvent):
        """Handle order events for trade execution tracking."""
        self.execution_engine.on_order_event(orderEvent)
        
        # Phase 5: Initialize position tracking when both legs of straddle are filled
        if orderEvent.Status == OrderStatus.Filled:
            order_details = self.execution_engine.active_orders.get(orderEvent.OrderId)
            if order_details:
                # Check if both legs are filled
                filled_orders = [o for o in self.execution_engine.active_orders.values() 
                               if o.order.Status == OrderStatus.Filled]
                if len(filled_orders) == 2:  # Both call and put are filled
                    total_cost = sum(o.order.AbsoluteQuantity * o.order.AverageFillPrice 
                                   for o in filled_orders)
                    call_order = next(o for o in filled_orders if o.option_type == "Call")
                    put_order = next(o for o in filled_orders if o.option_type == "Put")
                    self.trade_manager.initialize_position_tracking(
                        call_order.order.Symbol,
                        put_order.order.Symbol,
                        total_cost
                    )

   
    def CheckExpiringPositions(self):
        """Check for and liquidate positions expiring today."""
        self.trade_manager.check_expiring_positions()

    # --- Placeholder methods for future phases ---

    def ExecuteStraddlePosition(self): # Phase 4
        # Implementation needed
        pass

# region imports
from AlgorithmImports import *
# endregion
# region imports
import math
import numpy as np
# endregion

class LogisticRegressionModel:
    """
    Implements the logistic regression model for SPX Long Straddle Strategy
    with exact coefficient values as specified in the requirements.
    """
    
    def __init__(self):
        """Initialize model with exact coefficient values."""
        self.intercept = -0.60
        self.term_slope_coefficient = -0.57
        self.iv_zscore_coefficient = 0.36
    
    def calculate_probability(self, term_slope, iv_zscore):
        """
        Calculate trade probability using exact logistic regression formula.
        
        Parameters:
        - term_slope: The term structure slope
        - iv_zscore: The implied volatility Z-score
        
        Returns:
        - probability: Value between 0 and 1 representing trade probability
        """
        # Calculate logistic regression value with exact coefficients
        z = self.intercept + (self.term_slope_coefficient * term_slope) + (self.iv_zscore_coefficient * iv_zscore)
        
        # Apply sigmoid function with exact implementation
        probability = 1 / (1 + math.exp(-z))
        
        return z, probability 

# region imports
from datetime import datetime
from AlgorithmImports import *
# endregion

class TradeManagementEngine:
    """
    Handles trade management functionality including stop loss monitoring
    and order tagging for SPX Long Straddle Strategy.
    """
    
    def __init__(self, algorithm):
        """Initialize with reference to main algorithm."""
        self.algorithm = algorithm
        self.initial_straddle_premium = None
        self.stop_loss_level = None
    
    def initialize_position_tracking(self, call_symbol, put_symbol, total_cost):
        """Initialize tracking for a new straddle position."""
        self.initial_straddle_premium = total_cost
        self.stop_loss_level = total_cost * 0.10  # 10% of initial (90% loss)
        self.algorithm.CustomLog(f"[PHASE5] Stop loss set: Initial premium=${total_cost:.2f}, "
                               f"Stop level=${self.stop_loss_level:.2f} (10% of initial)")
    
    def check_stop_loss(self):
        """Monitor position value and trigger stop loss if threshold is breached."""
        if self.stop_loss_level is None or not self.algorithm.Portfolio.Invested:
            return False
            
        current_value = self.get_current_position_value()
        if current_value <= 0:  # Invalid value
            return False
            
        buffer = current_value - self.stop_loss_level
        self.algorithm.CustomLog(f"[PHASE5] Stop loss check: Current value=${current_value:.2f}, "
                               f"Stop level=${self.stop_loss_level:.2f}, "
                               f"Remaining buffer=${buffer:.2f}")
        
        if current_value <= self.stop_loss_level:
            self.algorithm.CustomLog(f"[PHASE5] STOP LOSS TRIGGERED: "
                                   f"Current value=${current_value:.2f} below stop level ${self.stop_loss_level:.2f}")
            return True
        return False
    
    def get_current_position_value(self):
        """Calculate current value of the straddle position."""
        if not self.algorithm.Portfolio.Invested:
            return 0
            
        total_value = 0
        for holding in self.algorithm.Portfolio.Values:
            if holding.Invested:
                total_value += holding.HoldingsValue
        return total_value
    
    def tag_exit_orders(self, orders):
        """Tag exit orders with P&L information."""
        if not self.initial_straddle_premium or not self.algorithm.Portfolio.Invested:
            return
            
        current_value = self.get_current_position_value()
        pnl_pct = ((current_value / self.initial_straddle_premium) - 1) * 100
        tag = f"{'WIN' if pnl_pct > 0 else 'LOSS'} {abs(pnl_pct):.2f}%"
        
        for order in orders:
            order.Tag = tag
            position = self.algorithm.Portfolio[order.Symbol]
            position_type = "Call" if order.Symbol.ID.OptionRight == OptionRight.Call else "Put"
            self.algorithm.CustomLog(f"[PHASE5] Tagged {position_type} exit order {order.Id} with: {tag}")

    def check_expiring_positions(self):
        """Check for and liquidate positions expiring today."""
        if not self.algorithm.Portfolio.Invested:
            return False
            
        today = self.algorithm.Time.date()
        expiring_positions = []
        
        # Check each position in the portfolio for expiration
        for holding in self.algorithm.Portfolio.Values:
            if holding.Invested and holding.Symbol.ID.Date.date() == today:
                expiring_positions.append(holding)
        
        # Liquidate expiring positions
        if expiring_positions:
            self.algorithm.CustomLog(f"[PHASE5] Found {len(expiring_positions)} positions expiring today, liquidating...")
            
            # Calculate P&L using portfolio values
            total_cost_basis = sum(pos.HoldingsCost for pos in expiring_positions)
            total_value = sum(pos.HoldingsValue for pos in expiring_positions)
            
            if total_cost_basis != 0:
                pnl_pct = ((total_value / abs(total_cost_basis)) - 1) * 100
                tag = f"{'WIN' if pnl_pct > 0 else 'LOSS'} {abs(pnl_pct):.2f}%"
            else:
                tag = "EOD Exit"
            
            # Liquidate with tag
            for pos in expiring_positions:
                position_type = "Call" if pos.Symbol.ID.OptionRight == OptionRight.Call else "Put"
                self.algorithm.CustomLog(f"[PHASE5] Liquidating {position_type} position with tag: {tag}")
                self.algorithm.Liquidate(pos.Symbol, tag)
            return True
        return False