from AlgorithmImports import *

class DeltaHedgedStraddleAlgo(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2024, 9, 1)
        self.set_end_date(2024, 12, 31)
        self.set_cash(100000)
        self.settings.seed_initial_prices = True
        self._spy = self.add_equity("SPY", data_normalization_mode=DataNormalizationMode.RAW)
        # Initialize the straddle selector with 7-30 day DTE range.
        self._straddle_selector = StraddleSelector(min_dte=7, max_dte=30)
        # Initialize the delta hedger with 30 share minimum, 20 share rehedge band, and 4 hour delay.
        self._hedger = DeltaHedger(30, 20, timedelta(hours=4))
        self._canonical_option = Symbol.create_canonical_option(self._spy)
        # Set the straddle position size to 50% of portfolio value.
        self._straddle_weight = 0.5
        self._short_straddle = None
        self._exit_day = None
        # Schedule the rebalance method to run 30 minutes after market open.
        self.schedule.on(
            self.date_rules.every_day(self._spy),
            self.time_rules.after_market_open(self._spy, 30),
            self._rebalance
        )
        # Schedule the close expiring method to run 60 minutes before market close.
        self.schedule.on(
            self.date_rules.every_day(self._spy),
            self.time_rules.before_market_close(self._spy, 60),
            self._close_expiring
        )
        self.set_warm_up(timedelta(45))

    def on_data(self, data):
        if self.is_warming_up or not data.bars:
            return
        chain = data.option_chains.get(self._canonical_option)
        if not chain or not self._short_straddle:
            return
        # Accumulate the delta from all option legs in the straddle.
        for leg in self._short_straddle.option_legs:
            contract = chain.contracts.get(leg.symbol)
            if contract:
                self._delta += contract.greeks.delta
        # Calculate the net portfolio delta including underlying and options.
        net_delta = self._hedger.compute_net_delta(self, self._short_straddle.option_legs, self._spy.holdings.quantity, self._delta)
        # Reset the delta accumulator for the next iteration.
        self._delta = 0
        # Exit if it's not time to hedge or net delta is zero.
        if not (self._hedger.should_hedge(self.time) and net_delta):
            return
        # Calculate the hedge quantity needed to neutralize delta.
        hedge_qty = self._hedger.get_quantity(net_delta)
        if hedge_qty:
            self.market_order(self._spy, hedge_qty, tag=f"Daily delta hedge ({net_delta:.2f} share delta)")
            # Record the hedge operation with timestamp and delta.
            self._hedger.record_hedge(self.time, net_delta)
            # Update the net delta after the hedge.
            net_delta += hedge_qty
        self.plot("Portfolio Delta", "Delta", net_delta)

    def _rebalance(self):
        if self.is_warming_up or self.portfolio.invested:
            return
        chain = self.option_chain(self._spy)
        if not chain:
            return
        # Use the straddle selector to find the best ATM straddle.
        result = self._straddle_selector.select(chain, self._spy.price, self.time.date())
        if not result:
            return
        self._delta, expiry, strike = result
        self._short_straddle = OptionStrategies.short_straddle(self._canonical_option, strike, expiry)
        # Calculate the number of contracts based on portfolio value and weight.
        self.buy(self._short_straddle, max(1, int(self.portfolio.total_portfolio_value * self._straddle_weight / (self._spy.price * 100.0))), tag=f"Sell ATM straddle exp {expiry.date()}")
        self._hedger.reset()
        self._exit_day = expiry.date()

    def _close_expiring(self):
        if not self._short_straddle or not self._exit_day or self.time.date() < self._exit_day:
            return
        # Liquidate all positions before expiration.
        self.liquidate(tag="Expiry hedge liquidation")

    def on_order_event(self, order_event):
        # Handle option assignment events.
        if order_event.status == OrderStatus.FILLED and order_event.is_assignment:
            # Liquidate all option legs of the straddle.
            for leg in self._short_straddle.option_legs:
                self.liquidate(leg.symbol, tag="Assignment liquidation")
            # Determine the direction based on whether it's a call or put assignment.
            direction = 1 if self.securities[order_event.symbol].right == OptionRight.CALL else -1
            self.market_order(self._spy, -self._spy.holdings.quantity + order_event.quantity * direction * 100)
            self._short_straddle = None


class StraddleSelector:

    def __init__(self, min_dte, max_dte):
        # Store the minimum days to expiration for contract selection.
        self._min_dte = min_dte
        # Store the maximum days to expiration for contract selection.
        self._max_dte = max_dte

    def select(self, chain, spot_price, current_date):
        # Filter contracts to those within the specified DTE range.
        contracts = [
            contract for contract in chain
            if self._min_dte <= (contract.expiry.date() - current_date).days <= self._max_dte
        ]
        if not contracts:
            return
        # Select the farthest expiration date from the filtered contracts.
        expiry = max(contract.expiry for contract in contracts)
        # Filter to only contracts with the selected expiration.
        contracts = [contract for contract in contracts if contract.expiry == expiry]
        # Find the strike price closest to the current spot price.
        strike = min(contracts, key=lambda c: abs(c.strike - spot_price)).strike
        # Filter to only contracts with the selected strike.
        contracts = [contract for contract in contracts if contract.strike == strike]
        # Ensure we have both call and put for a straddle.
        if len(contracts) < 2:
            return
        # Return the combined delta, expiry, and strike for the straddle.
        return sum(contract.greeks.delta for contract in contracts), expiry, strike


class DeltaHedger:

    def __init__(self, min_shares, rehedge_band, hedge_delay):
        # Store the minimum number of shares required to execute a hedge.
        self._min_shares = min_shares
        # Store the rehedge band threshold to prevent excessive hedging.
        self._rehedge_band = rehedge_band
        # Store the time delay between hedging operations.
        self._hedge_delay = hedge_delay
        self.reset()

    def reset(self):
        self._next_hedge_time = datetime.min
        self._last_hedged_delta = None

    def should_hedge(self, current_time):
        # Check if enough time has passed since the last hedge.
        return current_time >= self._next_hedge_time

    def compute_net_delta(self, algorithm, legs, underlying_quantity, contract_delta):
        # Calculate the net portfolio delta by combining underlying and option positions.
        return underlying_quantity + sum(
            algorithm.securities[leg.symbol].holdings.quantity * contract_delta * 100
            for leg in legs
            if leg.symbol in algorithm.securities
        )

    def get_quantity(self, net_delta):
        if (
            self._last_hedged_delta is not None
            and abs(net_delta - self._last_hedged_delta) < self._rehedge_band
            and abs(net_delta) < self._min_shares + self._rehedge_band
        ):
            return
        # Calculate the required hedge quantity to neutralize the net delta.
        hedge_qty = int(round(-net_delta))
        if abs(hedge_qty) < self._min_shares:
            return
        # Return the calculated hedge quantity.
        return hedge_qty

    def record_hedge(self, current_time, net_delta):
        # Record the delta value after this hedge operation.
        self._last_hedged_delta = net_delta
        # Set the next allowed hedge time based on the configured delay.
        self._next_hedge_time = current_time + self._hedge_delay