from AlgorithmImports import *

class DynamicProtectivePutStrategy(QCAlgorithm):

    def Initialize(self):
        # === STRATEGY SETUP ===
        self.SetStartDate(2021, 1, 1)
        self.SetEndDate(2024, 1, 1)
        self.cash = 1000000
        self.SetCash(self.cash)

        # === STRATEGY CONFIG ===
        self.top_n_marketcap = 35 # top market caps worldwide
        self.top_n_momentum = 5   # number of stocks in the portfolio
        self.roc_period = 252  # 12-month ROC
        self.delta_put = 40    # Target delta to be mu

        #evenly distribute cash across the shares while keeping margin for the options
        self.stock_investment = self.cash / self.top_n_momentum * 0.7  

        self.active_symbols = []        # currently invested stocks
        self.stock_data = {}            # stock symbol -> {"put": option symbol}
        self.pending_hedges = {}        # stock symbol -> shares waiting for put hedge

        # === UNIVERSE SETTINGS ===
        self.UniverseSettings.Resolution = Resolution.Daily
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.Raw

        self.AddUniverse(self.CoarseSelection, self.FineSelection)

        # === MONTHLY REBALANCING SCHEDULE ===
        self.Schedule.On(
            self.DateRules.MonthStart("SPY"),
            self.TimeRules.At(9, 30),
            self.Rebalance
        )

    def CoarseSelection(self, coarse):
        # Filter for liquid stocks with fundamental data
        return [x.Symbol for x in coarse if x.HasFundamentalData and x.Price > 10 and x.Volume > 100000]

    def FineSelection(self, fine):
        # Select top 35 by market cap
        sorted_by_cap = sorted(fine, key=lambda x: x.MarketCap, reverse=True)
        self.market_cap_universe = [x.Symbol for x in sorted_by_cap[:self.top_n_marketcap]]
        return self.market_cap_universe

    def Momentum(self, symbol):
        # Compute 12-month rate of change
        history = self.History(symbol, self.roc_period, Resolution.Daily)
        if history.empty or len(history["close"]) < self.roc_period:
            return -1
        prices = history["close"]
        return (prices.iloc[-1] - prices.iloc[0]) / prices.iloc[0]

    def Rebalance(self):
        # Wait for universe to be ready
        if not hasattr(self, 'market_cap_universe') or not self.market_cap_universe:
            return

        # === SELECT TOP 5 MOMENTUM STOCKS ===
        ranked = sorted(
            self.market_cap_universe,
            key=lambda sym: self.Momentum(sym),
            reverse=True
        )
        new_top = ranked[:self.top_n_momentum]

        # === REMOVE OLD POSITIONS ===
        removed = [s for s in self.active_symbols if s not in new_top]
        for symbol in removed:
            self.Liquidate(symbol)
            if self.stock_data.get(symbol, {}).get("put"):
                self.Liquidate(self.stock_data[symbol]["put"])
                self.stock_data[symbol]["put"] = None

        self.active_symbols = new_top

        # === ENTER POSITIONS & SCHEDULE HEDGES ===
        for symbol in self.active_symbols:
            if symbol not in self.stock_data:
                self.stock_data[symbol] = {"put": None}

            price = self.Securities[symbol].Price
            if price == 0:
                continue

            # Buy ~$50K of stock in 100-share lots
            raw_qty = int(self.stock_investment / price)
            shares = int(round(raw_qty / 100.0)) * 100
            if shares < 100:
                self.Debug(f"Skipping {symbol.Value}: too expensive for 100 shares.")
                continue

            self.MarketOrder(symbol, shares)
            self.pending_hedges[symbol] = shares  # flag for delta-based put selection in OnData

    def OnData(self, data):
        if not self.pending_hedges:
            return

        for symbol, shares in list(self.pending_hedges.items()):
            if symbol not in data.OptionChains:
                continue

            chain = data.OptionChains[symbol]

            # Filter puts expiring in ~30-36 days
            puts = [x for x in chain if x.Right == OptionRight.Put and 29 <= (x.Expiry - self.Time).days <= 36]
            puts = [x for x in puts if x.Greeks and x.Greeks.Delta is not None]
            if not puts:
                continue

            # Select put with delta closest to -0.35
            target_delta = -self.delta_put / 100
            selected = sorted(puts, key=lambda x: abs(x.Greeks.Delta - target_delta))[0]

            if not selected:
                continue

            option_symbol = selected.Symbol

            if option_symbol not in self.Securities:
                self.AddOptionContract(option_symbol, Resolution.Daily)

            # Liquidate previous hedge if needed
            if self.stock_data[symbol]["put"] is not None:
                self.Liquidate(self.stock_data[symbol]["put"])

            contracts_to_buy = shares // 100
            self.Buy(option_symbol, contracts_to_buy)
            self.stock_data[symbol]["put"] = option_symbol

            # Mark as hedged
            del self.pending_hedges[symbol]