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

class IVSkewStraddleStrategy(QCAlgorithm):

    def Initialize(self):
        # Limit backtest to one month for memory optimization
        self.SetStartDate(2023, 5, 1)   # Start date
        self.SetEndDate(2023, 8, 1)    # End date
        self.SetCash(500000)            # Starting cash

        # Add SPX index
        index = self.AddIndex("SPX")
        # Add SPXW options (weekly options for SPX)
        option = self.AddIndexOption(index.Symbol, "SPXW")
        option.SetFilter(lambda u: u.IncludeWeeklys().Expiration(0, 0).Strikes(-30, 30))
        option.PriceModel = OptionPriceModels.BinomialCoxRossRubinstein()
        self.option_symbol = option.Symbol

        # Threshold for IV skew steepness
        self.iv_slope_threshold = 0.005

        # Risk management parameters
        self.max_portfolio_risk = 0.2  # 10% max risk per trade
        self.profit_target = 0.5        # 50% profit target
        self.stop_loss = 0.5            # 50% stop loss
        self.trade_open = False
        self.initial_debit = 0

        # Schedule the processing at market open
        self.Schedule.On(
            self.DateRules.EveryDay(),
            self.TimeRules.AfterMarketOpen("SPX", 1),
            self.ProcessMarketOpen
        )

    def ProcessMarketOpen(self):
        if self.trade_open:
            return  # Only one trade per day

        # Get the option chain
        chain = self.CurrentSlice.OptionChains.get(self.option_symbol, None)
        if not chain:
            self.Log("Option chain not available at market open.")
            return

        # Calculate IV skew data
        iv_skew_data = self.CalculateSteepestIVSkew(chain)
        if iv_skew_data is None:
            self.Log("IV skew data is None at market open.")
            return

        steepest_slope = iv_skew_data["steepest_slope"]
        avg_iv = iv_skew_data["average_iv"]
        self.Log(f"Date: {self.Time.date()} | Steepest Local Slope of Put IV Skew at Market Open: {steepest_slope}")
        self.Log(f"Date: {self.Time.date()} | Average IV of Puts at Market Open: {avg_iv}")

        if steepest_slope < self.iv_slope_threshold:
            self.Log("Steepest slope below threshold at market open. Skipping trade.")
            return

        # Execute straddle trade with hedge
        expiry = iv_skew_data["expiry"]
        atm_strike = iv_skew_data["atm_strike"]

        # Build the straddle strategy (buying ATM call and put)
        straddle = OptionStrategies.Straddle(self.option_symbol, atm_strike, expiry)
        self.Buy(straddle, 1)

        # Hedge with 15% OTM put (selling the put)
        otm_put = iv_skew_data["otm_put"]
        self.Sell(otm_put.Symbol, 1)

        # Calculate initial debit (total cost)
        initial_debit = self.CalculateInitialDebit(atm_strike, expiry, otm_put, chain)
        if initial_debit is None:
            self.Log("Unable to calculate initial debit. Skipping trade.")
            return

        self.initial_debit = initial_debit
        self.Log(f"Opened Straddle at strike {atm_strike} with hedge put at market open. Initial debit: ${self.initial_debit:.2f}")

        self.trade_open = False

    def OnData(self, data):
        if self.trade_open:
            self.CheckPositionManagement()

    def CheckPositionManagement(self):
        # Calculate total PnL of open positions
        total_pnl = sum([holding.UnrealizedProfit for holding in self.Portfolio.Values if holding.Invested])

        # Profit target reached
        if total_pnl >= self.initial_debit * self.profit_target:
            self.Liquidate()
            self.Log(f"Closed position at profit target on {self.Time}. Profit: ${total_pnl:.2f}")
            self.trade_open = False
        # Stop loss reached
        elif total_pnl <= -self.initial_debit * self.stop_loss:
            self.Liquidate()
            self.Log(f"Closed position at stop loss on {self.Time}. Loss: ${total_pnl:.2f}")
            self.trade_open = False

    def CalculateSteepestIVSkew(self, chain):
        """
        Calculate the steepest local slope of the put IV skew (IV vs Strike Price)
        and the average IV of the puts.
        """
        puts = [
            c for c in chain
            if c.Right == OptionRight.Put and c.ImpliedVolatility > 0 and c.ImpliedVolatility <= 5
        ]

        if len(puts) < 2:
            self.Log("Not enough puts to calculate steepest slope.")
            return None  # Not enough data to calculate slope

        # Calculate average IV
        avg_iv = sum(c.ImpliedVolatility for c in puts) / len(puts)

        # Sort puts by strike price
        puts = sorted(puts, key=lambda c: c.Strike)

        # Calculate local slopes
        steepest_slope = 0
        for i in range(len(puts) - 1):
            strike_diff = puts[i + 1].Strike - puts[i].Strike
            iv_diff = puts[i + 1].ImpliedVolatility - puts[i].ImpliedVolatility

            if strike_diff != 0:  # Avoid division by zero
                local_slope = abs(iv_diff / strike_diff)
                if local_slope > steepest_slope:
                    steepest_slope = local_slope

        expiry = puts[0].Expiry
        underlying_price = chain.Underlying.Price

        # Find ATM strike
        atm_strike = min([c.Strike for c in puts], key=lambda x: abs(x - underlying_price))

        # Find 15% OTM put
        otm_put = min(
            puts,
            key=lambda c: abs(c.Strike - 0.85 * underlying_price)
        )

        return {
            "steepest_slope": steepest_slope,
            "average_iv": avg_iv,
            "expiry": expiry,
            "atm_strike": atm_strike,
            "otm_put": otm_put
        }

    def CalculateInitialDebit(self, atm_strike, expiry, otm_put, chain):
        """
        Calculate the initial debit (total cost) of entering the position.
        """
        # Get ATM call and put contracts
        atm_call = [
            c for c in chain
            if c.Right == OptionRight.Call and c.Strike == atm_strike and c.Expiry == expiry
        ]
        atm_put = [
            c for c in chain
            if c.Right == OptionRight.Put and c.Strike == atm_strike and c.Expiry == expiry
        ]

        if not atm_call or not atm_put:
            self.Log("ATM contracts not found.")
            return None

        # Get option prices
        atm_call_price = atm_call[0].AskPrice
        atm_put_price = atm_put[0].AskPrice
        otm_put_price = otm_put.BidPrice

        # Calculate initial debit (cost to buy straddle minus premium received from selling OTM put)
        initial_debit = ((atm_call_price + atm_put_price) - otm_put_price) * 100  # Multiply by 100 for option multiplier
        return initial_debit