from AlgorithmImports import *
from datetime import datetime
import math
from scipy.stats import kurtosis
import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler, LabelEncoder
from xgboost import XGBClassifier
from collections import deque

class MySecurityInitializer(BrokerageModelSecurityInitializer):
    def __init__(self, brokerage_model: IBrokerageModel, security_seeder: ISecuritySeeder):
        super().__init__(brokerage_model, security_seeder)

    def Initialize(self, security: Security):
        # First call the base class initialization
        super().Initialize(security)

class CombinedOptionsAlpha(QCAlgorithm):
    def Initialize(self):

        self.SetStartDate(2022, 5, 1)
        self.SetEndDate(2024, 12, 1)
        self.SetCash(10000000)
        self.SetTimeZone(TimeZones.NewYork)

        self.SetWarmup(30)

        # for FactorAlpha model
        self.symbols = []

        self.straddle_alpha = DeltaHedgedStraddleAlpha(self)
        self.condor_alpha = IronCondorAlpha(self)
        self.factor_alpha = FactorAlpha(self)

        self.last_rebalance_time = None
        self.rebalance_period = TimeSpan.FromDays(14)  # 2 weeks
        self.strategy_returns = {
            self.straddle_alpha: [],
            self.condor_alpha: [],
            self.factor_alpha: []
        }
        self.min_strategy_weight = 0.10  # Minimum 10% allocation
        

        # Set Brokerage Model
        self.SetSecurityInitializer(MySecurityInitializer(
                    self.BrokerageModel, 
                    FuncSecuritySeeder(self.GetLastKnownPrice)
                ))
        
        # Set initial weights for each strat 
        self.alpha_weights = {
            self.straddle_alpha: 0.33,
            self.condor_alpha: 0.33,
            self.factor_alpha: 0.33,
        }

        # Track enabled status of strategies
        self.strategy_enabled = {
            self.straddle_alpha: True,
            self.condor_alpha: True,
            self.factor_alpha: True
        }

        self.profit_target = 1.5
        self.stop_loss = 0.75

        # Add these new tracking variables
        self.last_performance_check = None
        self.performance_window = 7
        self.performance_threshold = -0.01
        self.strategy_trades = {
            self.straddle_alpha: [],
            self.condor_alpha: [],
            self.factor_alpha: []
        }

        # Initialize strategy performance tracking
        self.strategy_performance = {
            self.straddle_alpha: 0,
            self.condor_alpha: 0,
            self.factor_alpha: 0
        }

        self.strategy_trade_counts = {
            self.straddle_alpha: 0,
            self.condor_alpha: 0,
            self.factor_alpha: 0
        }

        # Track trades and P&L for each strategy
        self.strategy_trades = {
            self.straddle_alpha: [],
            self.condor_alpha: [],
            self.factor_alpha: []
        }

        self.Log(f"[{self.Time}] Initialized CombinedOptionsAlpha with 2 strategies.")

    def OnData(self, slice):
        if self.IsWarmingUp:
            return

        self.ManagePositions()
    
        # Pass option chain data to alpha models
        if slice.OptionChains:
            if self.strategy_enabled[self.condor_alpha]:
                self.condor_alpha.OnOptionChainChanged(slice)
            if self.strategy_enabled[self.straddle_alpha]:
                self.straddle_alpha.OnOptionChainChanged(slice)
        
        if self.strategy_enabled[self.factor_alpha]:
            # Check if the universe has changed
            current_universe_symbols = set(self.symbols)
            if not current_universe_symbols:
                return  # Skip if no symbols are selected
            
            if not hasattr(self, 'last_universe_symbols'):
                self.last_universe_symbols = set()  # Initialize on the first run

            # Compare the current universe with the last known universe
            if current_universe_symbols != self.last_universe_symbols:
                self.last_universe_symbols = current_universe_symbols  # Update to the new universe
                self.ExecuteStrategyOrders(self.factor_alpha, slice)
                self.Debug("Executing FactorAlpha trades due to universe change.")
            else:
                self.Log("No universe change detected; skipping FactorAlpha execution.")


        current_time = self.Time

        # Straddle at 11:30
        if (current_time.hour == 11 and 
            current_time.minute == 30 and 
            self.strategy_enabled[self.straddle_alpha]):
            
            if self.straddle_alpha.ShouldTrade(slice):
                self.Log("Executing Straddle Strategy")
                self.ExecuteStrategyOrders(self.straddle_alpha, slice)

        # Iron Condor between 15:00 and 15:55
        if (current_time.hour == 15 and 
            0 <= current_time.minute <= 5):
            
            # self.Log(f"Checking Iron Condor at {current_time}")
            # self.Log(f"Strategy enabled: {self.strategy_enabled[self.condor_alpha]}")
            
            if self.strategy_enabled[self.condor_alpha]:
                # self.Log(f"Checking ShouldTrade for Iron Condor")
                # self.Log(f"Portfolio Invested: {self.Portfolio.Invested}")
                # self.Log(f"Kurtosis condition met: {self.condor_alpha.kurtosis_condition_met}")
                
                if self.condor_alpha.ShouldTrade(slice):
                    # self.Log("Iron Condor ShouldTrade returned True")
                    trade_orders = self.condor_alpha.GenerateOrders(slice)
                    if trade_orders:
                        # self.Log(f"Generated Iron Condor orders: {trade_orders}")
                        weight = self.alpha_weights[self.condor_alpha]
                        weighted_orders = self.WeightOrders(trade_orders, weight)
                        self.ExecuteOrders(weighted_orders)
                        self.Log(f"Executed Iron Condor orders with {weight} weight")
                    else:
                        self.Log("No valid Iron Condor orders generated")


    # def UpdateStrategyWeights(self):
    #     """Calculate new weights based on risk-adjusted returns."""
    #     try:
    #         # Calculate returns and volatility for each strategy
    #         strategy_metrics = {}
            
    #         for strategy in self.strategy_returns:
    #             returns = self.strategy_returns[strategy]
    #             if not returns:
    #                 strategy_metrics[strategy] = 0
    #                 continue
                    
    #             returns_array = np.array(returns)
    #             mean_return = np.mean(returns_array)
    #             std_return = np.std(returns_array) if len(returns_array) > 1 else 1
                
    #             # Calculate Sharpe ratio (using 0 as risk-free rate for simplicity)
    #             sharpe = mean_return / std_return if std_return != 0 else 0
    #             strategy_metrics[strategy] = max(sharpe, 0)  # Ensure non-negative
            
    #         # Calculate weights based on relative Sharpe ratios
    #         total_metric = sum(strategy_metrics.values())
            
    #         if total_metric == 0:
    #             # If no performance data, use equal weights
    #             new_weights = {strategy: 1/len(self.strategy_returns) 
    #                         for strategy in self.strategy_returns}
    #         else:
    #             # Calculate weights ensuring minimum allocation
    #             raw_weights = {strategy: metric/total_metric 
    #                         for strategy, metric in strategy_metrics.items()}
                
    #             # Adjust weights to ensure minimum allocation
    #             adjusted_weights = {}
    #             remaining_weight = 1.0
    #             strategies_below_min = []
                
    #             for strategy, weight in raw_weights.items():
    #                 if weight < self.min_strategy_weight:
    #                     adjusted_weights[strategy] = self.min_strategy_weight
    #                     remaining_weight -= self.min_strategy_weight
    #                     strategies_below_min.append(strategy)
                
    #             # Distribute remaining weight proportionally
    #             remaining_strategies = [s for s in raw_weights if s not in strategies_below_min]
    #             if remaining_strategies:
    #                 remaining_total = sum(raw_weights[s] for s in remaining_strategies)
    #                 for strategy in remaining_strategies:
    #                     if remaining_total > 0:
    #                         adjusted_weights[strategy] = (raw_weights[strategy] / remaining_total) * remaining_weight
    #                     else:
    #                         adjusted_weights[strategy] = remaining_weight / len(remaining_strategies)
                
    #             new_weights = adjusted_weights

    #         # Update weights
    #         self.alpha_weights = new_weights
            
    #         self.Log("Updated strategy weights:")
    #         for strategy, weight in self.alpha_weights.items():
    #             self.Log(f"{strategy.__class__.__name__}: {weight:.2%}")
            
    #         # Reset return tracking for next period
    #         self.strategy_returns = {strategy: [] for strategy in self.strategy_returns}
            
    #     except Exception as e:
    #         self.Error(f"Error updating strategy weights: {str(e)}")


    def UpdateStrategyWeights(self):
        """
        Calculate new weights with Factor model capped at 20% and remaining 80%
        distributed between Straddle and Condor based on performance.
        """
        try:
            # First, assign 20% to Factor Alpha
            self.alpha_weights[self.factor_alpha] = 0.20
            remaining_weight = 0.80
            
            # Calculate returns for Straddle and Condor only
            option_strategies = [self.straddle_alpha, self.condor_alpha]
            strategy_metrics = {}
            
            for strategy in option_strategies:
                returns = self.strategy_returns[strategy]
                if not returns:
                    strategy_metrics[strategy] = 0
                    continue
                
                avg_return = np.mean(returns)
                strategy_metrics[strategy] = max(avg_return, 0)  # Ensure non-negative
                
                self.Log(f"{strategy.__class__.__name__} Average Return: {avg_return:.2%}")

            # Calculate weights for option strategies
            total_return = sum(strategy_metrics.values())
            
            if total_return == 0:
                # If no returns data or all zero returns, split remaining weight equally
                for strategy in option_strategies:
                    self.alpha_weights[strategy] = remaining_weight / 2
                self.Log("No performance difference - splitting remaining 80% equally")
            else:
                # Calculate weights based on relative performance
                for strategy in option_strategies:
                    weight = (strategy_metrics[strategy] / total_return) * remaining_weight
                    # Ensure minimum 10% allocation
                    weight = max(weight, 0.10)
                    self.alpha_weights[strategy] = weight

            # Log weight changes
            self.Log("\n=== Weight Distribution Update ===")
            total_weight = sum(self.alpha_weights.values())
            self.Log(f"Total Weight: {total_weight:.2%}")
            for strategy, weight in self.alpha_weights.items():
                strategy_name = strategy.__class__.__name__
                allocated_capital = weight * self.Portfolio.TotalPortfolioValue
                self.Log(f"\n{strategy_name}:")
                self.Log(f"  Weight: {weight:.2%}")
                self.Log(f"  Allocated Capital: ${allocated_capital:,.2f}")
                if strategy in strategy_metrics:
                    self.Log(f"  Recent Return: {strategy_metrics.get(strategy, 0):.2%}")

            # Reset return tracking for next period
            self.strategy_returns = {strategy: [] for strategy in self.strategy_returns}
            
        except Exception as e:
            self.Error(f"Error updating strategy weights: {str(e)}")

    def OnOrderEvent(self, orderEvent):
        if orderEvent.Status == OrderStatus.Filled:
            for strategy in [self.straddle_alpha, self.condor_alpha, self.factor_alpha]:
                if strategy.trade_open:
                    # Calculate trade return
                    trade_return = (orderEvent.FillPrice * orderEvent.FillQuantity) / self.Portfolio.TotalPortfolioValue
                    self.strategy_returns[strategy].append(trade_return)
                    
                    # Check if it's time to rebalance
                    if (self.last_rebalance_time is None or 
                        self.Time - self.last_rebalance_time >= self.rebalance_period):
                        self.UpdateStrategyWeights()
                        self.last_rebalance_time = self.Time

                        
    def EvaluatePerformance(self, strategy):
        """Evaluate the performance of a strategy based on the last 10 trades."""
        trades = self.strategy_trades.get(strategy, [])
        if not trades:
            self.Log(f"No trades for {strategy.__class__.__name__}. Skipping evaluation.")
            return

        # Calculate total P&L from the trades
        total_pnl = sum(trade['pnl'] for trade in trades)
        performance = total_pnl / self.Portfolio.TotalPortfolioValue

        self.Log(f"Performance for {strategy.__class__.__name__}: {performance:.2%} over the last 10 trades")

        # Disable or enable strategies based on performance
        if performance < self.performance_threshold:
            self.Debug(f"Disabling {strategy.__class__.__name__} due to poor performance.")
        #     self.DisableStrategy(strategy)

        #     # Enable the alternate strategy
        #     if strategy == self.straddle_alpha:
        #         alternate_strategy = self.condor_alpha
        #     elif strategy == self.condor_alpha:
        #         alternate_strategy = self.factor_alpha
        #     elif strategy == self.factor_alpha:
        #         alternate_strategy = self.straddle_alpha
            
        #     self.EnableStrategy(alternate_strategy)
        # else:
        #     self.Log(f"{strategy.__class__.__name__} performance is acceptable. Strategy remains active.")


    def CheckWeeklyPerformance(self):
        if self.last_performance_check is None:
            self.last_performance_check = self.Time
            return

        # Check if a week has passed
        if (self.Time - self.last_performance_check) >= 15:
            self.Log("Performing weekly performance evaluation")
            
            # Calculate performance for each strategy
            for strategy in [self.straddle_alpha, self.condor_alpha, self.factor_alpha]:
                trades = self.strategy_trades[strategy]
                if trades:
                    total_pnl = sum(trade['pnl'] for trade in trades)
                    performance = total_pnl / self.Portfolio.TotalPortfolioValue
                    
                    # Disable the strategy if performance is below threshold
                    if performance < self.performance_threshold:
                        self.Log(f"Disabling {strategy.__class__.__name__} due to poor performance: {performance:.2%}")
                        self.DisableStrategy(strategy)
                        
                        # Enable the alternate strategy
                        if strategy == self.straddle_alpha:
                            alternate_strategy = self.condor_alpha
                        elif strategy == self.condor_alpha:
                            alternate_strategy = self.factor_alpha
                        elif strategy == self.factor_alpha:
                            alternate_strategy = self.straddle_alpha

                        self.EnableStrategy(alternate_strategy)
                    else:
                        # Enable the strategy if performance is acceptable
                        if not self.strategy_enabled[strategy]:
                            self.Log(f"Enabling {strategy.__class__.__name__}: {performance:.2%}")
                            self.EnableStrategy(strategy)

            # Reset tracking
            self.last_performance_check = self.Time
            self.strategy_trades = {
                self.straddle_alpha: [],
                self.condor_alpha: [],
                self.factor_alpha: []
            }

    def DisableStrategy(self, strategy):
        """Disable a strategy and liquidate its positions."""
        self.strategy_enabled[strategy] = False
        self.alpha_weights[strategy] = 0
        if strategy.trade_open:
            self.Liquidate()
            strategy.trade_open = False
        self.Log(f"Disabled {strategy.__class__.__name__}")

    def EnableStrategy(self, strategy):
        """Enable a strategy and restore its weight."""
        self.strategy_enabled[strategy] = True
        self.alpha_weights[strategy] = 1  # Restore default weight (adjust as needed)
        self.Log(f"Enabled {strategy.__class__.__name__}")


    def ExecuteStrategyOrders(self, strategy, slice):
        """Execute orders for a specific strategy with weight applied"""
        if strategy == self.factor_alpha:
            # Call GenerateOrders instead of ExecuteOrders
            trade_orders = strategy.GenerateOrders(slice)
            if trade_orders:
                weight = self.alpha_weights[strategy]
                weighted_orders = self.WeightOrders(trade_orders, weight)
                self.ExecuteOrders(weighted_orders)
                self.Log(f"Executed Factor Alpha orders with {weight} weight")
        else:
            trade_orders = strategy.GenerateOrders(slice)
            if trade_orders:
                weight = self.alpha_weights[strategy]
                weighted_orders = self.WeightOrders(trade_orders, weight)
                self.ExecuteOrders(weighted_orders)

    def WeightOrders(self, orders, weight):
        """Apply strategy weight to order quantities"""
        weighted_orders = []
        for order in orders:
            if len(order) == 2:  # Iron Condor case: (strategy, quantity)
                strategy, quantity = order
                weighted_quantity = max(1, int(quantity * weight))  # Ensure minimum 1 contract
                weighted_orders.append((strategy, weighted_quantity))
            else:  # Straddle case: (symbol, quantity, is_buy)
                symbol, quantity, is_buy = order
                weighted_quantity = max(1, int(quantity * weight))  # Ensure minimum 1 contract
                weighted_orders.append((symbol, weighted_quantity, is_buy))
        return weighted_orders

    def ExecuteOrders(self, orders):
        """Execute the weighted orders"""
        for order in orders:
            try:
                if len(order) == 2:  # Iron Condor case
                    strategy, quantity = order
                    self.Buy(strategy, quantity)
                    self.Log(f"Executing Iron Condor order: {quantity} contracts")
                else:  # Straddle case
                    symbol, quantity, is_buy = order
                    if is_buy:
                        self.Buy(symbol, quantity)
                        self.Log(f"Buying {quantity} of {symbol}")
                    else:
                        self.Sell(symbol, quantity)
                        self.Log(f"Selling {quantity} of {symbol}")
            except Exception as e:
                self.Error(f"Order execution failed: {str(e)}")

    def ManagePositions(self):
        """Comprehensive position management for all strategies with percentage-based stop-loss and take-profit."""
        if not self.Portfolio.Invested:
            return

        if (self.last_rebalance_time is None or 
            self.Time - self.last_rebalance_time >= self.rebalance_period):
            self.UpdateStrategyWeights()
            self.last_rebalance_time = self.Time
            
            # Print detailed weight information
            self.Debug("=== Strategy Weight Update ===")
            self.Debug(f"Time: {self.Time}")
            for strategy, weight in self.alpha_weights.items():
                strategy_name = strategy.__class__.__name__
                allocated_capital = weight * self.Portfolio.TotalPortfolioValue
                self.Debug(f"{strategy_name}:")
                self.Debug(f"  Weight: {weight:.2%}")
                self.Debug(f"  Allocated Capital: ${allocated_capital:,.2f}")
                
                # Print recent returns if available
                if self.strategy_returns[strategy]:
                    avg_return = np.mean(self.strategy_returns[strategy])
                    self.Debug(f"  Average Return: {avg_return:.2%}")
            
            self.Debug(f"Total Portfolio Value: ${self.Portfolio.TotalPortfolioValue:,.2f}")
            self.Debug("===========================")

        # Manage StraddleAlpha positions
        if self.strategy_enabled[self.straddle_alpha] and self.straddle_alpha.trade_open:
            straddle_symbols = [self.straddle_alpha.spy_symbol, self.straddle_alpha.option_symbol]
            for symbol in straddle_symbols:
                holding = self.Portfolio[symbol]
                if not holding.Invested:
                    continue

                entry_price = holding.AveragePrice
                current_price = holding.Price
                pnl_percentage = (current_price - entry_price) / entry_price

                # self.Debug(f"StraddleAlpha {symbol}: Entry={entry_price:.2f}, Current={current_price:.2f}, PnL%={pnl_percentage:.2%}")

                if pnl_percentage >= self.profit_target:
                    self.Liquidate(symbol)
                    self.Debug(f"Liquidated {symbol} due to StraddleAlpha profit target reached")
                    self.straddle_alpha.trade_open = False
                elif pnl_percentage <= -self.stop_loss:
                    self.Liquidate(symbol)
                    self.Debug(f"Liquidated {symbol} due to StraddleAlpha stop loss reached")
                    self.straddle_alpha.trade_open = False

        # Manage IronCondorAlpha positions
        if self.strategy_enabled[self.condor_alpha] and self.condor_alpha.trade_open:
            condor_symbols = [self.condor_alpha._symbol1, self.condor_alpha._symbol2]
            for symbol in condor_symbols:
                holding = self.Portfolio[symbol]
                if not holding.Invested:
                    continue

                entry_price = holding.AveragePrice
                current_price = holding.Price
                pnl_percentage = (current_price - entry_price) / entry_price

                # self.Debug(f"IronCondorAlpha {symbol}: Entry={entry_price:.2f}, Current={current_price:.2f}, PnL%={pnl_percentage:.2%}")

                if pnl_percentage >= self.profit_target:
                    self.Liquidate(symbol)
                    self.Debug(f"Liquidated {symbol} due to IronCondorAlpha profit target reached")
                    self.condor_alpha.trade_open = False
                elif pnl_percentage <= -self.stop_loss:
                    self.Liquidate(symbol)
                    self.Debug(f"Liquidated {symbol} due to IronCondorAlpha stop loss reached")
                    self.condor_alpha.trade_open = False

        # Manage FactorAlpha positions
        if self.strategy_enabled[self.factor_alpha] and self.factor_alpha.trade_open:
            factor_symbols = self.symbols if hasattr(self, 'symbols') else []
            for symbol in factor_symbols:
                holding = self.Portfolio[symbol]
                if not holding.Invested:
                    continue

                entry_price = holding.AveragePrice
                current_price = holding.Price
                pnl_percentage = (current_price - entry_price) / entry_price

                # self.Debug(f"FactorAlpha {symbol}: Entry={entry_price:.2f}, Current={current_price:.2f}, PnL%={pnl_percentage:.2%}")

                if pnl_percentage >= self.profit_target:
                    self.Liquidate(symbol)
                    self.Debug(f"Liquidated {symbol} due to FactorAlpha profit target reached")
                elif pnl_percentage <= -self.stop_loss:
                    self.Liquidate(symbol)
                    self.Debug(f"Liquidated {symbol} due to FactorAlpha stop loss reached")

        self.Debug("Position management completed for all strategies.")
        
class FactorAlpha:
    def __init__(self, algorithm):
        self.algorithm = algorithm
        self.Initialize()

    def Initialize(self):
        self.algorithm.Log("Initializing FactorAlpha")
        self.algorithm.UniverseSettings.Resolution = Resolution.Daily
        self.algorithm.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)

        self.num_stocks = 500
        self.trade_open = False
        self.num_groups = 10
        self.current_month = -1
        self.model = None
        self.last_month_features = pd.DataFrame()
        self.label_encoder = LabelEncoder()
        self.predicted_stocks = {'long': [], 'short': []}
        self.position_size = 0.1  # 10% of portfolio per position

        # Initialize XGBoost model
        self.model = XGBClassifier(
            n_estimators=100,
            learning_rate=0.1,
            max_depth=5,
            random_state=42,
            objective='multi:softprob'
        )
        self.algorithm.Log("FactorAlpha initialization complete")

    def CoarseSelectionFunction(self, coarse):
        # self.algorithm.Log(f"Running CoarseSelectionFunction at {self.algorithm.Time}")
        
        if self.algorithm.Time.month == self.current_month:
            # self.algorithm.Log("Same month - returning unchanged universe")
            return Universe.Unchanged
        
        self.current_month = self.algorithm.Time.month
        
        try:
            sorted_by_volume = sorted(
                [x for x in coarse if x.HasFundamentalData], 
                key=lambda x: x.Market, 
                reverse=True
            )
            self.algorithm.Debug(f"Found {len(sorted_by_volume)} stocks with fundamental data")
            
            selected_symbols = [x.Symbol for x in sorted_by_volume[:self.num_stocks]]
            self.algorithm.Debug(f"Selected {len(selected_symbols)} symbols in coarse selection")
            return selected_symbols
            
        except Exception as e:
            # self.algorithm.Error(f"Error in CoarseSelectionFunction: {str(e)}")
            return []

    def FineSelectionFunction(self, fine):
        self.algorithm.Log(f"Running FineSelectionFunction at {self.algorithm.Time}")
        
        fine_list = list(fine)
        if not fine_list:
            self.algorithm.Log("Empty fine data received")
            return []

        try:
            current_month_features = pd.DataFrame()
            current_month_returns = pd.DataFrame()

            for stock in fine_list:
                try:
                    symbol = str(stock.Symbol)
                    
                    # Get historical data
                    history = self.algorithm.History(stock.Symbol, 20, Resolution.Daily)
                    if len(history) < 20:
                        continue

                    # Calculate features
                    daily_returns = history['close'].pct_change().dropna()
                    volatility = daily_returns.std() * np.sqrt(252)
                    momentum = stock.ValuationRatios.PriceChange1M
                    
                    # Value calculation
                    if stock.ValuationRatios.PERatio > 0 and stock.ValuationRatios.PERatio < 100:
                        value = 1 / stock.ValuationRatios.PERatio
                    else:
                        continue

                    size = np.log(stock.MarketCap) if stock.MarketCap > 0 else np.nan
                    quality = stock.OperationRatios.ROE.Value
                    pb = stock.ValuationRatios.PBRatio
                    margin = stock.OperationRatios.GrossMargin.OneMonth

                    # Store features
                    current_month_features.loc[symbol, 'Momentum'] = momentum
                    current_month_features.loc[symbol, 'Value'] = value
                    current_month_features.loc[symbol, 'Size'] = size
                    current_month_features.loc[symbol, 'Quality'] = quality
                    current_month_features.loc[symbol, 'Volatility'] = volatility
                    current_month_features.loc[symbol, 'PB'] = pb
                    current_month_features.loc[symbol, 'Margin'] = margin

                    # Calculate returns
                    first_price = history['close'].iloc[0]
                    last_price = history['close'].iloc[-1]
                    log_return = np.log(last_price / first_price)   
                    current_month_returns.loc[symbol, 'Returns'] = log_return

                except Exception as e:
                    self.algorithm.Log(f"Error processing individual stock {symbol}: {str(e)}")
                    continue

            if current_month_features.empty:
                self.algorithm.Log("No features collected this month")
                return []

            if self.last_month_features.empty:
                self.algorithm.Log("Storing first month's features")
                self.last_month_features = current_month_features
                return []

            self.algorithm.Log("Training model with previous month's data")
            
            # Prepare training data
            X_train = self.last_month_features
            y_train = current_month_returns
            common_symbols = X_train.index.intersection(y_train.index)
            X_train = X_train.loc[common_symbols]
            y_train = y_train.loc[common_symbols]

            # Process features
            X_train = X_train.fillna(X_train.median())
            y_classes = pd.qcut(y_train['Returns'], q=self.num_groups, labels=False)
            
            scaler = StandardScaler()
            X_train_scaled = scaler.fit_transform(X_train)

            # Train model
            self.model.fit(X_train_scaled, y_classes)
            self.algorithm.Log("Model training completed")

            # Make predictions
            predictions = self.PredictGroups(current_month_features)
            if predictions.empty:
                self.algorithm.Log("No predictions generated")
                return []

            # Update predicted stocks for trading
            self.predicted_stocks['long'] = list(predictions[predictions['predicted_group'] == self.num_groups - 1].index)
            self.predicted_stocks['short'] = list(predictions[predictions['predicted_group'] == 0].index)

            self.algorithm.Log(f"Selected {len(self.predicted_stocks['long'])} long and {len(self.predicted_stocks['short'])} short positions")

            # Convert string symbols back to Symbol objects
            selected_symbols = []
            for symbol_str in self.predicted_stocks['long'] + self.predicted_stocks['short']:
                for stock in fine_list:
                    if str(stock.Symbol) == symbol_str:
                        selected_symbols.append(stock.Symbol)
                        break

            self.algorithm.symbols = selected_symbols
            self.last_month_features = current_month_features
            return selected_symbols

        except Exception as e:
            self.algorithm.Error(f"Error in FineSelectionFunction: {str(e)}")
            return []

    def PredictGroups(self, features):
        self.algorithm.Log("Making predictions for current month")
        try:
            features = features.fillna(features.mean())
            scaler = StandardScaler()
            features_scaled = scaler.fit_transform(features)

            class_probs = self.model.predict_proba(features_scaled)
            predicted_classes = np.argmax(class_probs, axis=1)

            predictions = pd.DataFrame({
                'predicted_group': predicted_classes,
                'confidence': np.max(class_probs, axis=1)
            }, index=features.index)

            self.algorithm.Log(f"Generated predictions for {len(predictions)} stocks")
            return predictions
        except Exception as e:
            self.algorithm.Error(f"Error in PredictGroups: {str(e)}")
            return pd.DataFrame()

    def GenerateOrders(self, slice):
        """Generate orders based on predictions within the allocated capital."""
        self.algorithm.Log("Generating orders for Factor Alpha")
        
        if not hasattr(self.algorithm, 'symbols') or not self.algorithm.symbols:
            self.algorithm.Log("No symbols available for trading")
            return []
        
        try:
            # Get total allocated capital for the strategy
            allocated_capital = self.algorithm.alpha_weights[self] * self.algorithm.Portfolio.TotalPortfolioValue
            
            if allocated_capital <= 0:
                self.algorithm.Log("Capital allocation is zero; no trades will be executed")
                return []
            
            # Determine the total number of positions (long + short)
            total_positions = len(self.predicted_stocks['long']) + len(self.predicted_stocks['short'])
            if total_positions == 0:
                self.algorithm.Log("No positions to allocate; skipping orders")
                return []
            
            # Allocate equal capital to each position
            position_value = allocated_capital / total_positions
            self.algorithm.debug(f"Position Value:{position_value}")
            orders = []

            # Process long positions
            for symbol_str in self.predicted_stocks['long']:
                symbol = None
                for s in self.algorithm.symbols:
                    if str(s) == symbol_str:
                        symbol = s
                        break
                
                if symbol is None:
                    continue
                
                # Get current price
                security = self.algorithm.Securities[symbol]
                if security.Price == 0:
                    continue
                
                # Calculate position size
                quantity = int(position_value / security.Price)
                if quantity > 0:
                    orders.append((symbol, quantity, True))  # True for buy
                    self.algorithm.Log(f"Generated long order for {symbol}: {quantity} shares")
            
            # Process short positions
            for symbol_str in self.predicted_stocks['short']:
                symbol = None
                for s in self.algorithm.symbols:
                    if str(s) == symbol_str:
                        symbol = s
                        break
                
                if symbol is None:
                    continue
                
                # Get current price
                security = self.algorithm.Securities[symbol]
                if security.Price == 0:
                    continue
                
                # Calculate position size
                quantity = int(position_value / security.Price)
                if quantity > 0:
                    orders.append((symbol, quantity, False))  # False for sell/short
                    self.algorithm.Log(f"Generated short order for {symbol}: {quantity} shares")
            
            self.trade_open = True
            return orders

        except Exception as e:
            self.algorithm.Error(f"Error generating orders in FactorAlpha: {str(e)}")
            return []

class IronCondorAlpha:
    def __init__(self, algorithm):
        self.algorithm = algorithm  # Store reference to main algorithm
        self.Initialize()

    def Initialize(self):
        # Add SPX index
        self.index = self.algorithm.AddIndex("SPX")

        # Universe 1 (option1): Wide filter for kurtosis calculations
        self.option1 = self.algorithm.AddIndexOption(self.index.Symbol, "SPXW")
        self.option1.SetFilter(lambda universe: universe.IncludeWeeklys()
                             .Strikes(-30,30).Expiration(0, 0))
        self._symbol1 = self.option1.Symbol

        # Universe 2 (option2): Iron Condor filter for placing trades
        self.option2 = self.algorithm.AddIndexOption(self.index.Symbol, "SPXW")
        self.option2.SetFilter(lambda x: x.IncludeWeeklys().IronCondor(0, 20, 40))
        self._symbol2 = self.option2.Symbol

        # Risk and trade management parameters
        self.max_portfolio_risk = 0.05
        self.profit_target = 1.5
        self.stop_loss = 0.5
        self.trade_open = False
        self.initial_credit = 0
        self.max_potential_loss = 0
        self.target_delta = 0.25

        self.kurtosis_threshold = 2  # Changed to match original
        self.current_date = None
        self.kurtosis_condition_met = False
        self.computed_kurtosis_today = False

    def OnOptionChainChanged(self, slice):
        # Check if a new day has started
        if self.current_date != self.algorithm.Time.date():
            self.current_date = self.algorithm.Time.date()
            self.trade_open = False
            self.kurtosis_condition_met = False
            self.computed_kurtosis_today = False
            self.algorithm.Log(f"New day reset for Iron Condor at {self.algorithm.Time}")

        # Compute kurtosis at 9:31-9:36 AM
        if (not self.computed_kurtosis_today and 
            self.algorithm.Time.hour == 9 and 
            self.algorithm.Time.minute >= 31 and 
            self.algorithm.Time.minute <= 36):
            
            chain1 = slice.OptionChains.get(self._symbol1)
            if chain1:
                iv_values = [x.ImpliedVolatility for x in chain1 
                           if x.ImpliedVolatility and 0 < x.ImpliedVolatility < 5]
                if len(iv_values) > 10:  # Using 10 as in original
                    daily_kurtosis = kurtosis(iv_values)
                    self.algorithm.Log(f"Iron Condor Kurtosis: {daily_kurtosis}")
                    if daily_kurtosis > self.kurtosis_threshold:
                        self.kurtosis_condition_met = True
                        self.algorithm.Log("Iron Condor Kurtosis condition met")
                    self.computed_kurtosis_today = True

    def ShouldTrade(self, slice):
        # Only check if we should trade based on conditions, not time
        return (not self.algorithm.Portfolio.Invested and 
                self.kurtosis_condition_met)

    def GenerateOrders(self, slice):
        chain2 = slice.OptionChains.get(self._symbol2)

        if not chain2:
            return None

        expiry = max([x.Expiry for x in chain2])
        chain2 = sorted([x for x in chain2 if x.Expiry == expiry], 
                        key=lambda x: x.Strike)

        put_contracts = [x for x in chain2 
                        if x.Right == OptionRight.PUT and 
                        abs(x.Greeks.Delta) <= self.target_delta]
        call_contracts = [x for x in chain2 
                        if x.Right == OptionRight.CALL and 
                        abs(x.Greeks.Delta) <= self.target_delta]

        if len(call_contracts) < 2 or len(put_contracts) < 2:
            return None

        near_call = min(call_contracts, 
                        key=lambda x: abs(x.Greeks.Delta - self.target_delta))
        far_call = min([x for x in call_contracts if x.Strike > near_call.Strike], 
                    key=lambda x: abs(x.Greeks.Delta - self.target_delta))
        
        near_put = min(put_contracts, 
                    key=lambda x: abs(x.Greeks.Delta + self.target_delta))
        far_put = min([x for x in put_contracts if x.Strike < near_put.Strike], 
                    key=lambda x: abs(x.Greeks.Delta + self.target_delta))

        credit = (near_call.BidPrice - far_call.AskPrice) + (near_put.BidPrice - far_put.AskPrice)
        spread_width = max(far_call.Strike - near_call.Strike, 
                        near_put.Strike - far_put.Strike)
        max_potential_loss = spread_width * 100 - credit * 100

        # Use allocated capital instead of total portfolio value
        allocated_capital = self.algorithm.alpha_weights[self] * self.algorithm.Portfolio.TotalPortfolioValue
        max_trade_risk = allocated_capital * self.max_portfolio_risk
        contracts = int(max_trade_risk / max_potential_loss)

        if contracts > 0:
            iron_condor = OptionStrategies.IronCondor(
                self._symbol2,
                far_put.Strike,
                near_put.Strike,
                near_call.Strike,
                far_call.Strike,
                expiry
            )

            # Store trade parameters for position management
            self.initial_credit = credit * 100 * contracts
            self.max_potential_loss = max_potential_loss * contracts
            self.trade_open = True

            self.algorithm.Log(f"Generated iron condor at {self.algorithm.Time}, "
                            f"Contracts: {contracts}, Credit: ${self.initial_credit:.2f}")

            return [(iron_condor, contracts)]

        return None


class DeltaHedgedStraddleAlpha:
    def __init__(self, algorithm):
        self.algorithm = algorithm
        self.Initialize()

    def Initialize(self):
        # Add SPX index
        self.index = self.algorithm.AddIndex("SPX")

        # Add SPY for Delta Hedging
        self.spy = self.algorithm.AddEquity("SPY", Resolution.Minute)
        self.spy.SetLeverage(1)
        self.spy.SetDataNormalizationMode(DataNormalizationMode.Raw)
        self.spy_symbol = self.spy.Symbol

        # Add SPX options
        self.option = self.algorithm.AddIndexOption(self.index.Symbol, "SPXW")
        self.option.SetFilter(lambda universe: universe.IncludeWeeklys()
                            .Strikes(-30, 30).Expiration(0, 0))
        self.option_symbol = self.option.Symbol

        # Risk and trade management parameters
        self.max_portfolio_risk = 0.05
        self.profit_target = 1.5
        self.stop_loss = 0.75
        self.trade_open = False

        # Kurtosis calculation variables
        self.kurtosis_threshold = 0
        self.kurtosis_condition_met = False
        self.computed_kurtosis_today = False
        self.current_date = None

        # Variables for delta hedging
        self.hedge_order_ticket = None
        self.net_delta = 0.0
        self.max_potential_loss = 0.0

    def OnOptionChainChanged(self, slice):
        # Check if a new day has started
        if self.current_date != self.algorithm.Time.date():
            self.current_date = self.algorithm.Time.date()
            self.trade_open = False
            self.kurtosis_condition_met = False
            self.computed_kurtosis_today = False
            self.algorithm.Log(f"New day reset for Straddle at {self.algorithm.Time}")

            # Liquidate any existing hedge at the start of a new day
            if self.hedge_order_ticket and self.hedge_order_ticket.Status not in [OrderStatus.Filled, OrderStatus.Canceled]:
                self.algorithm.CancelOrder(self.hedge_order_ticket.OrderId)
            self.algorithm.Liquidate(self.spy_symbol)
            self.algorithm.Liquidate(self.option_symbol)

        # Compute kurtosis from option chain at 9:31-9:36 AM
        if (not self.computed_kurtosis_today and 
            self.algorithm.Time.hour == 9 and 
            self.algorithm.Time.minute >= 31 and 
            self.algorithm.Time.minute <= 36):
            
            chain = slice.OptionChains.get(self.option_symbol)
            if chain:
                iv_values = [x.ImpliedVolatility for x in chain 
                           if x.ImpliedVolatility and 0 < x.ImpliedVolatility < 5]
                if len(iv_values) > 3:
                    daily_kurtosis = kurtosis(iv_values)
                    if daily_kurtosis > self.kurtosis_threshold:
                        self.kurtosis_condition_met = True
                        self.algorithm.Log(f"Straddle Kurtosis met: {daily_kurtosis}")
                    self.computed_kurtosis_today = True

    def ShouldTrade(self, slice):
        return (not self.trade_open and 
                self.kurtosis_condition_met)

    def GenerateOrders(self, slice):
        chain = slice.OptionChains.get(self.option_symbol)

        if not chain:
            return None

        # Find ATM strike
        atm_strike = self.index.Price
        closest_option = min(chain, key=lambda x: abs(x.Strike - atm_strike))
        atm_strike = closest_option.Strike

        # Filter for ATM call and put contracts with the highest Vega
        atm_call_candidates = [x for x in chain 
                                if x.Strike == atm_strike and 
                                x.Right == OptionRight.CALL]
        atm_put_candidates = [x for x in chain 
                            if x.Strike == atm_strike and 
                            x.Right == OptionRight.PUT]

        if not atm_call_candidates or not atm_put_candidates:
            return None

        # Select contracts with highest Vega
        atm_call = max(atm_call_candidates, key=lambda x: x.Greeks.Vega)
        atm_put = max(atm_put_candidates, key=lambda x: x.Greeks.Vega)

        # Calculate credit received from selling the straddle
        credit = atm_call.BidPrice + atm_put.BidPrice
        max_loss = abs(atm_call.Strike - self.index.Price) * 100 + credit * 100

        if max_loss <= 0:
            return None

        # Use allocated capital instead of total portfolio value
        allocated_capital = self.algorithm.alpha_weights[self] * self.algorithm.Portfolio.TotalPortfolioValue
        max_trade_risk = allocated_capital * self.max_portfolio_risk
        contracts = int(max_trade_risk / max_loss)

        if contracts <= 0:
            return None

        # Calculate delta hedge - Converting SPX delta to SPY (dividing by 10)
        net_delta = (atm_call.Greeks.Delta + atm_put.Greeks.Delta) * contracts
        required_spy_position = int(-net_delta * 10)  # SPX/SPY ratio is roughly 10:1

        # Store trade parameters
        self.trade_open = True
        self.max_potential_loss = max_loss * contracts
        self.net_delta = net_delta

        # Log the hedge calculation
        self.algorithm.Log(f"Delta calculation: Call Delta={atm_call.Greeks.Delta}, "
                            f"Put Delta={atm_put.Greeks.Delta}, "
                            f"Contracts={contracts}, "
                            f"Net Delta={net_delta}, "
                            f"Required SPY Position={required_spy_position}")

        # Return orders as tuples: (symbol, quantity, is_buy)
        orders = [
            (atm_call.Symbol, contracts, False),  # Sell call
            (atm_put.Symbol, contracts, False),   # Sell put
            (self.spy_symbol, abs(required_spy_position), required_spy_position > 0)  # Hedge
        ]

        self.algorithm.Log(f"Generated straddle orders: Straddle Contracts={contracts}, "
                        f"Delta Hedge Size={abs(required_spy_position)}, "
                        f"Hedge Direction={'Long' if required_spy_position > 0 else 'Short'}")

        return orders