from AlgorithmImports import *
from datetime import timedelta

class OptimizedLongGammaStrategy(QCAlgorithm):
    def initialize(self) -> None:
        self.set_start_date(2022, 1, 1)
        self.set_end_date(2024, 9, 30)
        self.set_cash(100000)

        self.contracts = []
        self.portfolio_delta = 0
        self.hedge_threshold = 25  # Shares of delta required to trigger a hedge
        self.straddle_entry_value = 0
        self.total_straddles = 0
        self.total_hedges = 0
        self.next_hedge_time = self.time

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

        # 1. 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
        )
        
        # 2. Close straddle on Friday afternoon to avoid weekend theta & expiration
        self.schedule.on(
            self.date_rules.every(DayOfWeek.FRIDAY),
            self.time_rules.before_market_close(self.spy, 15),
            self.liquidate_options
        )

    # ──────────────────────────────────────────────────────────────────────────
    def on_data(self, data: Slice) -> None:
        if self.is_warming_up or not self.contracts:
            return
            
        # 3. Intraday Gamma Scalping: Check delta every 30 minutes
        if self.time >= self.next_hedge_time:
            self.hedge_delta()
            self.next_hedge_time = self.time + timedelta(minutes=30)

    # ──────────────────────────────────────────────────────────────────────────
    def enter_straddle(self) -> None:
        if self.is_warming_up:
            return

        self.liquidate_options()

        # Get this week's option chain
        chain = self.option_chain(self.spy)
        week_end = Expiry.end_of_week(self.time)
        
        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
        atm_strike = min(
            set(x.id.strike_price for x in valid_contracts),
            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

        # 4. Add contracts and explicitly set the pricing model for accurate Greeks
        for contract in [atm_call, atm_put]:
            symbol = self.add_option_contract(contract).symbol
            self.securities[symbol].set_price_model(OptionPriceModels.BlackScholes())
            self.contracts.append(symbol)

        self.market_order(self.contracts[0], 1)
        self.market_order(self.contracts[1], 1)
        self.total_straddles += 1

        cost = (self.securities[self.contracts[0]].price + self.securities[self.contracts[1]].price) * 100
        self.straddle_entry_value = cost
        self.log(f"[ENTRY #{self.total_straddles}] Strike={atm_strike} | SPY=${spy_price:.2f} | Cost=${cost:.2f}")

    # ──────────────────────────────────────────────────────────────────────────
    def hedge_delta(self) -> None:
        total_delta = 0
        
        for contract in self.contracts:
            if not self.portfolio.contains_key(contract) or not self.portfolio[contract].invested:
                continue
                
            sec = self.securities[contract]



            if sec.price > 0:
                # Try to access greeks through different attributes
                greeks = None
                if hasattr(sec, 'option_greeks') and sec.option_greeks:
                    greeks = sec.option_greeks
                elif hasattr(sec, 'greeks') and sec.greeks:
                    greeks = sec.greeks

                if greeks and hasattr(greeks, 'delta'):
                    total_delta += greeks.delta * self.portfolio[contract].quantity * 100

            if sec.price > 0:
                # Try to access greeks through different attributes
                greeks = None
                if hasattr(sec, 'option_greeks') and sec.option_greeks:
                    greeks = sec.option_greeks
                elif hasattr(sec, 'greeks') and sec.greeks:
                    greeks = sec.greeks

                if greeks and hasattr(greeks, 'delta'):
                    total_delta += greeks.delta * self.portfolio[contract].quantity * 100

        self.portfolio_delta = total_delta

        # Check if the delta breach exceeds our threshold
        if abs(total_delta) < self.hedge_threshold:
            return

        # Calculate how many underlying shares we need to buy/sell to neutralize
        target_spy   = -round(total_delta)
        current_spy  = self.portfolio[self.spy].quantity if self.portfolio.contains_key(self.spy) else 0
        shares_trade = target_spy - current_spy

        if shares_trade != 0:
            self.market_order(self.spy, shares_trade)
            self.total_hedges += 1
            self.log(f"[HEDGE #{self.total_hedges}] Delta={total_delta:.1f} | Traded {shares_trade:+d} SPY")

    # ──────────────────────────────────────────────────────────────────────────
    def liquidate_options(self) -> None:
        """Closes options and zeroes out the underlying hedge."""
        for contract in self.contracts:
            if self.portfolio.contains_key(contract) and self.portfolio[contract].invested:
                self.liquidate(contract)
                
        if self.portfolio.contains_key(self.spy) and self.portfolio[self.spy].invested:
            self.liquidate(self.spy)
            
        self.contracts = []
        self.portfolio_delta = 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}"
        )