from AlgorithmImports import *


class HedgeAction:
    NONE = 0
    HEDGE = 1
    ROLL = 2


class DeltaHedger:

    def __init__(self, min_shares=6, rehedge_band=3, hedge_delay=timedelta(minutes=60), roll_threshold=None):
        self._min_shares = min_shares
        self._rehedge_band = rehedge_band
        self._hedge_delay = hedge_delay
        self._roll_threshold = roll_threshold
        self.reset()

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

    def is_ready(self, current_time):
        return current_time >= self._next_hedge_time

    def compute_net_delta(self, chain, legs, equity_quantity):
        net_delta = float(equity_quantity)
        for leg in legs:
            contract = chain.contracts.get(leg.symbol)
            if contract is None or contract.greeks is None:
                return None
            net_delta += float(contract.greeks.delta) * leg.holdings.quantity * 100.0
        return net_delta

    def evaluate(self, net_delta):
        if self._roll_threshold and abs(net_delta) >= self._roll_threshold:
            return HedgeAction.ROLL
        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 HedgeAction.NONE
        if abs(int(round(-net_delta))) < self._min_shares:
            return HedgeAction.NONE
        return HedgeAction.HEDGE

    def record_hedge(self, current_time, net_delta, cooldown=None):
        self._last_hedged_delta = net_delta
        self._next_hedge_time = current_time + (cooldown or self._hedge_delay)

from AlgorithmImports import *
from straddle_select import StraddleSelector
from delta_hedge import HedgeAction, DeltaHedger


class DeltaHedgedStraddleAlgo(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(1_000_000)
        self.settings.seed_initial_prices = True
        self._spy = self.add_equity("SPY", Resolution.HOUR, data_normalization_mode=DataNormalizationMode.RAW)
        self._straddle = StraddleSelector(min_dte=7, max_dte=30)
        self._hedger = DeltaHedger(min_shares=6, rehedge_band=3, hedge_delay=timedelta(minutes=60), roll_threshold=None)
        self._canonical_option = None
        self._straddle_weight = 1.0
        self._legs = []
        self._exit_day = None
        self._reentry_time = None
        self.schedule.on(
            self.date_rules.every_day(self._spy),
            self.time_rules.after_market_open(self._spy, 30),
            self._rebalance
        )
        self.schedule.on(
            self.date_rules.every_day(self._spy),
            self.time_rules.before_market_close(self._spy, 60),
            self._close_expiring
        )
        self.schedule.on(
            self.date_rules.every_day(self._spy),
            self.time_rules.before_market_close(self._spy, 60),
            self._eod_hedge
        )
        self.set_warm_up(timedelta(45))

    def _rebalance(self):
        chain = self.option_chain(self._spy.symbol, flatten=True)
        result = self._straddle.select(chain, self._spy.price, self.time.date()) if chain else None
        if (self.is_warming_up or self._active_legs() or self.transactions.get_open_orders() or self._spy.price <= 0 or not chain or not result):
            return
        call_symbol, put_symbol, expiry = result
        self._legs = [self.add_option_contract(s) for s in (call_symbol, put_symbol)]
        self._canonical_option = call_symbol.canonical
        self.combo_market_order(
            [Leg.create(leg.symbol, -1) for leg in self._legs],
            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()
        exit_day = expiry.date()
        if not self._spy.exchange.hours.is_open(datetime(exit_day.year, exit_day.month, exit_day.day, 12, 0, 0), False):
            exit_day = self._spy.exchange.hours.get_previous_trading_day(exit_day)
        self._exit_day = exit_day

    def _close_expiring(self):
        legs = self._active_legs()
        if not legs or not self._exit_day or self.time.date() < self._exit_day:
            return
        self._legs, self._exit_day = [], None
        for leg in legs:
            self.liquidate(leg, tag="Expiry liquidation")
        if self._spy.holdings.invested:
            self.liquidate(self._spy, tag="Expiry hedge liquidation")

    def on_data(self, slice_):
        if self._reentry_time is not None:
            if self.transactions.get_open_orders() or self._active_legs():
                return
            self._reentry_time = None
            self._rebalance()
            return
        legs = self._active_legs()
        hedger = self._hedger
        chain = slice_.option_chains.get(self._canonical_option) if self._canonical_option else None
        net_delta = hedger.compute_net_delta(chain, legs, self._spy.holdings.quantity) if chain else None
        if (not legs or not self._spy.exchange.hours.is_open(self.time, False) or not hedger.is_ready(self.time) or self.transactions.get_open_orders(self._spy.symbol) or not chain or net_delta is None):
            return
        action = hedger.evaluate(net_delta)
        if action == HedgeAction.ROLL:
            for leg in legs:
                self.liquidate(leg, tag=f"Delta roll: net_delta={net_delta:.2f}")
            if self._spy.holdings.invested:
                self.liquidate(self._spy, tag=f"Delta roll hedge: net_delta={net_delta:.2f}")
            self._legs, self._exit_day = [], None
            self._reentry_time = self.time
            hedger.reset()
            return
        self.plot("Portfolio Delta", "Delta", net_delta)
        if action == HedgeAction.HEDGE:
            hedge_qty = int(round(-net_delta))
            self.market_order(self._spy, hedge_qty, tag=f"Daily delta hedge ({net_delta:.2f} share delta)")
            hedger.record_hedge(self.time, net_delta)

    def _eod_hedge(self):
        legs = self._active_legs()
        slice_ = self.current_slice
        chain = slice_.option_chains.get(self._canonical_option) if slice_ and self._canonical_option else None
        net_delta = self._hedger.compute_net_delta(chain, legs, self._spy.holdings.quantity) if chain else None
        hedge_qty = int(round(-net_delta)) if net_delta is not None else 0        
        if not legs or self.transactions.get_open_orders(self._spy.symbol) or not slice_ or not chain or net_delta is None or not hedge_qty:
            return
        self.market_order(self._spy, hedge_qty, tag=f"EOD delta hedge ({net_delta:.2f} share delta)")
        self._hedger.record_hedge(self.time, 0.0, cooldown=timedelta(minutes=1))

    def _active_legs(self):
        return [leg for leg in self._legs if leg.holdings.invested]

    def on_order_event(self, order_event):
        if order_event.status != OrderStatus.FILLED:
            return
        if order_event.symbol.security_type == SecurityType.OPTION:
            matching = [l for l in self._legs if l.symbol == order_event.symbol]
            leg = matching[0] if matching else None
            if leg and not leg.holdings.invested:
                for other in self._active_legs():
                    if other.symbol != order_event.symbol:
                        self.liquidate(other, tag="Close option after assignment")
                self._legs = [l for l in self._legs if l.holdings.invested]
        if (order_event.symbol.security_type == SecurityType.EQUITY and self._spy.holdings.invested and not self._active_legs()):
            self.liquidate(self._spy, tag="Assignment liquidation")

from AlgorithmImports import *


class StraddleSelector:

    def __init__(self, min_dte=7, max_dte=30):
        self._min_dte = min_dte
        self._max_dte = max_dte

    def select(self, chain, spot_price, current_date):
        contracts = [
            row for row in chain
            if self._min_dte <= (row.id.date.date() - current_date).days <= self._max_dte
        ]
        if not contracts:
            return None
        expiry = max(row.id.date for row in contracts)
        expiry_rows = [row for row in contracts if row.id.date == expiry]
        strike = min(expiry_rows, key=lambda r: abs(r.id.strike_price - spot_price)).id.strike_price
        symbol_by_right = {
            row.id.option_right: row.symbol
            for row in expiry_rows
            if row.id.strike_price == strike
            and row.id.option_right in (OptionRight.CALL, OptionRight.PUT)
        }
        if len(symbol_by_right) < 2:
            return None
        return symbol_by_right[OptionRight.CALL], symbol_by_right[OptionRight.PUT], expiry