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

class DeltaHedgedStraddleWithVegaFiltering(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2022, 5, 16)
        self.SetEndDate(2022, 12, 1)
        self.SetCash(2000000)
        self.SetTimeZone(TimeZones.NewYork)

        # Add SPX index
        self.index = self.AddIndex("SPX")

        # Add SPY for Delta Hedging
        self.spy = self.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.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.025  # Max 5% of portfolio risked in any trade
        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 OnData(self, slice):
        # Check if a new day has started
        if self.current_date != self.Time.date():
            self.current_date = self.Time.date()
            self.trade_open = False
            self.kurtosis_condition_met = False
            self.computed_kurtosis_today = False

            # 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.CancelOrder(self.hedge_order_ticket.OrderId)
            self.Liquidate(self.spy_symbol)
            self.Liquidate(self.option_symbol)

        # Compute kurtosis from option chain at 9:31 AM
        if not self.computed_kurtosis_today and self.Time.hour == 9 and self.Time.minute == 31:
            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.computed_kurtosis_today = True

        # If kurtosis condition is met, enter straddle at 9:35 AM
        if not self.trade_open and self.kurtosis_condition_met and self.Time.hour == 11 and self.Time.minute == 30:
            self.OpenStraddle(slice)

    def OpenStraddle(self, slice):
        chain = slice.OptionChains.get(self.option_symbol)
        if not chain:
            return

        # 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

        # Select the ATM call and put contracts with the highest Vega
        atm_call = max(atm_call_candidates, key=lambda x: x.Greeks.Vega, default=None)
        atm_put = max(atm_put_candidates, key=lambda x: x.Greeks.Vega, default=None)

        if not atm_call or not atm_put:
            return

        # Calculate credit received from selling the straddle
        credit = atm_call.BidPrice + atm_put.BidPrice

        # Calculate maximum potential loss
        max_loss = abs(atm_call.Strike - self.index.Price) * 100 + credit * 100

        if max_loss <= 0:
            return

        # Position size calculation
        total_portfolio_value = self.Portfolio.TotalPortfolioValue
        max_trade_risk = total_portfolio_value * self.max_portfolio_risk
        contracts = int(max_trade_risk / max_loss)

        if contracts <= 0:
            return

        # Sell straddle
        self.Sell(atm_call.Symbol, contracts)
        self.Sell(atm_put.Symbol, contracts)

        # Calculate net delta
        net_delta = (atm_call.Greeks.Delta + atm_put.Greeks.Delta) * contracts

        # Implement Delta Hedge using SPY
        required_spy_position = -net_delta
        existing_spy_holding = self.Portfolio[self.spy_symbol].Quantity
        spy_order_quantity = math.floor(required_spy_position - existing_spy_holding)

        if spy_order_quantity != 0:
            spy_order = self.MarketOrder(self.spy_symbol, spy_order_quantity)
            self.hedge_order_ticket = spy_order

        # Update trade state
        self.trade_open = True
        self.max_potential_loss = max_loss * contracts

    def CheckPositionManagement(self):
        # Calculate total PnL from options and SPY
        total_pnl = 0.0
        for holding in self.Portfolio.Values:
            if holding.Invested:
                total_pnl += holding.UnrealizedProfit

        if total_pnl >= self.max_potential_loss * self.profit_target:
            self.Liquidate()
            self.trade_open = False
        elif total_pnl <= -self.max_potential_loss * self.stop_loss:
            self.Liquidate()
            self.trade_open = False