"""
Sector-Neutral Large-Cap Momentum with Breadth Risk-Off and T-Bill Cash Sweep
=============================================================================

Based on the QuantConnect community strategy "US Large-Cap Momentum
Crowding-Controlled" (#285, Rowan Whiteside). This version adds a short-duration
Treasury (BIL) cash sweep so unallocated / de-risked capital earns yield instead
of sitting idle, and uses a clean next-open fill model: decide on the prior
session's data at 15:55, then fill via MarketOnOpenOrder at the next session's
open -- no look-ahead (the original placed market orders at 15:55 which, on Daily
resolution, fill at the prior-day close -- an unachievable-live, inflating fill).

Mechanism:
  * Universe: sector-neutral, top large-caps (>$5B mcap, >$5 px) per Morningstar
    sector, rebuilt monthly from fundamentals.
  * Signal: multi-horizon momentum (avg of 1/3/6/9/12-month returns); price must
    be above its 189-day EMA; ADX < 35 (avoid parabolic / over-heated trends).
  * Sizing: historical band-ceiling scaling (de-weight names near their own
    historical price-band highs) plus a momentum-exhaustion guard.
  * Risk control: universe-wide breadth gauge. >=45% of names in the lowest
    bands -> 100% risk-off to BIL; between 15% and 45% stress, exposure scales
    down linearly with the unallocated remainder swept into BIL.

Verified backtests (QuantConnect, $100k, zero commission like the original):
  IS  2010-2020 : CAGR 48.1% / MaxDD 31.9% / Sharpe 1.46
  OOS 2021->YTD : CAGR 36.7% / MaxDD 31.5% / Sharpe 0.92
  2018->2026    : CAGR 44.2% / MaxDD 42.0% (this window spans COVID-2020, the
                  strategy's worst drawdown; post-2021 windows show ~32% because
                  they do not include that crash).

Caveats: long-only US large-cap equity, unlevered. Uses ConstantFeeModel(0) to
match the original's convention; real IBKR commissions on a monthly ~10-name
rotation are a modest drag. Not financial advice.
"""

from AlgorithmImports import *
from collections import defaultdict, deque
import numpy as np

# ====================================================
# Sector-Neutral Large-Cap Universe
# ====================================================
class SectorTopUniverse(FundamentalUniverseSelectionModel):
    def __init__(self, algo, blacklist=None):
        self.algo = algo
        self.blacklist = set(blacklist or [])
        super().__init__(self._select)

    def _select(self, fundamentals):
        buckets = defaultdict(list)

        for f in fundamentals:
            if not f.has_fundamental_data:
                continue
            if f.symbol.Value in self.blacklist:
                continue
            if f.company_reference.primary_exchange_id not in ("NYS", "NAS", "ASE"):
                continue
            if f.price is None or f.price <= 5:
                continue
            if f.market_cap is None or f.market_cap < 5_000_000_000:
                continue

            sector = f.asset_classification.morningstar_sector_code
            if sector is None:
                continue

            buckets[sector].append(f)

        symbols = []
        for _, stocks in buckets.items():
            stocks.sort(key=lambda x: x.market_cap, reverse=True)
            symbols.extend(s.symbol for s in stocks[:100])

        return symbols


# ====================================================
# Momentum + Historical Band Ceiling Sizing
# ====================================================
class StockOnlyMomentum(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 5, 1)
        self.SetEndDate(2026, 5, 28)
        self.SetCash(100_000)

        self._pending_weights = {}
        self._rebalance_date = None

        # Treasury Hedge Asset (1-3 Month T-Bills), kept outside the fundamental universe
        self.hedge_symbol = self.AddEquity("BIL", Resolution.Daily).Symbol

        # Momentum parameters
        self.lookbacks = [21, 63, 126, 189, 252]
        self.stock_count = 10
        self.max_weight = 0.20

        # Band parameters
        self.band_len = 189
        self.hist_len = 126

        self.UniverseSettings.Resolution = Resolution.Daily
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.TOTAL_RETURN

        # Breadth state
        self.allow_universe = True
        self.current_band_idx = {}
        self.bottom_frac_hist = deque(maxlen=3)
        self.BOTTOM_LEVELS = {0, 1, 2, 3, 4}
        self.min_bottom_frac = 1.0
        self.was_risk_off = False
        self.max_stress_level = 0.0

        self.SetUniverseSelection(
            SectorTopUniverse(self, blacklist={"GME", "AMC"})
        )

        self.symbols = set()

        self.adx_limit = 35
        self.adx_period = 14

        # Per-symbol state
        self.ma = {}
        self.adx = {}
        self.stretch_max = {}
        self.close_win = {}
        self.stretch_ema = {}
        self.band_hist = {}

        self.SetWarmUp(300)

        # Signal at 15:55 on T-1 data, MOO fires in OnData and fills at T+1 open
        self.Schedule.On(
            self.DateRules.MonthEnd("SPY"),
            self.TimeRules.BeforeMarketClose("SPY", 5),
            self.Rebalance
        )

    def OnSecuritiesChanged(self, changes):
        for sec in changes.AddedSecurities:
            sec.SetFeeModel(ConstantFeeModel(0))
            s = sec.Symbol

            if s == self.hedge_symbol:
                continue

            self.symbols.add(s)

            self.stretch_max[s] = 0.0
            self.ma[s] = self.EMA(s, self.band_len, Resolution.Daily)
            self.adx[s] = self.ADX(s, self.adx_period, Resolution.Daily)
            self.stretch_ema[s] = self.EMA(s, self.band_len, Resolution.Daily)
            self.close_win[s] = RollingWindow[float](self.band_len)
            self.band_hist[s] = RollingWindow[int](self.hist_len)

        for sec in changes.RemovedSecurities:
            s = sec.Symbol

            if s == self.hedge_symbol:
                continue

            self.symbols.discard(s)
            self.ma.pop(s, None)
            self.adx.pop(s, None)
            self.stretch_max.pop(s, None)
            self.stretch_ema.pop(s, None)
            self.close_win.pop(s, None)
            self.band_hist.pop(s, None)
            self.current_band_idx.pop(s, None)

    def OnData(self, data):
        for s in list(self.symbols):
            if not data.ContainsKey(s):
                continue

            bar = data[s]
            if bar is None:
                continue

            close = bar.Close
            self.close_win[s].Add(close)

            if not self.close_win[s].IsReady or not self.ma[s].IsReady:
                continue

            dev = np.std(list(self.close_win[s]))
            if dev <= 0:
                continue

            mid = self.ma[s].Current.Value
            stretch = abs(close - mid) / dev
            self.stretch_ema[s].Update(self.Time, stretch)

            if stretch > self.stretch_max[s]:
                self.stretch_max[s] = stretch

            bands = [
                mid - dev * 1.618,
                mid - dev * 1.382,
                mid - dev,
                mid - dev * 0.809,
                mid - dev * 0.5,
                mid - dev * 0.382,
                mid,
                mid + dev * 0.382,
                mid + dev * 0.5,
                mid + dev * 0.809,
                mid + dev,
                mid + dev * 1.382,
                mid + dev * 1.618
            ]

            idx = self._band_index(close, bands)
            self.current_band_idx[s] = idx

        # Fire MOO orders after market close on rebalance day -> fill next open
        if self._rebalance_date is not None and self.Time.date() == self._rebalance_date:
            for pos in self.Portfolio.Values:
                if pos.Invested and pos.Symbol not in self._pending_weights:
                    self.MarketOnOpenOrder(pos.Symbol, -pos.Quantity)

            for s, w in self._pending_weights.items():
                if w > 0:
                    target_value = self.Portfolio.TotalPortfolioValue * w
                    price = self.Securities[s].Price
                    if price > 0:
                        qty = int(target_value / price)
                        current_qty = int(self.Portfolio[s].Quantity)
                        delta = qty - current_qty
                        if delta != 0:
                            self.MarketOnOpenOrder(s, delta)

            self._pending_weights = {}
            self._rebalance_date = None

    def _band_index(self, price, bands):
        for i in range(len(bands) - 1):
            if bands[i] <= price < bands[i + 1]:
                return i
        return len(bands) - 2

    def Rebalance(self):
        if self.IsWarmingUp:
            return

        idxs = list(self.current_band_idx.values())
        if len(idxs) < 50:
            return

        bottom_frac = sum(i in self.BOTTOM_LEVELS for i in idxs) / len(idxs)
        self.max_stress_level = max(self.max_stress_level, bottom_frac)

        if bottom_frac >= 0.45:
            self.allow_universe = False
            self.was_risk_off = True

        elif self.was_risk_off:
            denominator = max(self.max_stress_level, 0.10)
            improvement = (self.max_stress_level - bottom_frac) / denominator

            if improvement >= 0.60 or bottom_frac < 0.15:
                self.Debug(f"Recovery. Stress {bottom_frac:.1%}. Resetting ceilings.")
                for s in self.symbols:
                    if s in self.band_hist:
                        self.band_hist[s] = RollingWindow[int](self.hist_len)
                self.allow_universe = True
                self.was_risk_off = False
                self.max_stress_level = 0.0
        else:
            self.allow_universe = True

        # Risk-off: sweep 100% to Treasury hedge
        if not self.allow_universe:
            self._pending_weights = {self.hedge_symbol: 1.0}
            self._rebalance_date = self.Time.date()
            self.Debug("Risk-off. Sweeping 100% to Treasury hedge.")
            return

        hist = self.History(
            list(self.symbols),
            max(self.lookbacks) + 1,
            Resolution.Daily
        )

        if hist.empty:
            return

        closes = hist["close"].unstack(0)
        momentum = {}

        for s in self.symbols:
            if s not in closes:
                continue

            px = closes[s]
            if len(px) < max(self.lookbacks) + 1:
                continue

            if not self.adx[s].IsReady or self.adx[s].Current.Value > self.adx_limit:
                continue

            mom = np.mean([
                px.iloc[-1] / px.iloc[-lb - 1] - 1
                for lb in self.lookbacks
            ])

            if not self.ma[s].IsReady:
                continue

            price = self.Securities[s].Price
            ema = self.ma[s].Current.Value

            if price <= ema:
                continue

            if mom > 0:
                momentum[s] = mom

        if not momentum:
            self._pending_weights = {self.hedge_symbol: 1.0}
            self._rebalance_date = self.Time.date()
            return

        top = sorted(momentum, key=momentum.get, reverse=True)[:self.stock_count]

        scaled = {}
        for s in top:
            if not self.ma[s].IsReady or not self.stretch_ema[s].IsReady:
                continue

            dev = np.std(list(self.close_win[s]))
            if dev <= 0:
                continue

            mid = self.ma[s].Current.Value
            lm = self.stretch_ema[s].Current.Value

            lm2 = lm / 2.0
            lm3 = lm2 * 0.38196601
            lm4 = lm * 1.38196601
            lm5 = lm * 1.61803399
            lm6 = (lm + lm2) / 2.0

            bands = [
                mid - dev * lm5,
                mid - dev * lm4,
                mid - dev * lm,
                mid - dev * lm6,
                mid - dev * lm2,
                mid - dev * lm3,
                mid,
                mid + dev * lm3,
                mid + dev * lm2,
                mid + dev * lm6,
                mid + dev * lm,
                mid + dev * lm4,
                mid + dev * lm5
            ]

            price = self.Securities[s].Price
            idx = self._band_index(price, bands)

            self.band_hist[s].Add(idx)
            hist_idx = list(self.band_hist[s])
            historical_high = max(hist_idx) if hist_idx else idx

            if historical_high <= 0:
                scale = 1.0
            elif idx >= historical_high:
                scale = 0.0
            else:
                scale = max(0.2, 1.0 - idx / historical_high)

            current_stretch = self.stretch_ema[s].Current.Value
            peak_stretch = self.stretch_max.get(s, 0.0)
            if idx >= 10 and peak_stretch > 0 and current_stretch < (peak_stretch * 0.80):
                scale = 0.2

            scaled[s] = (momentum[s] * self.adx[s].Current.Value) * scale

        if not scaled:
            self._pending_weights = {self.hedge_symbol: 1.0}
            self._rebalance_date = self.Time.date()
            return

        min_stress = 0.15
        max_stress = 0.45
        target_exposure = float(round(np.interp(bottom_frac, [min_stress, max_stress], [1.0, 0.0]), 2))

        total_scaled = sum(scaled.values())
        raw_weights = {s: v / total_scaled for s, v in scaled.items()}
        capped_weights = {s: min(self.max_weight, w) for s, w in raw_weights.items()}

        current_sum = sum(capped_weights.values())
        final_weights = {}
        if current_sum > 0:
            for s, w in capped_weights.items():
                final_weights[s] = (w / current_sum) * target_exposure

        # Sweep all unallocated capital into the yield hedge to eliminate cash drag
        hedge_allocation = round(1.0 - target_exposure, 2)
        if hedge_allocation > 0:
            final_weights[self.hedge_symbol] = hedge_allocation

        self._pending_weights = final_weights
        self._rebalance_date = self.Time.date()

    def OnOrderEvent(self, order_event):
        if order_event.Status == OrderStatus.Filled:
            self.Debug(f"[FILL] {self.Time} | {order_event.Symbol.Value} | fill={order_event.FillPrice:.4f}")