# 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):
    """Sector-neutral large-cap universe: top-N by market cap per Morningstar sector."""
    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)
        self.lookbacks = [21, 42, 63, 126, 189]
        self.mom_weights_bull = [0.25, 0.20, 0.20, 0.20, 0.15]
        self.mom_weights_choppy = [0.10, 0.15, 0.20, 0.25, 0.30]
        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
        self.prev_bottom_frac = None
        self.bull_streak = 0
        self.portfolio_peak = 10_000.0
        self.risk_off_months = 0
        self.riskoff_threshold = 0.47
        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.bond_syms = set()
        _ief = self.AddEquity("IEF", Resolution.Daily)
        _ief.SetFeeModel(InteractiveBrokersFeeModel()); _ief.SetSlippageModel(ConstantSlippageModel(0.001))
        self.bond_etf = _ief.Symbol
        self.bond_syms.add(self.bond_etf)
        _spy = self.AddEquity("SPY", Resolution.Daily)
        _spy.SetFeeModel(InteractiveBrokersFeeModel())
        self.spy_symbol = _spy.Symbol
        self.bond_syms.add(self.spy_symbol)
        self.spy_sma200 = self.SMA("SPY", 200, Resolution.Daily)
        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:
            s = sec.Symbol
            if s in self.bond_syms:
                continue
            sec.SetFeeModel(InteractiveBrokersFeeModel()); sec.SetSlippageModel(ConstantSlippageModel(0.001))
            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_code = sec.Fundamentals.AssetClassification.MorningstarSectorCode
                if sector_code is not None and sector_code != 0:
                    self.symbol_to_sector[s] = sector_code
            except Exception:
                pass
        for sec in changes.RemovedSecurities:
            s = sec.Symbol
            if s in self.bond_syms:
                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)
            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_bottom_frac = sum(i in self.BOTTOM_LEVELS for i in idxs) / len(idxs)
        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)
            bottom_frac = 0.5 * stock_bottom_frac + 0.5 * sector_bottom_frac
        else:
            bottom_frac = stock_bottom_frac
        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)
        if bottom_frac < 0.20:
            self.bull_streak += 1
        else:
            self.bull_streak = 0
        stress_improvement_velocity = max(0.0, -stress_roc)
        velocity_threshold = max(0.25, 0.55 - 3.0 * stress_improvement_velocity)
        time_threshold = max(0.35, 0.55 - 0.10 * max(0, self.risk_off_months - 1))
        adaptive_recovery_threshold = min(velocity_threshold, time_threshold)
        current_value = self.Portfolio.TotalPortfolioValue
        self.portfolio_peak = max(self.portfolio_peak, current_value)
        if bottom_frac >= self.riskoff_threshold:
            self.allow_universe = False
            self.was_risk_off = True
            self.risk_off_months += 1
        elif self.was_risk_off:
            denominator = max(self.max_stress_level, 0.10)
            improvement = (self.max_stress_level - bottom_frac) / denominator
            if improvement >= adaptive_recovery_threshold 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
                self.bull_streak = 0
                self.risk_off_months = 0
                self.portfolio_peak = current_value
            else:
                self.risk_off_months += 1
        else:
            self.allow_universe = True
            self.risk_off_months = 0
        if not self.allow_universe:
            self.Liquidate()
            return
        # CONVEX exposure curve: quadratic de-risk keeps more capital deployed
        # during moderate stress (false-positive breadth signals in 2011-12, 2015-16)
        # while still reaching zero at the risk-off threshold.
        # At midpoint stress, exposure = 0.75 instead of linear 0.50.
        min_stress_eff = 0.25 if self.bull_streak >= 3 else 0.15
        stress_range = self.riskoff_threshold - min_stress_eff
        if stress_range > 0 and bottom_frac > min_stress_eff:
            raw_frac = min(1.0, (bottom_frac - min_stress_eff) / stress_range)
            target_exposure = max(0.0, 1.0 - raw_frac * raw_frac)
        else:
            target_exposure = 1.0
        if stress_roc > 0.12:
            target_exposure *= 0.60
        eq_dd = (self.portfolio_peak - current_value) / self.portfolio_peak if self.portfolio_peak > 0 else 0.0
        if eq_dd > 0.15:
            brake_factor = max(0.50, 1.0 - (eq_dd - 0.15) / 0.10)
            target_exposure *= brake_factor
        spy_below_sma = False
        if self.spy_sma200.IsReady:
            spy_price = self.Securities[self.spy_symbol].Price
            spy_sma_val = self.spy_sma200.Current.Value
            if spy_sma_val > 0 and spy_price < spy_sma_val:
                spy_below_sma = True
                pct_below = (spy_sma_val - spy_price) / spy_sma_val
                spy_brake = max(0.40, 1.0 - pct_below * 6.0)
                target_exposure *= spy_brake
        pre_hist = self.History(list(self.symbols), 22, Resolution.Daily)
        if not pre_hist.empty:
            pre_closes = pre_hist["close"].unstack(0)
            if pre_closes.shape[0] >= 22:
                med_ret_21d = float((pre_closes.iloc[-1] / pre_closes.iloc[0]).median()) - 1.0
                if med_ret_21d < -0.05:
                    crash_mult = max(0.30, 1.0 + med_ret_21d * 5.0)
                    target_exposure *= crash_mult
        target_exposure = float(round(max(0.0, min(1.0, target_exposure)), 2))
        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)
        mom_weights = self.mom_weights_choppy if bottom_frac >= self.choppy_threshold else self.mom_weights_bull
        recovering = stress_roc < 0.0
        ema_penalty = 1.0 if recovering else 0.5
        momentum = {}
        for s in self.symbols:
            if s not in closes:
                continue
            px = closes[s]
            if len(px) < max(self.lookbacks) + 1:
                continue
            adx_limit_eff = 45 if bottom_frac >= self.choppy_threshold else self.adx_limit
            if not self.adx[s].IsReady or self.adx[s].Current.Value > adx_limit_eff:
                continue
            lb_returns = [px.iloc[-1] / px.iloc[-lb - 1] - 1 for lb in self.lookbacks]
            mom = sum(w * r for w, r in zip(mom_weights, lb_returns))
            if not self.ma[s].IsReady:
                continue
            price = self.Securities[s].Price
            ema = self.ma[s].Current.Value
            if mom > 0:
                if price <= ema:
                    momentum[s] = mom * ema_penalty
                else:
                    momentum[s] = mom
        if not momentum:
            self.Liquidate()
            return
        INCUMBENT_BUFFER = 3
        current_invested = {pos.Symbol for pos in self.Portfolio.Values
                            if pos.Invested and pos.Quantity > 0}
        ranked = sorted(momentum, key=momentum.get, reverse=True)
        top_definite = ranked[:stock_count]
        top_incumbent = [s for s in ranked[stock_count:stock_count + INCUMBENT_BUFFER]
                         if s in current_invested]
        top = top_definite + top_incumbent
        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:
                if stress_roc < -0.05:
                    dyn_floor = 0.35
                elif stress_roc < 0:
                    dyn_floor = 0.20
                else:
                    dyn_floor = 0.10
                scale = max(dyn_floor, 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
            scaled[s] = momentum[s] * scale
            px_s = closes[s]
            if len(px_s) >= 21:
                vol_20d[s] = float(px_s.pct_change().iloc[-20:].std())
            else:
                vol_20d[s] = 0.015
        if not scaled:
            self.Liquidate()
            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()}
        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
        ief_share = 0.75 if spy_below_sma else 0.50
        bond_alloc = round((1.0 - target_exposure) * ief_share, 2)
        if bond_alloc > 0.01:
            final_weights[self.bond_etf] = bond_alloc
        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