from AlgorithmImports import *
from datetime import timedelta


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


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 = 30  # Higher threshold = fewer hedges
        self._straddle_entry_cost = 0.0
        self._total_straddles = 0
        self._total_hedges = 0
        self._hedge_order_ids: set[int] = set()

        # Risk management
        self._profit_target_multiple = 1.4   # Close at 40% profit
        self._stop_loss_multiple = 0.5       # Close at 50% loss
        self._vix_threshold = 30             # VIX filter (raised from 20)

        self.settings.seed_initial_prices = True
        self._spy = self.add_equity(
            "SPY", data_normalization_mode=DataNormalizationMode.RAW
        ).symbol

        # Add VIX for volatility filtering
        self._vix = self.add_equity(
            "VIX", data_normalization_mode=DataNormalizationMode.RAW
        ).symbol

        # Enter new straddle every Monday morning
        self.schedule.on(
            self.date_rules.every(DayOfWeek.MONDAY),
            self.time_rules.after_market_open(self._spy, 30),
            self._enter_straddle
        )

        # Hedge delta once per day at midday to minimize transaction costs
        self.schedule.on(
            self.date_rules.every_day(self._spy),
            self.time_rules.after_market_open(self._spy, 120),  # ~11:30 AM
            self._hedge_delta
        )

        # Close straddle on Friday afternoon
        self.schedule.on(
            self.date_rules.every(DayOfWeek.FRIDAY),
            self.time_rules.before_market_close(self._spy, 30),
            self._liquidate_options
        )

    def on_data(self, data: Slice) -> None:
        """Check profit/loss targets on each data update."""
        if self.is_warming_up or not self._contracts:
            return
        self._check_risk_limits()

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

        current_value = 0.0
        for cd in self._contracts:
            holding = self.portfolio[cd.symbol]
            if holding.invested:
                current_value += holding.holdings_value

        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()
        elif 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()

    def _enter_straddle(self) -> None:
        if self.is_warming_up:
            return

        # VIX filter: only enter when implied vol is reasonable
        vix_price = self.securities[self._vix].price
        if vix_price > self._vix_threshold:
            self.log(f"[SKIP ENTRY] VIX={vix_price:.2f} > {self._vix_threshold}")
            return

        self._liquidate_options()

        # Get this week's option chain
        chain = self.option_chain(self._spy)
        days_until_friday = (4 - self.time.weekday()) % 7
        if days_until_friday == 0:
            days_until_friday = 7
        week_end = self.time + timedelta(days=days_until_friday)

        valid_contracts = [x for x in chain if x.id.date < week_end]
        if not valid_contracts:
            return

        spy_price = self.securities[self._spy].price
        if spy_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 - spy_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),
            ContractData(put_symbol, OptionRight.PUT),
        ]

        call_price = self.securities[call_symbol].price
        put_price = self.securities[put_symbol].price

        self.limit_order(call_symbol, 1, call_price * 1.01)
        self.limit_order(put_symbol, 1, put_price * 1.01)
        self._total_straddles += 1

        cost = (call_price + put_price) * 100
        self._straddle_entry_cost = cost
        self.log(
            f"[ENTRY #{self._total_straddles}] Strike={atm_strike} "
            f"| SPY=${spy_price:.2f} | VIX={vix_price:.2f} | Cost=${cost:.2f}"
        )

    def _hedge_delta(self) -> None:
        """Calculate portfolio delta from options and hedge with underlying."""
        if not self._contracts:
            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_spy = -round(total_delta)
        current_spy = int(self.portfolio[self._spy].quantity)
        shares_trade = target_spy - current_spy

        if shares_trade != 0 and self.is_market_open(self._spy):
            ticket = self.market_order(self._spy, 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} SPY"
            )

    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(self._spy, 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)

        if self.portfolio[self._spy].invested:
            self.liquidate(self._spy)

        self._contracts = []
        self._portfolio_delta = 0.0
        self._straddle_entry_cost = 0.0

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