# region imports
from AlgorithmImports import *
# endregion


class ContractData:

    def __init__(self, symbol: Symbol, right: OptionRight, expiry, underlying: Symbol):
        self.symbol = symbol
        self.right = right
        self.expiry = expiry
        self.underlying = underlying


class StraddleSelector:

    def __init__(self, min_dte: int, max_dte: int, target_dte: int):
        self._min_dte = min_dte
        self._max_dte = max_dte
        self._target_dte = target_dte

    def select(self, chain, underlying_price, today):
        valid_contracts = [x for x in chain if self._min_dte <= (x.id.date.date() - today).days <= self._max_dte]
        if not valid_contracts:
            return None
        target_expiry = sorted(set(x.id.date for x in valid_contracts), key=lambda d: abs((d.date() - today).days - self._target_dte))[0]
        valid_contracts = [x for x in valid_contracts if x.id.date == target_expiry]
        if underlying_price == 0:
            return None
        atm_strike = min(set(x.id.strike_price for x in valid_contracts), key=lambda s: abs(s - underlying_price))
        calls_list = [x for x in valid_contracts if x.id.strike_price == atm_strike and x.id.option_right == OptionRight.CALL]
        atm_call = calls_list[0] if calls_list else None
        puts_list = [x for x in valid_contracts if x.id.strike_price == atm_strike and x.id.option_right == OptionRight.PUT]
        atm_put = puts_list[0] if puts_list else None
        if not atm_call or not atm_put:
            return None
        return atm_call, atm_put, target_expiry, atm_strike


class DeltaHedger:

    def __init__(self, threshold: float):
        self._threshold = threshold

    def compute_shares(self, algorithm, contracts, current_slice, underlying):
        total_delta = 0.0
        for cd in contracts:
            holding = algorithm.portfolio[cd.symbol]
            if not holding.invested:
                continue
            chain = current_slice.option_chains.get(cd.symbol.canonical)
            if chain is None:
                continue
            contract = chain.contracts.get(cd.symbol)
            if contract is None:
                continue
            total_delta += contract.greeks.delta * holding.quantity * 100
        if abs(total_delta) < self._threshold:
            return 0, None
        shares_trade = -round(total_delta) - int(algorithm.portfolio[underlying].quantity)
        return shares_trade, total_delta

# region imports
from AlgorithmImports import *
from helpers import ContractData, StraddleSelector, DeltaHedger
# endregion


class OptimizedLongGammaStrategy(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2024, 1, 1)
        self.set_end_date(2026, 5, 1)
        self.set_cash(100000)
        self.settings.seed_initial_prices = True
        self._contracts: list[ContractData] = []
        self._straddle_entry_cost = 0.0
        self._total_straddles = 0
        self._total_hedges = 0
        self._hedge_order_ids: set[int] = set()
        self._entry_filled = False
        self._entry_date = None
        self._trailing_stop_activated = False
        self._peak_portfolio_value = 100000.0
        self._circuit_breaker_until = None
        self._straddle_selector = StraddleSelector(14, 30, 21)
        self._delta_hedger = DeltaHedger(15)
        self._current_underlying = None
        self._underlying_symbols: list[Symbol] = []
        # Daily close windows for realized vol — consolidate to daily
        # regardless of subscription resolution to guarantee correct annualization
        self._rv_windows: dict = {}
        for ticker in ["SPY", "QQQ"]:
            equity = self.add_equity(ticker, Resolution.MINUTE, data_normalization_mode=DataNormalizationMode.RAW)
            self._rv_windows[equity.symbol] = []
            self.consolidate(equity.symbol, Resolution.DAILY, self._on_equity_daily_bar)
            self._underlying_symbols.append(equity.symbol)
        self._vix = self.add_index("VIX", Resolution.MINUTE)
        # Daily consolidator keeps a proper 510-day lookback at MINUTE subscription
        # — mirrors the _rv_windows pattern for SPY/QQQ
        self._vix_closes: list = []
        self.consolidate(self._vix, Resolution.DAILY, self._on_vix_daily_bar)
        self._underlying_index = 0
        self.schedule.on(
            self.date_rules.every(DayOfWeek.MONDAY),
            self.time_rules.after_market_open(self._underlying_symbols[0], 30),
            self._maybe_enter_straddle
        )
        self.schedule.on(
            self.date_rules.every_day(self._underlying_symbols[0]),
            self.time_rules.after_market_open(self._underlying_symbols[0], 120),
            self._hedge_delta
        )
        self.schedule.on(
            self.date_rules.every_day(self._underlying_symbols[0]),
            self.time_rules.before_market_close(self._underlying_symbols[0], 60),
            self._hedge_delta
        )
        self.set_warm_up(510, Resolution.DAILY)

    def _on_equity_daily_bar(self, bar) -> None:
        window = self._rv_windows.get(bar.symbol)
        if window is None:
            return
        window.append(bar.close)
        if len(window) > 25:
            window.pop(0)

    def _on_vix_daily_bar(self, bar) -> None:
        self._vix_closes.append(bar.close)
        if len(self._vix_closes) > 510:
            self._vix_closes.pop(0)

    def on_data(self, data: Slice) -> None:
        if self.is_warming_up:
            return
        current_value = self.portfolio.total_portfolio_value
        if current_value > self._peak_portfolio_value:
            self._peak_portfolio_value = current_value
        drawdown = (self._peak_portfolio_value - current_value) / self._peak_portfolio_value
        if drawdown >= 0.05 and self._circuit_breaker_until is None:
            self._circuit_breaker_until = self.time + timedelta(days=5)
            self.log(f"[CIRCUIT BREAKER] Drawdown={drawdown:.2%}, pausing until {self._circuit_breaker_until}")
            self._liquidate_options()
        if not self._contracts:
            return
        if not self._entry_filled:
            if sum(1 for cd in self._contracts if self.portfolio[cd.symbol].invested) == len(self._contracts):
                self._entry_filled = True
                self._straddle_entry_cost = abs(sum(
                    self.portfolio[cd.symbol].holdings_value
                    for cd in self._contracts
                ))
                self.log(f"[FILLED] Entry cost=${self._straddle_entry_cost:.2f}")
            return
        if self._straddle_entry_cost != 0:
            # Exit at 7 DTE to avoid theta cliff; prevents expiry-to-zero losses
            dte = (self._contracts[0].expiry.date() - self.time.date()).days
            if dte <= 7:
                self.log(f"[DTE EXIT] {dte} DTE remaining, closing to avoid theta cliff")
                self._liquidate_options()
                return
            limits_current_value = abs(sum(
                self.portfolio[cd.symbol].holdings_value
                for cd in self._contracts
                if self.portfolio[cd.symbol].invested
            ))
            gain_multiple = limits_current_value / self._straddle_entry_cost
            if gain_multiple >= 2.0:
                self._trailing_stop_activated = True
            liquidated = False
            if self._trailing_stop_activated:
                floor_value = self._straddle_entry_cost * 1.5
                if limits_current_value < floor_value:
                    self.log(f"[TRAILING STOP] Value=${limits_current_value:.2f} fell below floor=${floor_value:.2f} (peak was {gain_multiple:.2f}x)")
                    self._liquidate_options()
                    liquidated = True
            # Tightened stop: 75% (was 70%) — cuts losses 5% sooner
            if not liquidated and limits_current_value < self._straddle_entry_cost * 0.75:
                self.log(f"[STOP LOSS] Value=${limits_current_value:.2f} vs Entry=${self._straddle_entry_cost:.2f}")
                self._liquidate_options()

    def _vix_percentile(self) -> float:
        if len(self._vix_closes) < 100:
            return 50.0
        current_vix = self._vix_closes[-1]
        return (sum(1 for v in self._vix_closes if v < current_vix) / len(self._vix_closes)) * 100

    def _realized_vol(self, underlying: Symbol) -> float:
        window = self._rv_windows.get(underlying)
        if window is None or len(window) < 22:
            return 0.0
        closes = window[-22:]
        returns = [math.log(closes[i] / closes[i - 1]) for i in range(1, len(closes))]
        if len(returns) < 20:
            return 0.0
        mean_r = sum(returns) / len(returns)
        variance = sum((r - mean_r) ** 2 for r in returns) / (len(returns) - 1)
        return math.sqrt(variance * 252) * 100

    def _maybe_enter_straddle(self) -> None:
        if self.is_warming_up or self._contracts:
            return
        if self._circuit_breaker_until and self.time < self._circuit_breaker_until:
            return
        vix_price = float(self.securities[self._vix].price) if self.securities[self._vix].has_data else 20.0
        if vix_price < 12 or vix_price > 35:
            self.log(f"[SKIP] VIX={vix_price:.2f} outside 12-35")
            return
        self._underlying_index = (self._underlying_index + 1) % len(self._underlying_symbols)
        underlying = self._underlying_symbols[self._underlying_index]
        realized_vol = self._realized_vol(underlying)
        if realized_vol > 0 and realized_vol < vix_price * 0.7:
            self.log(f"[SKIP] RVol={realized_vol:.1f}% < 70% VIX={vix_price:.1f}%")
            return
        self._liquidate_options()
        today = self.time.date()
        underlying_price = self.securities[underlying].price
        if underlying_price == 0:
            return
        chain = self.option_chain(underlying)
        res = self._straddle_selector.select(chain, underlying_price, today)
        if not res:
            self.log("[SKIP] No contracts in 14-30 DTE range")
            return
        atm_call, atm_put, target_expiry, atm_strike = res
        call_symbol = self.add_option_contract(atm_call).symbol
        put_symbol = self.add_option_contract(atm_put).symbol
        self._contracts = [
            ContractData(call_symbol, OptionRight.CALL, target_expiry, underlying),
            ContractData(put_symbol, OptionRight.PUT, target_expiry, underlying),
        ]
        self._current_underlying = underlying
        self._entry_filled = False
        self._straddle_entry_cost = 0.0
        self._entry_date = self.time.date()
        self._trailing_stop_activated = False
        straddle_price = self.securities[call_symbol].price + self.securities[put_symbol].price
        if straddle_price == 0:
            return
        num_contracts = min(max(1, int((self.portfolio.total_portfolio_value * 0.02) / (straddle_price * 100))), 10)
        if self.is_market_open(underlying):
            self.market_order(call_symbol, num_contracts)
            self.market_order(put_symbol, num_contracts)
            self._total_straddles += 1
            self.log(
                f"[ENTRY #{self._total_straddles}] {underlying} K={atm_strike} "
                f"Price={underlying_price:.2f} VIX={vix_price:.1f} "
                f"RVol={realized_vol:.1f}% DTE={(target_expiry.date() - today).days} "
                f"Qty={num_contracts} Cost={straddle_price * 100 * num_contracts:.0f}"
            )

    def _hedge_delta(self) -> None:
        if not self._contracts or not self._entry_filled:
            return
        underlying = self._current_underlying
        if underlying is None:
            return
        shares_trade, total_delta = self._delta_hedger.compute_shares(self, self._contracts, self.current_slice, underlying)
        if total_delta is None:
            return
        if shares_trade != 0 and self.is_market_open(underlying):
            self._hedge_order_ids.add(self.market_order(underlying, shares_trade, tag="Delta hedge").order_id)
            self._total_hedges += 1
            self.log(
                f"[HEDGE #{self._total_hedges}] Delta={total_delta:.1f} "
                f"Traded {shares_trade:+d} {underlying}"
            )

    def on_order_event(self, order_event: OrderEvent) -> None:
        if (order_event.symbol.security_type != SecurityType.EQUITY
                or order_event.status != OrderStatus.FILLED):
            return
        if order_event.order_id in self._hedge_order_ids:
            return
        # In backtesting, market orders fill synchronously inside market_order(), so
        # on_order_event fires before the returned order_id is added to _hedge_order_ids.
        # Fall back to tag check to handle this race condition.
        order = self.transactions.get_order_by_id(order_event.order_id)
        if order is not None and order.tag == "Delta hedge":
            self._hedge_order_ids.add(order_event.order_id)
            return
        self.liquidate(order_event.symbol, tag="Assignment liquidation")

    def _liquidate_options(self) -> None:
        for cd in self._contracts:
            if self.portfolio[cd.symbol].invested:
                self.liquidate(cd.symbol)
        for underlying in self._underlying_symbols:
            if self.portfolio[underlying].invested:
                self.liquidate(underlying)
        self._contracts = []
        self._straddle_entry_cost = 0.0
        self._entry_filled = False
        self._entry_date = None
        self._current_underlying = None
        self._trailing_stop_activated = False

    def on_end_of_algorithm(self) -> None:
        final = self.portfolio.total_portfolio_value
        self.log(
            f"\n{'='*50}\n"
            f"  Final Value  : ${final:,.2f}\n"
            f"  Total Return : {(final - 100000) / 100000 * 100:+.2f}%\n"
            f"  Straddles    : {self._total_straddles}\n"
            f"  Hedges       : {self._total_hedges}\n"
            f"{'='*50}"
        )