from AlgorithmImports import *


class BinanceTopCapMomentum(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2021, 1, 1)
        self.SetCash(100000)

        self.SetBrokerageModel(BrokerageName.Binance, AccountType.Margin)

        self.UniverseSettings.Resolution = Resolution.Daily
        self.UniverseSettings.Leverage = 1.0  # exposure controlled by weights

        # --- Strategy parameters
        self.N_HOLD = 5
        self.WEIGHT = 0.05          # 5% each => 25% total exposure
        self.MOM_LOOKBACK = 252     # ~12 months (trading-days convention)

        # --- "Top 10 market cap USDT pairs"
        # NOTE: QC Binance universe doesn't provide market cap ranking,
        # so we use a static list that approximates "top market cap".
        # You can modify this list anytime.
        self.top10_usdt_tickers = [
            "BTCUSDT", "ETHUSDT", "BNBUSDT", "XRPUSDT", "SOLUSDT",
            "ADAUSDT", "DOGEUSDT", "TRXUSDT", "DOTUSDT", "MATICUSDT"
        ]

        # --- Add stablecoin pairs on top of the universe
        self.extra_pairs = ["USDCUSDT", "TUSDUSDT"]

        self.rebalance = False
        self.current_universe = []

        # Add a Binance universe, but we will *only* select our chosen symbols from it
        self.AddUniverse(CryptoUniverse.Binance(self.SelectUniverse))

        # Also explicitly add the stablecoin pairs so they exist even if universe data is sparse
        for ticker in self.extra_pairs:
            self.AddCrypto(ticker, Resolution.Daily, Market.Binance)

        # Monthly rebalance: first trading day of each month at 00:00
        self.Schedule.On(
            self.DateRules.MonthStart(),
            self.TimeRules.At(0, 0),
            self.TriggerRebalance
        )

        self.SetWarmUp(self.MOM_LOOKBACK + 5, Resolution.Daily)

    def TriggerRebalance(self):
        self.rebalance = True

    def SelectUniverse(self, universe):
        """
        Universe contains Binance crypto pair snapshots.
        We select:
          - the static top-10 "market cap" USDT pairs, plus
          - USDCUSDT and DAIUSDT
        And we also avoid non-USDT quotes to prevent FX conversion issues (e.g., NGN).
        """
        # Build a set of available symbols from the incoming universe slice
        available = set()
        for c in universe:
            sym = getattr(c, "Symbol", None)
            if sym is None:
                continue

            # Avoid NGN and other unsupported quote currencies; keep USDT pairs only
            if not sym.Value.endswith("USDT"):
                continue
            available.add(sym.Value)

        desired = list(self.top10_usdt_tickers) + list(self.extra_pairs)

        # Keep only what is actually available in this runtime/feed
        selected_values = [v for v in desired if v in available]

        # Convert to Symbol objects that QC expects from universe selection
        selected_symbols = [Symbol.Create(v, SecurityType.Crypto, Market.Binance) for v in selected_values]

        self.current_universe = selected_symbols
        return selected_symbols

    def OnSecuritiesChanged(self, changes: SecurityChanges):
        for sec in changes.AddedSecurities:
            sec.SetLeverage(1.0)

    def OnData(self, data: Slice):
        if self.IsWarmingUp or not self.rebalance:
            return
        self.rebalance = False

        symbols = list(self.current_universe)
        # Ensure stablecoin pairs are included even if not selected via universe
        for ticker in self.extra_pairs:
            sym = self.Symbol(ticker)
            if sym not in symbols:
                symbols.append(sym)

        # Pull history for momentum ranking
        hist = self.History(symbols, self.MOM_LOOKBACK + 1, Resolution.Daily)
        if hist.empty:
            return

        mom = []
        for sym in symbols:
            try:
                df = hist.loc[sym]
            except Exception:
                continue

            if df is None or len(df.index) < (self.MOM_LOOKBACK + 1):
                continue
            if "close" not in df.columns:
                continue

            closes = df["close"].dropna()
            if len(closes) < (self.MOM_LOOKBACK + 1):
                continue

            close_now = float(closes.iloc[-1])
            close_then = float(closes.iloc[-(self.MOM_LOOKBACK + 1)])
            if close_then <= 0:
                continue

            momentum_12m = close_now / close_then - 1.0
            mom.append((sym, momentum_12m))

        if len(mom) == 0:
            return

        # Top 5 by 12-month momentum
        mom.sort(key=lambda x: x[1], reverse=True)
        selected = [x[0] for x in mom[:self.N_HOLD]]

        # Liquidate anything not selected
        for sym in self.Securities.Keys:
            if self.Portfolio[sym].Invested and sym not in selected:
                self.Liquidate(sym)

        # Allocate 5% each (25% total exposure)
        for sym in selected:
            if sym in self.Securities and self.Securities[sym].HasData:
                self.SetHoldings(sym, self.WEIGHT)

        self.Debug(f"{self.Time.date()} Selected: {', '.join([s.Value for s in selected])}")