from AlgorithmImports import *
from datetime import timedelta
import math
from collections import deque


class ContractData:
    """DTO to hold option contract references."""
    def __init__(self, symbol: Symbol, right: OptionRight, expiry, underlying: Symbol):
        self.symbol = symbol
        self.right = right
        self.expiry = expiry
        self.underlying = underlying


class OptimizedLongGammaStrategy(QCAlgorithm):
    def initialize(self) -> None:
        self.set_start_date(2024, 1, 1)
        self.set_end_date(2026, 5, 16)
        self.set_cash(100000)

        self._contracts: list[ContractData] = []
        self._portfolio_delta = 0.0
        self._hedge_threshold = 15
        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._num_contracts = 1
        self._entry_date = None

        # Risk management
        self._profit_target_multiple = 3.0
        self._stop_loss_multiple = 0.7
        self._risk_per_trade = 0.02
        self._holding_days_max = 18

        # NEW: Trailing stop - once value hits 2x entry, floor at 1.5x
        self._trailing_activation_multiple = 2.0
        self._trailing_floor_multiple = 1.5
        self._trailing_stop_activated = False

        # NEW: Circuit breaker - pause trading after 5% drawdown
        self._max_drawdown_threshold = 0.05  # 5% from peak
        self._circuit_breaker_days = 5
        self._peak_portfolio_value = 100000.0
        self._circuit_breaker_until = None

        # VIX history for percentile filter
        self._vix_history: deque[float] = deque(maxlen=504)

        # Track current underlying for this straddle
        self._current_underlying = None

        # Underlyings to trade for diversification
        self._underlying_symbols: list[Symbol] = []

        self.settings.seed_initial_prices = True

        # Add SPY and QQQ for diversification
        for ticker in ["SPY", "QQQ"]:
            equity = self.add_equity(ticker, data_normalization_mode=DataNormalizationMode.RAW)
            self._underlying_symbols.append(equity.symbol)
            option = self.add_option(ticker, Resolution.MINUTE)
            option.set_filter(-5, 5, 0, 45)

        # VIX for volatility regime filtering
        self._vix = self.add_index("VIX", Resolution.DAILY).symbol

        # Round-robin index for rotating underlyings
        self._underlying_index = 0

        # Weekly entry
        self.schedule.on(
            self.date_rules.every(DayOfWeek.MONDAY),
            self.time_rules.after_market_open(self._underlying_symbols[0], 30),
            self._maybe_enter_straddle
        )

        # Hedge delta mid-morning and before close
        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
        )

        # Check holding period exit daily
        self.schedule.on(
            self.date_rules.every_day(self._underlying_symbols[0]),
            self.time_rules.before_market_close(self._underlying_symbols[0], 30),
            self._check_holding_period
        )

    def on_data(self, data: Slice) -> None:
        if self.is_warming_up:
            return

        # Track VIX history for percentile filter
        if data.contains_key(self._vix):
            vix_price = float(self.securities[self._vix].price)
            if vix_price > 0:
                self._vix_history.append(vix_price)

        # NEW: Update peak portfolio value and check circuit breaker
        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 >= self._max_drawdown_threshold and self._circuit_breaker_until is None:
            self._circuit_breaker_until = self.time + timedelta(days=self._circuit_breaker_days)
            self.log(f"[CIRCUIT BREAKER] Drawdown={drawdown:.2%}, pausing until {self._circuit_breaker_until}")
            self._liquidate_options()

        if not self._contracts:
            return

        # Wait for legs to fill, then record cost basis
        if not self._entry_filled:
            invested_count = sum(
                1 for cd in self._contracts
                if self.portfolio[cd.symbol].invested
            )
            if invested_count == 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

        self._check_risk_limits()

    def _get_realized_vol(self, underlying: Symbol) -> float:
        """Calculate 21-day annualized realized volatility."""
        history = self.history(underlying, 30, Resolution.DAILY)
        bars = list(history)
        if len(bars) < 22:
            return 0.0

        closes = [float(bar.close) for bar in bars[-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)
        var_r = sum((r - mean_r) ** 2 for r in returns) / (len(returns) - 1)
        return math.sqrt(var_r * 252) * 100

    def _get_vix_percentile(self) -> float:
        """Calculate current VIX percentile over the last 2 years."""
        if len(self._vix_history) < 100:
            return 50.0  # Default to median if not enough data

        current_vix = self._vix_history[-1]
        count_below = sum(1 for v in self._vix_history if v < current_vix)
        return (count_below / len(self._vix_history)) * 100

    def _check_risk_limits(self) -> None:
        """Close straddle if profit target, stop loss, or trailing stop is hit."""
        if not self._contracts or self._straddle_entry_cost == 0:
            return

        current_value = abs(sum(
            self.portfolio[cd.symbol].holdings_value
            for cd in self._contracts
            if self.portfolio[cd.symbol].invested
        ))

        # Trailing stop logic
        gain_multiple = current_value / self._straddle_entry_cost
        
        if gain_multiple >= self._trailing_activation_multiple:
            self._trailing_stop_activated = True
        
        if self._trailing_stop_activated:
            # Floor at 1.5x entry to lock in gains
            floor_value = self._straddle_entry_cost * self._trailing_floor_multiple
            if current_value <= floor_value:
                self.log(
                    f"[TRAILING STOP] Value=${current_value:.2f} "
                    f"fell below floor=${floor_value:.2f} (peak was {gain_multiple:.2f}x)"
                )
                self._liquidate_options()
                return

        # Traditional stop loss
        if current_value <= self._straddle_entry_cost * self._stop_loss_multiple:
            self.log(
                f"[STOP LOSS] Value=${current_value:.2f} "
                f"vs Entry=${self._straddle_entry_cost:.2f}"
            )
            self._liquidate_options()
            return

        # Traditional profit target
        if current_value >= self._straddle_entry_cost * self._profit_target_multiple:
            self.log(
                f"[PROFIT TARGET] Value=${current_value:.2f} "
                f"vs Entry=${self._straddle_entry_cost:.2f}"
            )
            self._liquidate_options()

    def _maybe_enter_straddle(self) -> None:
        if self.is_warming_up or self._contracts:
            return

        # NEW: Check circuit breaker
        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
        )

        # VIX regime filter: avoid extremes
        if vix_price < 12 or vix_price > 30:
            self.log(f"[SKIP] VIX={vix_price:.2f} outside 12-30")
            return

        # VIX percentile filter: only enter when IV is relatively cheap
        vix_percentile = self._get_vix_percentile()
        if vix_percentile > 40:
            self.log(f"[SKIP] VIX percentile={vix_percentile:.1f}% > 40%")
            return

        # Rotate underlying for diversification
        self._underlying_index = (self._underlying_index + 1) % len(self._underlying_symbols)
        underlying = self._underlying_symbols[self._underlying_index]

        # Realized vol filter
        realized_vol = self._get_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

        # Clean up any stale positions
        self._liquidate_options()

        # Find options with 14-30 DTE
        chain = self.option_chain(underlying)
        today = self.time.date()

        valid_contracts = [
            x for x in chain
            if 14 <= (x.id.date.date() - today).days <= 30
        ]
        if not valid_contracts:
            self.log("[SKIP] No contracts in 14-30 DTE range")
            return

        # Pick expiry closest to 21 DTE target
        expiries = sorted(
            set(x.id.date for x in valid_contracts),
            key=lambda d: abs((d.date() - today).days - 21)
        )
        target_expiry = expiries[0]
        target_expiry_date = target_expiry.date()
        valid_contracts = [x for x in valid_contracts if x.id.date == target_expiry]

        underlying_price = self.securities[underlying].price
        if underlying_price == 0:
            return

        # Find ATM strike
        strikes = set(x.id.strike_price for x in valid_contracts)
        atm_strike = min(strikes, key=lambda s: abs(s - underlying_price))

        atm_call = next(
            (x for x in valid_contracts
             if x.id.strike_price == atm_strike
             and x.id.option_right == OptionRight.CALL),
            None
        )
        atm_put = next(
            (x for x in valid_contracts
             if x.id.strike_price == atm_strike
             and x.id.option_right == OptionRight.PUT),
            None
        )

        if not atm_call or not atm_put:
            return

        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  # Reset trailing stop for new trade

        call_price = self.securities[call_symbol].price
        put_price = self.securities[put_symbol].price
        straddle_price = call_price + put_price
        if straddle_price == 0:
            return

        # Position sizing
        portfolio_value = self.portfolio.total_portfolio_value
        max_risk = portfolio_value * self._risk_per_trade
        num_contracts = max(1, int(max_risk / (straddle_price * 100)))
        num_contracts = min(num_contracts, 10)
        self._num_contracts = num_contracts

        if self.is_market_open(underlying):
            self.market_order(call_symbol, num_contracts)
            self.market_order(put_symbol, num_contracts)
            self._total_straddles += 1

            estimated_cost = straddle_price * 100 * num_contracts
            dte = (target_expiry_date - today).days
            ticker = underlying.value
            self.log(
                f"[ENTRY #{self._total_straddles}] {ticker} K={atm_strike} "
                f"Price={underlying_price:.2f} VIX={vix_price:.1f} "
                f"VIX%ile={vix_percentile:.0f}% "
                f"RVol={realized_vol:.1f}% DTE={dte} "
                f"Qty={num_contracts} Cost={estimated_cost:.0f}"
            )

    def _check_holding_period(self) -> None:
        """Exit after holding for max days."""
        if self._entry_date and self._contracts and self._entry_filled:
            days_held = (self.time.date() - self._entry_date).days
            if days_held >= self._holding_days_max:
                self.log(f"[HOLD EXPIRY] {days_held} days held, closing")
                self._liquidate_options()

    def _hedge_delta(self) -> None:
        """Calculate portfolio delta from options and hedge with underlying."""
        if not self._contracts or not self._entry_filled:
            return

        underlying = self._current_underlying
        if underlying is None:
            return

        total_delta = 0.0
        for cd in self._contracts:
            holding = self.portfolio[cd.symbol]
            if not holding.invested:
                continue

            chain = self.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

        self._portfolio_delta = total_delta

        if abs(total_delta) < self._hedge_threshold:
            return

        target_shares = -round(total_delta)
        current_shares = int(self.portfolio[underlying].quantity)
        shares_trade = target_shares - current_shares

        if shares_trade != 0 and self.is_market_open(underlying):
            ticket = self.market_order(underlying, shares_trade, tag="Delta hedge")
            self._hedge_order_ids.add(ticket.order_id)
            self._total_hedges += 1
            self.log(
                f"[HEDGE #{self._total_hedges}] Delta={total_delta:.1f} "
                f"Traded {shares_trade:+d} {underlying.value}"
            )

    def on_order_event(self, order_event: OrderEvent) -> None:
        """Handle assignment by liquidating delivered shares."""
        if (order_event.symbol.security_type == SecurityType.EQUITY
                and order_event.status == OrderStatus.FILLED
                and order_event.order_id not in self._hedge_order_ids):
            self.liquidate(order_event.symbol, tag="Assignment liquidation")

    def _liquidate_options(self) -> None:
        """Closes options and zeroes out the underlying hedge."""
        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._portfolio_delta = 0.0
        self._straddle_entry_cost = 0.0
        self._entry_filled = False
        self._entry_date = None
        self._current_underlying = None

    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}"
        )