from AlgorithmImports import *
from datetime import datetime
import math
from scipy.stats import kurtosis

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, 16)
        self.SetEndDate(2024, 12, 1)
        self.SetCash(2000000)
        self.SetTimeZone(TimeZones.NewYork)

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

        # Set Brokerage Model
        self.SetSecurityInitializer(MySecurityInitializer(
                    self.BrokerageModel, 
                    FuncSecuritySeeder(self.GetLastKnownPrice)
                ))
        
        # Set initial weights for each alpha (matching original 60-40 split)
        self.alpha_weights = {
            self.straddle_alpha: 1,
            self.condor_alpha: 0,
        }

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


        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: []
        }


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

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

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

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

    def OnData(self, slice):
        self.ManagePositions()
        # self.CheckWeeklyPerformance()
    
        # 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)

        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 OnOrderEvent(self, orderEvent):
        """Track trades and evaluate performance every 10 trades."""
        if orderEvent.Status == OrderStatus.Filled:
            # Determine which strategy the order belongs to
            for strategy in [self.straddle_alpha, self.condor_alpha]:
                if strategy.trade_open:
                    # Append trade details
                    self.strategy_trades[strategy].append({
                        'time': self.Time,
                        'pnl': orderEvent.FillPrice * orderEvent.FillQuantity
                    })

                    # Increment trade count
                    self.strategy_trade_counts[strategy] += 1
                    self.Log(f"{strategy.__class__.__name__} trade count: {self.strategy_trade_counts[strategy]}")

                    # Check performance after every 10 trades
                    if self.strategy_trade_counts[strategy] >= 8:
                        self.EvaluatePerformance(strategy)
                        # Reset trade count and trades after evaluation
                        self.strategy_trade_counts[strategy] = 0
                        self.strategy_trades[strategy] = []

    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.Log(f"Disabling {strategy.__class__.__name__} due to poor performance.")
            self.DisableStrategy(strategy)

            # Enable the alternate strategy
            alternate_strategy = self.straddle_alpha if strategy == self.condor_alpha else self.condor_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]:
                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
                        alternate_strategy = self.straddle_alpha if strategy == self.condor_alpha else self.condor_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: []
            }

    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"""
        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 {strategy.__class__.__name__} orders with {weight} weight")

    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):
        """Centralized position management for both strategies"""
        if not self.Portfolio.Invested:
            return

        total_pnl = sum([holding.UnrealizedProfit 
                        for holding in self.Portfolio.Values 
                        if holding.Invested])

        # For each strategy, check if its positions need management
        for alpha in [self.straddle_alpha, self.condor_alpha]:
            if hasattr(alpha, 'trade_open') and alpha.trade_open:
                if hasattr(alpha, 'initial_credit'):  # Iron Condor case
                    if total_pnl >= alpha.initial_credit * self.profit_target:
                        self.Liquidate()
                        alpha.trade_open = False
                        self.Log(f"Closed position at profit target on {self.Time}")
                    elif total_pnl <= -alpha.max_potential_loss * self.stop_loss:
                        self.Liquidate()
                        alpha.trade_open = False
                        self.Log(f"Closed position at stop loss on {self.Time}")
                
                elif hasattr(alpha, 'max_potential_loss'):  # Straddle case
                    if total_pnl >= alpha.max_potential_loss * self.profit_target:
                        self.Liquidate()
                        alpha.trade_open = False
                    elif total_pnl <= -alpha.max_potential_loss * self.stop_loss:
                        self.Liquidate()
                        alpha.trade_open = False
                
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.75
        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

        total_portfolio_value = self.algorithm.Portfolio.TotalPortfolioValue
        max_trade_risk = total_portfolio_value * 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

        # Position size calculation
        total_portfolio_value = self.algorithm.Portfolio.TotalPortfolioValue
        max_trade_risk = total_portfolio_value * 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)  # Multiply by 10 as 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