# region imports
from AlgorithmImports import *
from collections import defaultdict, deque
import calendar
import numpy as np
import pandas as pd
# endregion

# =============================================================================
# FROZEN — imports only.  The mutator owns the ENTIRE algorithm: universe model,
# every hook, every method, the class itself.  Maximum flexibility.
#
# Conventions (not hash-enforced; sys_msg.md flags them as expectations):
#   - Read backtest window from QC parameters via self.get_parameter so one
#     file covers IS / OOS (evaluate.py drives the window).
#   - Keep a deployable cost model + PLAN long-only / no-margin posture; the
#     score is what is judged.
# =============================================================================

# EVOLVE-BLOCK-START
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 StockOnlyMomentum(QCAlgorithm):
    def Initialize(self):
        sy = int(self.get_parameter("start_year", "2010"))
        sm = int(self.get_parameter("start_month", "1"))
        ey = int(self.get_parameter("end_year", "2020"))
        em = int(self.get_parameter("end_month", "12"))
        ed = calendar.monthrange(ey, em)[1]
        self.SetStartDate(sy, sm, 1)
        self.SetEndDate(ey, em, ed)
        self.SetCash(10_000)
        # Gen32 proven short-weighted lookbacks
        self.lookbacks = [21, 42, 63, 126, 189]
        self.mom_weights = [0.25, 0.20, 0.20, 0.20, 0.15]
        # Gen100 proven: concentrate in choppy (5), base 10
        self.stock_count_base = 10
        self.stock_count_choppy = 5
        self.choppy_threshold = 0.25
        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.was_risk_off = False
        self.max_stress_level = 0.0
        # Clean scalar for stress RoC
        self.prev_bottom_frac = None
        # Portfolio equity curve tracking for conditional brake
        self.portfolio_peak = 10_000.0
        self.SetUniverseSelection(
            SectorTopUniverse(self, blacklist={"GME", "AMC"})
        )
        self.symbols = set()
        self.symbol_to_sector = {}
        self.adx_limit = 35
        self.adx_period = 14
        self.ma = {}
        self.adx = {}
        self.stretch_max = {}
        self.close_win = {}
        self.stretch_ema = {}
        self.band_hist = {}
        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()); sec.SetSlippageModel(ConstantSlippageModel(0.001))
            s = sec.Symbol
            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)
            try:
                sector = sec.Fundamentals.AssetClassification.MorningstarSectorCode
                if sector is not None and sector != 0:
                    self.symbol_to_sector[s] = sector
            except Exception:
                pass
        for sec in changes.RemovedSecurities:
            s = sec.Symbol
            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)
            self.symbol_to_sector.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
            ]
            self.current_band_idx[s] = self._band_index(close, bands)
    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
        # Stock-level breadth
        stock_bottom_frac = sum(i in self.BOTTOM_LEVELS for i in idxs) / len(idxs)
        # Sector-level breadth: fraction of sectors where majority in bottom bands
        sector_bottoms = defaultdict(list)
        for s, idx in self.current_band_idx.items():
            sector = self.symbol_to_sector.get(s)
            if sector is not None:
                sector_bottoms[sector].append(idx in self.BOTTOM_LEVELS)
        if sector_bottoms:
            sector_stress_count = sum(
                1 for flags in sector_bottoms.values()
                if len(flags) > 0 and sum(flags) / len(flags) > 0.50
            )
            sector_bottom_frac = sector_stress_count / len(sector_bottoms)
            # 50/50 blend: sector signal smooths intra-sector noise (gen82 proven)
            bottom_frac = 0.5 * stock_bottom_frac + 0.5 * sector_bottom_frac
        else:
            bottom_frac = stock_bottom_frac
        # Stress rate of change (clean scalar, not deque)
        stress_roc = 0.0
        if self.prev_bottom_frac is not None:
            stress_roc = bottom_frac - self.prev_bottom_frac
        self.prev_bottom_frac = bottom_frac
        self.max_stress_level = max(self.max_stress_level, bottom_frac)
        # Portfolio equity curve tracking for conditional brake
        current_value = self.Portfolio.TotalPortfolioValue
        self.portfolio_peak = max(self.portfolio_peak, current_value)
        portfolio_dd = (self.portfolio_peak - current_value) / self.portfolio_peak if self.portfolio_peak > 0 else 0.0
        # Risk-off latch
        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:
                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.prev_bottom_frac = None  # fresh RoC after recovery
                # Reset peak to current (post-crash) level so brake starts fresh
                self.portfolio_peak = current_value
                # Begin graduated re-entry: cap exposure for 3 months post-crash
                self.in_recovery = True
                self.recovery_months = 0
        else:
            self.allow_universe = True
        if not self.allow_universe:
            self.Liquidate()
            return
        # Dynamic position count: concentrate in choppy regimes (gen100 proven)
        stock_count = self.stock_count_choppy if bottom_frac >= self.choppy_threshold else self.stock_count_base
        hist = self.History(list(self.symbols), max(self.lookbacks) + 1, Resolution.Daily)
        if hist.empty:
            return
        closes = hist["close"].unstack(0)
        momentum = {}
        consistency_map = {}
        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
            # Compute per-lookback returns once (reused for both momentum and consistency)
            lb_returns = [px.iloc[-1] / px.iloc[-lb - 1] - 1 for lb in self.lookbacks]
            mom = sum(w * r for w, r in zip(self.mom_weights, lb_returns))
            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
                # Consistency: fraction of lookback horizons with positive return
                n_positive = sum(1 for r in lb_returns if r > 0)
                consistency_map[s] = n_positive / len(self.lookbacks)
        if not momentum:
            self.Liquidate()
            return
        top = sorted(momentum, key=momentum.get, reverse=True)[:stock_count]
        scaled = {}
        vol_20d = {}
        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:
                if current_stretch < (peak_stretch * 0.80):
                    scale = 0.2
            # Momentum consistency soft multiplier (NEW):
            # Stocks confirmed positive across all lookbacks get full weight;
            # stocks with mixed horizon signals are down-weighted.
            # Range [0.6, 1.0] — tilts toward multi-horizon-confirmed trends
            # without over-penalizing short-term entries in new bull legs.
            # Only affects RELATIVE weights within top-K, not total exposure.
            consistency = consistency_map.get(s, 0.6)
            consist_factor = 0.6 + 0.4 * consistency
            scaled[s] = momentum[s] * scale * consist_factor
            # 20-day vol for position cap
            px = closes[s]
            if len(px) >= 21:
                vol_20d[s] = float(px.pct_change().iloc[-20:].std())
            else:
                vol_20d[s] = 0.015
        if not scaled:
            self.Liquidate()
            return
        # Base exposure from breadth interpolation
        min_stress = 0.15
        max_stress = 0.45
        target_exposure = float(np.interp(bottom_frac, [min_stress, max_stress], [1.0, 0.0]))
        # RoC early warning: accelerating stress → cut exposure 30%
        if stress_roc >= 0.15:
            target_exposure *= 0.70
        # Conditional portfolio equity brake:
        # Only fires when BOTH portfolio is materially down AND breadth is worsening.
        # Threshold lowered from 8%→7% to catch 2011/2018 crashes ~1 month earlier;
        # stress_roc > 0.03 gate prevents false positives in stable-breadth corrections.
        if portfolio_dd > 0.07 and stress_roc > 0.03:
            # Smooth reduction: no cut at dd=7%, full 75% cut at dd=29%
            dd_brake = max(0.25, 1.0 - (portfolio_dd - 0.07) / 0.22)
            target_exposure = min(target_exposure, dd_brake)
        target_exposure = float(round(max(0.0, min(1.0, target_exposure)), 2))
        # Weight computation
        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()}
        # Volatility cap without renormalization (gen100 proven)
        for s in list(capped_weights.keys()):
            vol = vol_20d.get(s, 0.015)
            if vol > 0:
                vol_cap_factor = min(1.0, 0.015 / vol)
                capped_weights[s] = min(capped_weights[s], self.max_weight * vol_cap_factor)
        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
        for pos in self.Portfolio.Values:
            if pos.Invested and pos.Symbol not in final_weights:
                self.Liquidate(pos.Symbol)
        for s, w in final_weights.items():
            if w > 0:
                self.SetHoldings(s, w)
# EVOLVE-BLOCK-END