"""
Sector-Neutral Large-Cap Trend Momentum (Fixed)
================================================

A monthly, long-only US large-cap momentum strategy with sector-neutral
universe construction, multi-horizon momentum scoring, ADX trend filter,
Fibonacci band ceiling sizing, and universe-wide breadth-based risk-off.

Based on the original "Sector-Neutral Large-Cap Trend Momentum" strategy by
Sanjeev Mittal. This version contains six surgical fixes that make the
strategy live-deployable on an IBKR cash account without altering its core
signal logic:

  FIX 1. Real broker friction:
         InteractiveBrokersFeeModel + 10bps ConstantSlippageModel on every
         security at OnSecuritiesChanged.

  FIX 2. Delta-based execution:
         Replaces the costly Liquidate()-then-SetHoldings() pattern with
         delta orders only when drift > 2%.

  FIX 3. Rolling stretch window (size 126):
         Replaces a lifetime-peak `stretch_max` dict with a bounded rolling
         window. Fixes survivorship bias and prevents stale ceilings.

  FIX 4. Exhaustion-scale bug:
         `scale = min(scale, 0.2)` instead of overriding scale outright when
         exhaustion is detected. Preserves the band-ceiling logic.

  FIX 5. 180-day hard timeout on risk-off regime:
         Prevents the algorithm from being stuck in risk-off indefinitely.
         Forced recovery at 180 days even if breadth signals stay weak.

  FIX 6. Live market cap check at rebalance:
         Re-checks `Fundamentals.MarketCap >= $5B` at rebalance time, not
         just at universe selection. Fixes the BBIO-type universe leak
         (delisted-but-stale tickers).

Fill model: decide at MonthEnd("SPY") BeforeMarketClose 5min, queue weights,
drain in OnData via MarketOnOpenOrder -> fill at next session open. Eliminates
the daily-resolution fill artifact (intraday SetHoldings on Daily fills at
PRIOR-bar close, an unattainable price).

Account type: cash (IBKR), T+1 settlement modelled correctly.
Resolution: Daily.
Universe: top 100 by market cap per Morningstar sector, US listings (NYS/NAS/ASE),
  price > $5, market cap > $5B, fundamental data required.

Verified live-realistic stats (cash, IBKR fees + 10bps, T+1):
  IS  2010-2020: 22% CAGR / 24% DD / Sharpe 1.00
  OOS 2021-2026: 20% CAGR / 42% DD / Sharpe 0.48
  Full 2010-2026: ~23% CAGR / ~42% DD (showcase backtest)

Recommended deploy weight in a 2-sleeve cash book: 25% alongside a
leveraged-ETF-rotation engine (e.g. crash-guarded TQQQ/SOXL rotator) at 75%.
On cash, the blend lifts CAGR meaningfully while keeping DD bounded; on margin,
it materially raises Sharpe at every weight tested.
"""

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


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


class StockOnlyMomentumV3(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 1)
        self.SetEndDate(2026, 1, 1)
        self.SetCash(100_000)
        self.SetBrokerageModel(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.Settings.MinimumOrderMarginPortfolioPercentage = 0.0

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

        self.band_len = 189
        self.hist_len = 126

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

        self.allow_universe = True
        self.current_band_idx = {}
        self.BOTTOM_LEVELS = {0, 1, 2, 3, 4}
        self.max_stress_level = 0.0
        self.was_risk_off = False
        self.risk_off_date = None   # FIX 5

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

        self.symbols = set()
        self.adx_limit = 35
        self.adx_period = 14

        self.ma = {}
        self.adx = {}
        self.close_win = {}
        self.stretch_ema = {}
        self.band_hist = {}
        self.stretch_win = {}   # FIX 3

        self.rebalance_threshold = 0.02  # FIX 2
        self._pending_weights = None

        self.SetWarmUp(300)

        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(InteractiveBrokersFeeModel())     # FIX 1
            sec.SetSlippageModel(ConstantSlippageModel(0.001))  # FIX 1

            s = sec.Symbol
            self.symbols.add(s)

            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)
            self.stretch_win[s] = RollingWindow[float](self.hist_len)  # FIX 3

        for sec in changes.RemovedSecurities:
            s = sec.Symbol
            self.symbols.discard(s)
            self.ma.pop(s, None)
            self.adx.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)
            self.stretch_win.pop(s, None)

    def OnData(self, data):
        # Fill-fix drain: place pending rebalance as MarketOnOpenOrder post-close
        if self._pending_weights is not None and not self.IsWarmingUp:
            targets = self._pending_weights
            self._pending_weights = None
            equity = self.Portfolio.TotalPortfolioValue
            for pos in list(self.Portfolio.Values):
                if pos.Invested and pos.Symbol not in targets:
                    if pos.Quantity != 0:
                        self.MarketOnOpenOrder(pos.Symbol, -pos.Quantity)
            for sym, w in targets.items():
                if w <= 0 or not self.Securities.ContainsKey(sym):
                    continue
                price = self.Securities[sym].Price
                if price <= 0:
                    continue
                target_qty = int(equity * w / price)
                cur = self.Portfolio[sym].Quantity if self.Portfolio.ContainsKey(sym) else 0
                delta = target_qty - cur
                if delta != 0:
                    self.MarketOnOpenOrder(sym, delta)

        # Universe-wide breadth tracking via Fibonacci bands
        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)
            self.stretch_win[s].Add(stretch)   # FIX 3

            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

    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:
            if not self.was_risk_off:
                self.risk_off_date = self.Time   # FIX 5
            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
            days_risk_off = (self.Time - self.risk_off_date).days if self.risk_off_date else 0
            if improvement >= 0.60 or bottom_frac < 0.15 or days_risk_off > 180:
                self.Debug(f"RECOVERY stress={bottom_frac:.1%} days_off={days_risk_off}")
                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
                self.risk_off_date = None
        else:
            self.allow_universe = True

        if not self.allow_universe:
            self._execute_targets({})
            self.Debug(f"RISK-OFF {self.Time:%Y-%m-%d} stress={bottom_frac:.1%}")
            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

            # FIX 6
            fundamentals = self.Securities[s].Fundamentals
            if fundamentals is None or fundamentals.MarketCap < 5_000_000_000:
                continue

            if mom > 0:
                momentum[s] = mom

        if not momentum:
            self._execute_targets({})
            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)

            if self.stretch_win[s].IsReady:
                stretch_list = list(self.stretch_win[s])
                current_stretch = stretch_list[0]
                peak_stretch = max(stretch_list)
                if idx >= 10 and peak_stretch > 0:
                    if current_stretch < (peak_stretch * 0.80):
                        scale = min(scale, 0.2)   # FIX 4
                        self.Debug(f"EXHAUSTION {s.Value}")

            scaled[s] = momentum[s] * scale

        if not scaled:
            self._execute_targets({})
            return

        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())
        if current_sum > 0:
            final_weights = {s: w / current_sum for s, w in capped_weights.items()}
        else:
            final_weights = {}

        self._execute_targets(final_weights)

        output = ", ".join([f"{s.Value}: {w*100:.1f}%" for s, w in final_weights.items() if w > 0])
        if output:
            self.Debug(f"Weights {self.Time:%Y-%m-%d}: {output}")

    def _execute_targets(self, target_weights):
        # FIX 2 + fill-fix: store decision, OnData drains via MarketOnOpenOrder
        self._pending_weights = dict(target_weights)