"""
Momentum and Historical Band Ceiling Sizing Algorithm (v3 - Surgical Fixes)

Changes from v1 (minimal, critical only):
1. Real fees + slippage (InteractiveBrokersFeeModel + 10bps)
2. Delta-based execution (no more full liquidate-then-rebuy)
3. stretch_max replaced with rolling stretch_win (fixes survivorship bias)
4. Exhaustion scaling bug fixed (min() instead of override)
5. 180-day hard timeout on risk-off regime (prevents getting permanently stuck)
6. Live market cap check at rebalance time (fixes BBIO-type universe leak)

Everything else — Fibonacci bands, ADX filter, breadth logic, momentum scoring,
universe construction — is identical to v1.
"""

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

# ====================================================
# Sector-Neutral Large-Cap Universe (unchanged from v1)
# ====================================================
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


# ====================================================
# Main Algorithm
# ====================================================
class StockOnlyMomentumV3(QCAlgorithm):

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

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

        # --------------------
        # Band parameters (unchanged)
        # --------------------
        self.band_len = 189
        self.hist_len = 126

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

        # -------- BREADTH STATE (unchanged) --------
        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: track entry date for hard timeout

        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.close_win = {}
        self.stretch_ema = {}
        self.band_hist = {}
        self.stretch_win = {}   # FIX 3: rolling window replaces stretch_max dict

        # Delta execution threshold
        self.rebalance_threshold = 0.02  # FIX 2: only trade if drift > 2%

        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:
            # FIX 1: Realistic fees and slippage
            sec.SetFeeModel(InteractiveBrokersFeeModel())
            sec.SetSlippageModel(ConstantSlippageModel(0.001))  # 10bps round-trip

            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)   # FIX 3

    # --------------------------------------------------
    def OnData(self, data):
        """Unchanged from v1 — uses fixed Fibonacci bands for breadth tracking."""
        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)

            # FIX 3: rolling stretch window instead of lifetime peak
            self.stretch_win[s].Add(stretch)

            # Fibonacci bands (unchanged from v1)
            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):
        """Unchanged from v1."""
        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

        # -------- UNIVERSE-WIDE BREADTH (unchanged logic) --------
        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)

        # -------- BREADTH REGIME --------
        if bottom_frac >= 0.45:
            if not self.was_risk_off:
                self.risk_off_date = self.Time   # FIX 5: record entry date
            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

            # FIX 5: hard 180-day timeout — never stuck risk-off indefinitely
            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}. 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
                self.risk_off_date = None

        else:
            self.allow_universe = True

        if not self.allow_universe:
            self._execute_targets({})   # FIX 2: delta execution even for full exit
            self.Debug(f"RISK-OFF @ {self.Time.strftime('%Y-%m-%d')}: stress={bottom_frac:.1%}")
            return

        # -------- MOMENTUM SCORING (unchanged) --------
        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: live market cap check — universe selection can go stale
            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]

        # -------- BAND SIZING (unchanged from v1) --------
        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)

            # FIX 3 + FIX 4: exhaustion check using rolling stretch_win + min() bug fix
            if self.stretch_win[s].IsReady:
                stretch_list = list(self.stretch_win[s])
                current_stretch = stretch_list[0]       # most recent value
                peak_stretch = max(stretch_list)        # rolling peak, not lifetime peak

                if idx >= 10 and peak_stretch > 0:
                    if current_stretch < (peak_stretch * 0.80):
                        scale = min(scale, 0.2)         # FIX 4: min() not override
                        self.Debug(f"EXHAUSTION: Scaling down {s.Value}")

            scaled[s] = momentum[s] * scale

        # -------- WEIGHTING (unchanged) --------
        if not scaled:
            self._execute_targets({})
            self.Debug("No assets to trade.")
            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 = {}

        # FIX 2: delta execution
        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.strftime('%Y-%m-%d')}: {output}")

    # --------------------------------------------------
    def _execute_targets(self, target_weights):
        """
        FIX 2: Delta-based execution.
        Only trades positions where drift from target exceeds 2%.
        Positions not in target_weights are closed (target = 0).
        Replaces the costly Liquidate()-then-SetHoldings() pattern.
        """
        total_value = self.Portfolio.TotalPortfolioValue
        if total_value <= 0:
            return

        current_weights = {}
        for kvp in self.Portfolio:
            s = kvp.Key
            holding = kvp.Value
            if holding.Invested:
                current_weights[s] = holding.HoldingsValue / total_value

        all_symbols = set(list(current_weights.keys()) + list(target_weights.keys()))

        for s in all_symbols:
            target  = target_weights.get(s, 0.0)
            current = current_weights.get(s, 0.0)
            if abs(target - current) > self.rebalance_threshold:
                self.SetHoldings(s, target)