# region imports
from AlgorithmImports import *
from datetime import timedelta, date
# endregion

class SpyRiskReversal2M(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2012, 1, 1)
        self.SetEndDate(2015, 1, 1)
        self.SetCash(100000)
        self.UniverseSettings.Resolution = Resolution.Minute
        self.SetWarmUp(timedelta(days=3))  # let greeks populate

        # Underlying
        self.spy = self.AddEquity("SPY", Resolution.Minute).Symbol

        # Options
        opt = self.AddOption("SPY", Resolution.Minute)
        opt.SetFilter(self.UniverseFilter)
        self.optionSym = opt.Symbol

        # State
        self.ready = False
        self.currentExpiry = None
        self.callSymbol = None
        self.putSymbol = None

        # Scheduled flags -> handled in OnData (where OptionChains exist)
        self.needs_open = False
        self.needs_roll = False
        self.needs_exit_check = False

        # Schedules
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday),
                         self.TimeRules.At(10, 5),
                         self.FlagWeeklyRoll)
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                         self.TimeRules.BeforeMarketClose("SPY", 10),
                         self.FlagExitCheck)

    def OnWarmupFinished(self):
        self.ready = True
        self.needs_open = True
        self.Debug("Warmup finished — will open first cycle from OnData.")

    def UniverseFilter(self, u: OptionFilterUniverse) -> OptionFilterUniverse:
        # 35–90 DTE so we can choose ~60
        return (u.IncludeWeeklys()
                .Expiration(timedelta(days=35), timedelta(days=90))
                .Strikes(-60, +60))

    # ---------- Scheduler flag setters ----------
    def FlagWeeklyRoll(self):
        if not self.ready:
            return
        self.needs_roll = True

    def FlagExitCheck(self):
        if not self.ready:
            return
        self.needs_exit_check = True

    # ---------- Core selection utilities (run in OnData) ----------
    def select_expiry_closest_to(self, chain, target_days: int):
        expiries = sorted({c.Expiry.date() for c in chain})
        if not expiries:
            return None
        today = self.Time.date()
        return min(expiries, key=lambda d: abs((d - today).days - target_days))

    def safe_delta(self, c):
        try:
            g = c.Greeks
            if g and g.Delta is not None:
                return float(g.Delta)
        except Exception:
            pass
        # quick fallback if greeks are temporarily missing
        spot = float(self.Securities[self.spy].Price)
        if c.Right == OptionRight.Call:
            return 0.25 if c.Strike >= spot else 0.5
        else:
            return -0.25 if c.Strike <= spot else -0.5

    def select_25d_rr(self, chain, expiry_date):
        contracts = [c for c in chain if c.Expiry.date() == expiry_date and c.BidPrice is not None and c.AskPrice is not None]
        if not contracts:
            return None, None
        calls = [c for c in contracts if c.Right == OptionRight.Call]
        puts  = [c for c in contracts if c.Right == OptionRight.Put]
        if not calls or not puts:
            return None, None

        call = min(calls, key=lambda k: abs(self.safe_delta(k) - 0.25))
        put  = min(puts,  key=lambda k: abs(self.safe_delta(k) + 0.25))
        return call, put

    # ---------- Trading helpers ----------
    def set_short_shares(self, target_short_qty=-50):
        cur = self.Portfolio[self.spy].Quantity
        if cur != target_short_qty:
            self.MarketOrder(self.spy, target_short_qty - cur)
            self.Debug(f"Set SPY shares to {target_short_qty}")

    def close_current_legs(self):
        for sym in (self.callSymbol, self.putSymbol):
            if sym and self.Portfolio[sym].Invested:
                self.MarketOrder(sym, -self.Portfolio[sym].Quantity)

    def close_all_spy_options(self):
        for kv in self.Portfolio:
            sym = kv.Key
            if sym.SecurityType == SecurityType.Option and sym.HasUnderlying and sym.Underlying == self.spy:
                h = self.Portfolio[sym]
                if h.Invested and h.Quantity != 0:
                    self.MarketOrder(sym, -h.Quantity)

    # ---------- Main logic in OnData ----------
    def OnData(self, data: Slice):
        if not self.ready:
            return

        chain = data.OptionChains.get(self.optionSym)
        if chain is None or len(chain) == 0:
            return  # wait until we have contracts

        today = self.Time.date()

        # 1) Exit check at 14 DTE
        if self.needs_exit_check and self.currentExpiry is not None:
            dte = (self.currentExpiry - today).days
            if dte <= 14:
                self.Debug(f"Exit (DTE={dte}) — closing current cycle.")
                self.close_current_legs()
                # keep the stock overlay; we'll re-affirm below
                self.currentExpiry = None
                self.callSymbol = None
                self.putSymbol = None
                self.needs_open = True  # immediately start the next cycle
            self.needs_exit_check = False

        # 2) If no active cycle, open a new ~60 DTE RR
        if self.needs_open and self.currentExpiry is None:
            expiry = self.select_expiry_closest_to(chain, 60)
            if expiry:
                c, p = self.select_25d_rr(chain, expiry)
                if c and p:
                    # flatten any leftover options just in case
                    self.close_all_spy_options()
                    # stock overlay
                    self.set_short_shares(-50)
                    # enter RR: +1 call, -1 put
                    self.MarketOrder(c.Symbol, +1)
                    self.MarketOrder(p.Symbol, -1)
                    self.currentExpiry = expiry
                    self.callSymbol = c.Symbol
                    self.putSymbol = p.Symbol
                    self.Debug(f"Opened RR exp {expiry}: +C{c.Strike} (~Δ{self.safe_delta(c):.2f}) / -P{p.Strike} (~Δ{self.safe_delta(p):.2f}); short 50 SPY")
                    self.needs_open = False
            else:
                # no expiry yet — try again next OnData
                pass

        # 3) Weekly roll to keep legs near 25Δ while >14 DTE
        if self.needs_roll and self.currentExpiry is not None:
            dte = (self.currentExpiry - today).days
            if dte > 14:
                newC, newP = self.select_25d_rr(chain, self.currentExpiry)
                if newC and newP and (self.callSymbol != newC.Symbol or self.putSymbol != newP.Symbol):
                    self.Debug(f"Rolling strikes to keep ~25Δ (DTE={dte}).")
                    self.close_current_legs()
                    self.MarketOrder(newC.Symbol, +1)
                    self.MarketOrder(newP.Symbol, -1)
                    self.callSymbol = newC.Symbol
                    self.putSymbol = newP.Symbol
                    self.Debug(f"Rolled to +C{newC.Strike} / -P{newP.Strike} exp {self.currentExpiry}")
            # maintain stock overlay
            self.set_short_shares(-50)
            self.needs_roll = False