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(2022, 5, 25)
        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: 0.6,
            self.condor_alpha: 0.4,
        }

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

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

        self.profit_target = 1.5
        self.stop_loss = 0.75

    def OnData(self, slice):
        self.ManagePositions()

        # Check for strategy switch on day 20
        if self.Time.day == 20:
            if not getattr(self, 'switch_executed', False):  # Ensure we only switch once
                self.Debug("Switching strategies on day 20")
                self.strategy_enabled[self.condor_alpha] = True
                self.strategy_enabled[self.straddle_alpha] = False
                
                self.alpha_weights[self.condor_alpha] = 1
                self.alpha_weights[self.straddle_alpha] = 0
                
                # Liquidate any existing positions
                self.Liquidate()
                
                # Mark switch as executed
                self.switch_executed = True
                
                self.Debug(f"Strategy switch completed. Condor enabled: {self.strategy_enabled[self.condor_alpha]}, "
                        f"Straddle enabled: {self.strategy_enabled[self.straddle_alpha]}")
        
    
        # 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.Debug("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 <= 55):
            
            self.Debug(f"Checking Iron Condor at {current_time}")
            self.Debug(f"Strategy enabled: {self.strategy_enabled[self.condor_alpha]}")
            
            if self.strategy_enabled[self.condor_alpha]:
                self.Debug(f"Checking ShouldTrade for Iron Condor")
                self.Debug(f"Portfolio Invested: {self.Portfolio.Invested}")
                self.Debug(f"Kurtosis condition met: {self.condor_alpha.kurtosis_condition_met}")
                
                if self.condor_alpha.ShouldTrade(slice):
                    self.Debug("Iron Condor ShouldTrade returned True")
                    trade_orders = self.condor_alpha.GenerateOrders(slice)
                    if trade_orders:
                        self.Debug(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.Debug(f"Executed Iron Condor orders with {weight} weight")
                    else:
                        self.Debug("No valid Iron Condor orders generated")

    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.Debug(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.Debug(f"Executing Iron Condor order: {quantity} contracts")
                else:  # Straddle case
                    symbol, quantity, is_buy = order
                    if is_buy:
                        self.Buy(symbol, quantity)
                        self.Debug(f"Buying {quantity} of {symbol}")
                    else:
                        self.Sell(symbol, quantity)
                        self.Debug(f"Selling {quantity} of {symbol}")
            except Exception as e:
                self.Error(f"Order execution failed: {str(e)}")

    def CheckAndUpdateWeights(self):
        """Check and update strategy weights based on various conditions"""
        # Example: Adjust weights on the 15th of each month
        if self.Time.day == 15:
            if self.Time.hour == 0 and self.Time.minute == 0:  # Beginning of day
                self.UpdateStrategyWeights({
                    self.straddle_alpha: 0.7,  # Increase straddle allocation
                    self.condor_alpha: 0.3     # Decrease condor allocation
                })
                self.Log("Monthly weight rebalancing executed")

        # Example: Adjust weights based on VIX level (you could add this)
        # if self.Time.hour == 9 and self.Time.minute == 31:  # After market open
        #     if some_market_condition:
        #         self.UpdateStrategyWeights({...})

    def UpdateStrategyWeights(self, new_weights):
        """Update strategy weights and log changes"""
        old_weights = self.alpha_weights.copy()
        self.alpha_weights.update(new_weights)
        
        for strategy, new_weight in new_weights.items():
            old_weight = old_weights[strategy]
            self.Log(f"Updated {strategy.__class__.__name__} weight: {old_weight:.2f} -> {new_weight:.2f}")
            
        # Verify weights sum to 1
        total_weight = sum(self.alpha_weights.values())
        if not 0.99 <= total_weight <= 1.01:
            self.Warning(f"Strategy weights sum to {total_weight}, should be close to 1.0")

    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.Debug(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.Debug(f"Iron Condor Kurtosis: {daily_kurtosis}")
                    if daily_kurtosis > self.kurtosis_threshold:
                        self.kurtosis_condition_met = True
                        self.algorithm.Debug("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.Debug(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.Debug(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.Debug(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

        # Debug the hedge calculation
        self.algorithm.Debug(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.Debug(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