from AlgorithmImports import *

class LongGammaStrategy(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  # hedge if delta exceeds 25 shares
        self.straddle_entry_value = 0
        self.total_straddles = 0
        self.total_hedges = 0

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

        # Enter new straddle every Monday
        self.schedule.on(
            self.date_rules.every(DayOfWeek.MONDAY),
            self.time_rules.after_market_open(self.spy, 30),
            self.enter_straddle
        )
        # Delta hedge once per day near close
        self.schedule.on(
            self.date_rules.every_day(self.spy),
            self.time_rules.before_market_close(self.spy, 15),
            self.hedge_delta
        )

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

        # Close existing positions
        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 = []

        # Get this week's option chain
        chain = self.option_chain(self.spy)
        week_end = Expiry.end_of_week(self.time)
        self.contracts = [
            self.add_option_contract(x).symbol for x in chain
            if x.id.date < week_end
        ]

        spy_price = self.securities[self.spy].price
        if spy_price == 0 or not self.contracts:
            return

        # Find ATM strike
        atm_strike = min(
            set(self.securities[c].symbol.id.strike_price for c in self.contracts),
            key=lambda s: abs(s - spy_price)
        )

        atm_call = next((c for c in self.contracts
            if self.securities[c].symbol.id.strike_price == atm_strike
            and self.securities[c].symbol.id.option_right == OptionRight.CALL), None)

        atm_put = next((c for c in self.contracts
            if self.securities[c].symbol.id.strike_price == atm_strike
            and self.securities[c].symbol.id.option_right == OptionRight.PUT), None)

        if not atm_call or not atm_put:
            return

        self.market_order(atm_call, 1)
        self.market_order(atm_put, 1)
        self.total_straddles += 1

        cost = (self.securities[atm_call].price + self.securities[atm_put].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:
        if self.is_warming_up:
            return

        total_delta = 0
        for contract in self.contracts:
            if not self.portfolio.contains_key(contract):
                continue
            if not self.portfolio[contract].invested:
                continue
            sec = self.securities[contract]
            if sec.price > 0 and hasattr(sec, 'option_greeks') and sec.option_greeks:
                total_delta += sec.option_greeks.delta * self.portfolio[contract].quantity * 100

        self.portfolio_delta = total_delta

        if abs(total_delta) < self.hedge_threshold:
            return

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