"""
Multi-strategy portfolio blender.

Ensemble consists of:
  * Macro book: Tactical All Weather (dynamic to cash using momentum+trend regime)
  * Alpha book: Market-neutral Long/Short equity momentum (AQR 12-minus-1), long momentum
    strength filter, SPY regime throttle, and trailing vol targeting on the L/S return stream.
  * Optional sector sleeve: SPDR sector ETFs, monthly top-``k`` by 12-minus-1 momentum, equal weight.
  * Optional defensive MR sleeve: when SPY daily-return z-score is extremely negative, tilt the
    sleeve toward long bonds (TLT); exit when z normalizes — aims to add ballast in equity panics.
  * Optional **buy-the-dip** sleeve: long-only **RSI dip** entries (e.g. RSI(5) < 20) in an SMA200
    uptrend, ranked by normalized short-term ATR; each trade held a fixed number of trading days;
    remainder in cash at the risk-free rate.

The sleeves are combined in ``EnsembleManager`` (configurable overlay weights).
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict

import numpy as np
import pandas as pd

from RenTech.strategy_stack.portfolio_risk_manager import BASE_WEIGHTS, TacticalAllWeatherManager
from RenTech.strategy_stack.statarb_engine import rolling_zscore

# SPDR Sector ETFs — one per GICS sector (Yahoo tickers).
SPDR_SECTOR_ETFS: tuple[tuple[str, str], ...] = (
    ("XLE", "Energy"),
    ("XLF", "Financials"),
    ("XLV", "Health Care"),
    ("XLI", "Industrials"),
    ("XLB", "Materials"),
    ("XLK", "Technology"),
    ("XLY", "Consumer Discretionary"),
    ("XLP", "Consumer Staples"),
    ("XLU", "Utilities"),
    ("XLRE", "Real Estate"),
    ("XLC", "Communication Services"),
)
SPDR_SECTOR_TICKERS: list[str] = [t[0] for t in SPDR_SECTOR_ETFS]


def _resolve_spy_frame(spy_df: pd.DataFrame | None, equity_dict: Dict[str, pd.DataFrame]) -> pd.DataFrame:
    """Prefer explicit SPY panel; else use equity_dict['SPY'] if present."""
    if spy_df is not None and not spy_df.empty:
        out = spy_df.copy()
    elif "SPY" in equity_dict and not equity_dict["SPY"].empty:
        out = equity_dict["SPY"].copy()
    else:
        raise ValueError(
            "Regime throttle requires SPY: pass spy_df=... or include 'SPY' in equity_dict with columns "
            "close (or Close), ret (optional), sma_200 (optional)."
        )
    if "close" not in out.columns:
        if "Close" in out.columns:
            out["close"] = out["Close"].astype(np.float64)
        else:
            raise KeyError("SPY frame must include a 'close' or 'Close' price column.")
    return out


def _regime_exposure_series(spy: pd.DataFrame, index: pd.DatetimeIndex) -> pd.Series:
    """
    Daily regime exposure in (0, 1], executed next bar (caller shifts).

    Trend: SPY close > SMA200 -> factor 1.0; else 0.5.
    Vol: 20d SPY return vol > 1.5 × trailing 1y mean of that vol -> factor 0.5; else 1.0.

    Missing SMA / insufficient vol history: no penalty (factor 1.0) for that leg.
    """
    s = spy.copy()
    s.index = pd.to_datetime(s.index).tz_localize(None)
    s = s.sort_index()
    s = s.reindex(index).ffill()

    close = s["close"].astype(np.float64)
    if "sma_200" in s.columns:
        sma = s["sma_200"].astype(np.float64)
    else:
        sma = close.rolling(window=200, min_periods=200).mean()

    ret = s["ret"].astype(np.float64) if "ret" in s.columns else close.pct_change()

    market_vol_20 = ret.rolling(window=20, min_periods=20).std(ddof=1)
    vol_avg_1y = market_vol_20.rolling(window=252, min_periods=60).mean()

    ok_trend = close.notna().to_numpy() & sma.notna().to_numpy()
    trend_bull = ok_trend & (close > sma).to_numpy()
    trend_factor = np.where(ok_trend, np.where(trend_bull, 1.0, 0.5), 1.0)

    ok_vol = (
        market_vol_20.notna().to_numpy()
        & vol_avg_1y.notna().to_numpy()
        & (vol_avg_1y.to_numpy() > 1e-12)
    )
    vol_high = ok_vol & (market_vol_20.to_numpy() > 1.5 * vol_avg_1y.to_numpy())
    vol_factor = np.where(vol_high, 0.5, 1.0)

    g = trend_factor * vol_factor
    return pd.Series(g, index=index, name="regime_exposure", dtype=np.float64)


def _apply_ls_vol_targeting(
    ls_pre_vol: pd.Series,
    *,
    daily_rf: float,
    target_annual_vol: float = 0.10,
    max_scale: float = 2.0,
    smooth_window: int = 5,
    vol_window: int = 20,
) -> tuple[pd.Series, pd.Series]:
    """
    Scale pre-target L/S daily returns toward ``target_annual_vol`` using trailing realized vol.

    Returns (ls_daily_ret_scaled, scale_executed) where ``scale_executed`` is the lagged,
    smoothed, capped factor applied each day.
    """
    r = ls_pre_vol.astype(np.float64).fillna(0.0)
    realized_d = r.rolling(window=vol_window, min_periods=vol_window).std(ddof=1)
    realized_ann = realized_d * np.sqrt(252.0)

    eps = 1e-8
    rv = realized_ann
    raw = target_annual_vol / rv
    raw = raw.where(rv.notna() & (rv > eps), 1.0)
    raw = raw.replace([np.inf, -np.inf], 1.0)
    capped = raw.clip(upper=float(max_scale))

    sw = max(1, int(smooth_window))
    smoothed = capped.rolling(window=sw, min_periods=sw).mean()
    smoothed = smoothed.ffill().fillna(1.0)

    scale_exec = smoothed.shift(1).fillna(1.0).clip(upper=float(max_scale))
    rf = float(daily_rf)
    s = scale_exec.to_numpy(dtype=np.float64)
    rp = r.to_numpy(dtype=np.float64)
    cash_w = np.maximum(0.0, 1.0 - s)
    out = s * rp + cash_w * rf

    out_s = pd.Series(out, index=ls_pre_vol.index, dtype=np.float64)
    return out_s, scale_exec


def _collect_spell_lengths_days(w: np.ndarray, *, positive: bool, eps: float = 1e-9) -> list[int]:
    """Lengths of consecutive days with long (positive) or short (negative) weight."""
    if positive:
        active = w > eps
    else:
        active = w < -eps
    out: list[int] = []
    i = 0
    n = len(active)
    while i < n:
        if not active[i]:
            i += 1
            continue
        j = i
        while j < n and active[j]:
            j += 1
        out.append(j - i)
        i = j
    return out


def _populate_cross_sectional_ls_diagnostics(
    diagnostics: dict,
    *,
    weights_throttled: pd.DataFrame,
    ret_filled: pd.DataFrame,
    ls_daily_ret: pd.Series,
    long_mask: np.ndarray,
    short_mask: np.ndarray,
    top_n: int,
) -> None:
    """Fill ``diagnostics`` in place (cross-sectional L/S sleeve, post vol-targeting returns)."""
    long_contrib = (weights_throttled.clip(lower=0.0) * ret_filled).sum(axis=1).astype(np.float64)
    short_contrib = (weights_throttled.clip(upper=0.0) * ret_filled).sum(axis=1).astype(np.float64)

    r = ls_daily_ret.fillna(0.0).astype(np.float64).to_numpy(dtype=np.float64)
    pos = r[r > 0]
    neg = r[r < 0]
    diagnostics["engine"] = "cross_sectional_momentum"
    diagnostics["trading_days"] = int(len(r))
    diagnostics["daily_win_rate"] = float((r > 0).mean()) if len(r) else float("nan")
    diagnostics["daily_avg_win"] = float(pos.mean()) if pos.size else float("nan")
    diagnostics["daily_avg_loss"] = float(neg.mean()) if neg.size else float("nan")
    loss_sum = float(neg.sum()) if neg.size else float("nan")
    diagnostics["profit_factor"] = (
        float(pos.sum() / abs(loss_sum)) if pos.size and loss_sum < 0 and np.isfinite(loss_sum) else float("nan")
    )
    mu = float(r.mean()) if len(r) else float("nan")
    sd = float(r.std(ddof=1)) if len(r) > 1 else float("nan")
    diagnostics["sharpe_ratio_approx"] = mu / sd * np.sqrt(252.0) if sd > 1e-12 else float("nan")

    eq = np.cumprod(1.0 + r)
    peak = np.maximum.accumulate(eq)
    dd = eq / peak - 1.0
    diagnostics["max_drawdown"] = float(dd.min()) if len(dd) else float("nan")

    lg = long_contrib.to_numpy(dtype=np.float64)
    sg = short_contrib.to_numpy(dtype=np.float64)
    long_gross = weights_throttled.clip(lower=0.0).sum(axis=1).to_numpy(dtype=np.float64)
    short_gross = -weights_throttled.clip(upper=0.0).sum(axis=1).to_numpy(dtype=np.float64)
    lm = long_gross > 1e-9
    sm = short_gross > 1e-9
    diagnostics["long_leg_daily_win_rate"] = (
        float((lg[lm] > 0).mean()) if np.any(lm) else float("nan")
    )
    diagnostics["short_leg_daily_win_rate"] = (
        float((sg[sm] > 0).mean()) if np.any(sm) else float("nan")
    )
    diagnostics["days_with_long_exposure"] = int(lm.sum())
    diagnostics["days_with_short_exposure"] = int(sm.sum())

    M = int(long_mask.shape[0])
    if M >= 2:
        long_turns: list[int] = []
        short_turns: list[int] = []
        for m in range(1, M):
            lp = set(np.flatnonzero(long_mask[m - 1]).tolist())
            lc = set(np.flatnonzero(long_mask[m]).tolist())
            long_turns.append(len(lp ^ lc))
            sp = set(np.flatnonzero(short_mask[m - 1]).tolist())
            sc = set(np.flatnonzero(short_mask[m]).tolist())
            short_turns.append(len(sp ^ sc))
        diagnostics["monthly_rebalance_dates"] = int(M)
        diagnostics["avg_long_names_changed_per_rebalance"] = float(np.mean(long_turns))
        diagnostics["avg_short_names_changed_per_rebalance"] = float(np.mean(short_turns))
        diagnostics["total_long_slot_changes"] = int(sum(long_turns))
        diagnostics["total_short_slot_changes"] = int(sum(short_turns))
        diagnostics["approx_long_leg_trades"] = int(sum(long_turns) // 2)
        diagnostics["approx_short_leg_trades"] = int(sum(short_turns) // 2)
    else:
        diagnostics["monthly_rebalance_dates"] = int(M)
        for k in (
            "avg_long_names_changed_per_rebalance",
            "avg_short_names_changed_per_rebalance",
            "total_long_slot_changes",
            "total_short_slot_changes",
            "approx_long_leg_trades",
            "approx_short_leg_trades",
        ):
            diagnostics[k] = float("nan")

    diagnostics["top_n_per_leg"] = int(top_n)

    all_long_spells: list[int] = []
    all_short_spells: list[int] = []
    for col in weights_throttled.columns:
        w = weights_throttled[col].to_numpy(dtype=np.float64)
        all_long_spells.extend(_collect_spell_lengths_days(w, positive=True))
        all_short_spells.extend(_collect_spell_lengths_days(w, positive=False))
    diagnostics["avg_days_in_long_spell"] = (
        float(np.mean(all_long_spells)) if all_long_spells else float("nan")
    )
    diagnostics["avg_days_in_short_spell"] = (
        float(np.mean(all_short_spells)) if all_short_spells else float("nan")
    )
    diagnostics["num_long_spells_all_tickers"] = int(len(all_long_spells))
    diagnostics["num_short_spells_all_tickers"] = int(len(all_short_spells))
    diagnostics["rebalance_calendar"] = "BME"


def print_cross_sectional_ls_diagnostics(d: dict) -> None:
    """Pretty-print dict from :func:`_populate_cross_sectional_ls_diagnostics`."""
    if not d:
        return
    if d.get("engine") != "cross_sectional_momentum":
        return
    if d.get("note"):
        print(f"\n  --- L/S sleeve (AQR cross-sectional): {d['note']} ---")
        return
    print("\n  --- L/S sleeve (AQR cross-sectional) — trade & return stats ---")
    print(
        f"  Trading days: {d.get('trading_days', '—')}  |  "
        f"Sharpe (approx, ann.): {d.get('sharpe_ratio_approx', float('nan')):.3f}  |  "
        f"max drawdown: {d.get('max_drawdown', float('nan')) * 100:.2f}%"
    )
    print(
        f"  Daily win rate: {d.get('daily_win_rate', 0) * 100:.2f}%  |  "
        f"avg win day: {d.get('daily_avg_win', float('nan')) * 100:.4f}%  |  "
        f"avg loss day: {d.get('daily_avg_loss', float('nan')) * 100:.4f}%  |  "
        f"profit factor: {d.get('profit_factor', float('nan')):.3f}"
    )
    print(
        f"  Long leg (pre–vol-target PnL): win rate on days w/ longs: "
        f"{d.get('long_leg_daily_win_rate', float('nan')) * 100:.2f}%  "
        f"({d.get('days_with_long_exposure', 0)} days)"
    )
    print(
        f"  Short leg (pre–vol-target PnL): win rate on days w/ shorts: "
        f"{d.get('short_leg_daily_win_rate', float('nan')) * 100:.2f}%  "
        f"({d.get('days_with_short_exposure', 0)} days)"
    )
    print(
        f"  Monthly rebalance rows: {d.get('monthly_rebalance_dates', '—')} ({d.get('rebalance_calendar', 'BME')})  |  "
        f"top_n per leg: {d.get('top_n_per_leg', '—')}"
    )
    print(
        f"  Avg names changed/rebal — long: {d.get('avg_long_names_changed_per_rebalance', float('nan')):.2f}  |  "
        f"short: {d.get('avg_short_names_changed_per_rebalance', float('nan')):.2f}"
    )
    print(
        f"  Cumulative slot changes — long: {d.get('total_long_slot_changes', '—')}  |  "
        f"short: {d.get('total_short_slot_changes', '—')}  "
        f"(~½ = one-sided entries/exits per leg)"
    )
    print(
        f"  Approx. completed round-turns per leg — long: {d.get('approx_long_leg_trades', '—')}  |  "
        f"short: {d.get('approx_short_leg_trades', '—')}"
    )
    print(
        f"  Avg days per spell — long: {d.get('avg_days_in_long_spell', float('nan')):.1f}  |  "
        f"short: {d.get('avg_days_in_short_spell', float('nan')):.1f}  "
        f"(spells: {d.get('num_long_spells_all_tickers', '—')} long / {d.get('num_short_spells_all_tickers', '—')} short)"
    )
    print(
        "  Note: Leg win rates use pre–vol-targeting daily attribution; "
        "Sharpe/drawdown use final L/S daily returns after regime + vol scale."
    )


def _wilder_rsi(close: pd.Series, period: int) -> pd.Series:
    """RSI with Wilder smoothing (``period`` matches common RSI parameterization)."""
    p = int(period)
    if p < 1:
        raise ValueError("RSI period must be >= 1")
    delta = close.diff()
    gain = delta.clip(lower=0.0)
    loss = (-delta).clip(lower=0.0)
    avg_g = gain.ewm(alpha=1.0 / float(p), min_periods=p, adjust=False).mean()
    avg_l = loss.ewm(alpha=1.0 / float(p), min_periods=p, adjust=False).mean()
    rs = avg_g / avg_l.replace(0.0, np.nan)
    return (100.0 - (100.0 / (1.0 + rs))).astype(np.float64)


def _wilder_atr(high: pd.Series, low: pd.Series, close: pd.Series, period: int) -> pd.Series:
    """Average True Range with Wilder smoothing."""
    p = int(period)
    if p < 1:
        raise ValueError("ATR period must be >= 1")
    prev = close.shift(1)
    tr = pd.concat(
        [(high - low).abs(), (high - prev).abs(), (low - prev).abs()],
        axis=1,
    ).max(axis=1)
    return tr.ewm(alpha=1.0 / float(p), min_periods=p, adjust=False).mean().astype(np.float64)


def _high_low_for_atr(df: pd.DataFrame, close: pd.Series) -> tuple[pd.Series, pd.Series]:
    """Use OHLC when present; otherwise approximate range from close changes (ATR proxy)."""
    idx = close.index
    if "high" in df.columns:
        hi = df["high"].astype(np.float64).reindex(idx)
    elif "High" in df.columns:
        hi = df["High"].astype(np.float64).reindex(idx)
    else:
        hi = close
    if "low" in df.columns:
        lo = df["low"].astype(np.float64).reindex(idx)
    elif "Low" in df.columns:
        lo = df["Low"].astype(np.float64).reindex(idx)
    else:
        lo = close
    return hi.astype(np.float64), lo.astype(np.float64)


def _distribute_slots(n_slots: int, n_bins: int) -> list[int]:
    """Split ``n_slots`` into ``n_bins`` nonnegative integers that sum to ``n_slots`` (largest first)."""
    if n_bins <= 0:
        return []
    if n_slots <= 0:
        return [0] * n_bins
    base = n_slots // n_bins
    rem = n_slots % n_bins
    return [base + (1 if i < rem else 0) for i in range(n_bins)]


def _sector_neutral_ls_masks(
    aqr_arr: np.ndarray,
    eligible_long: np.ndarray,
    eligible_short: np.ndarray,
    sector_ids: np.ndarray,
    top_n: int,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Each rebalance month, split ``top_n`` long seats and ``top_n`` short seats across sectors
    that have eligible names on that leg; within a sector, rank by ``aqr_mom`` (long: high,
    short: low) with eligibility already encoded.
    """
    M, N = aqr_arr.shape
    long_mask = np.zeros((M, N), dtype=bool)
    short_mask = np.zeros((M, N), dtype=bool)
    for m in range(M):
        long_sectors = sorted({sector_ids[j] for j in range(N) if eligible_long[m, j]})
        if long_sectors:
            slots = _distribute_slots(top_n, len(long_sectors))
            for si, s in enumerate(long_sectors):
                ks = slots[si]
                if ks <= 0:
                    continue
                idxs = np.where((sector_ids == s) & eligible_long[m])[0]
                if idxs.size == 0:
                    continue
                take = min(int(ks), int(idxs.size))
                aq = aqr_arr[m, idxs].astype(np.float64)
                part = np.argpartition(-aq, take - 1)[:take]
                long_mask[m, idxs[part]] = True

        short_sectors = sorted({sector_ids[j] for j in range(N) if eligible_short[m, j]})
        if short_sectors:
            slots = _distribute_slots(top_n, len(short_sectors))
            for si, s in enumerate(short_sectors):
                ks = slots[si]
                if ks <= 0:
                    continue
                idxs = np.where((sector_ids == s) & eligible_short[m])[0]
                if idxs.size == 0:
                    continue
                take = min(int(ks), int(idxs.size))
                aq = aqr_arr[m, idxs].astype(np.float64)
                part = np.argpartition(aq, take - 1)[:take]
                short_mask[m, idxs[part]] = True
    return long_mask, short_mask


def _align_panel_frames(frames: list[pd.Series]) -> tuple[pd.DataFrame, list[str]]:
    """
    Concatenate a list of per-asset Series into a single DataFrame.

    Returns (panel_df, columns).
    """
    if not frames:
        raise ValueError("No frames provided")
    cols = [s.name for s in frames]
    panel = pd.concat(frames, axis=1)
    panel.columns = cols
    panel = panel.sort_index()
    return panel, cols


@dataclass
class CrossSectionalMomentum:
    """
    Market-neutral Long/Short momentum strategy (monthly rebalance, daily PnL).

    Uses AQR 12-minus-1 momentum (`aqr_mom`) to rank a universe, then forms:
      * Long leg = top_n names (highest aqr_mom)
      * Short leg = bottom_n names (lowest aqr_mom), smart-short filter

    With ``sector_map``, the same ``top_n`` long/short *names* are split across sectors
    (within-sector momentum rank). That cuts sector crash risk but usually **hurts raw momentum
    PnL**: the factor often earns from concentrating in the strongest names wherever they sit,
    and sector quotas add lower-momentum longs and less extreme shorts.

    Optional SPY-based regime throttle scales gross exposure (trend + vol penalties)
    and parks the remainder in cash at ``cash_annual_yield``.

    Longs require ``aqr_mom > long_aqr_momentum_min`` (default 5%) to avoid weak/sideways names.

    Post-PnL volatility targeting: 20d realized vol of the L/S return stream, scale toward
    ``target_annual_vol`` (cap + smoothed scale, shifted 1 bar; cash on under-allocated sleeve).

    Lookahead avoidance:
      * monthly weights forward-filled, shifted 1 bar
      * regime exposure shifted 1 bar after SPY filters
      * vol scale shifted 1 bar after smoothing
    """

    def generate_ls_returns(
        self,
        equity_dict: Dict[str, pd.DataFrame],
        top_n: int = 10,
        *,
        spy_df: pd.DataFrame | None = None,
        cash_annual_yield: float = 0.04,
        target_annual_vol: float = 0.10,
        vol_scale_cap: float = 2.0,
        vol_scale_smooth_days: int = 5,
        long_aqr_momentum_min: float = 0.05,
        sector_map: Dict[str, str] | None = None,
        verbose: bool = True,
        diagnostics: dict | None = None,
    ) -> pd.Series:
        """
        Parameters
        ----------
        equity_dict
            ticker -> DataFrame containing columns:
              * aqr_mom, ret, close, sma_200
        top_n
            Number of names on each leg (long + short).
        spy_df
            SPY daily panel (``close``, optional ``ret``, ``sma_200``). If omitted, uses
            ``equity_dict['SPY']`` when present.
        cash_annual_yield
            Annual yield on the capital not deployed to the L/S book when exposure < 1,
            and on the sleeve not used after vol scaling when scale < 1.
        target_annual_vol
            Annualized vol target for the L/S return stream (after regime throttle).
        vol_scale_cap
            Maximum vol-scaling leverage (e.g. 2.0).
        vol_scale_smooth_days
            Moving-average window on the capped scaling factor.
        long_aqr_momentum_min
            Minimum AQR momentum (fraction) required to be eligible for the long leg.
        sector_map
            If set, allocate long/short seats per sector (even split of ``top_n`` per leg);
            tickers missing from the map use sector "Unknown".
        verbose
            If True, print diagnostics (regime gross exposure, vol targeting).
        diagnostics
            If a dict is passed, it is cleared and filled with L/S trade and return
            statistics (see :func:`print_cross_sectional_ls_diagnostics`).

        Returns
        -------
        pd.Series
            Daily Long/Short returns indexed by date.
        """
        if top_n <= 0:
            raise ValueError("top_n must be > 0")
        if not equity_dict:
            raise ValueError("equity_dict cannot be empty")

        tickers = list(equity_dict.keys())
        # Build aligned daily panels. We only loop over tickers (not over dates).
        aqr_pan: list[pd.Series] = []
        ret_pan: list[pd.Series] = []
        close_pan: list[pd.Series] = []
        sma_pan: list[pd.Series] = []

        for t in tickers:
            df = equity_dict[t]
            required = {"close", "ret", "sma_200", "aqr_mom"}
            missing = required.difference(df.columns)
            if missing:
                raise KeyError(f"{t} missing required columns: {sorted(missing)}")

            idx = pd.to_datetime(df.index).tz_localize(None)
            df2 = df.copy()
            df2.index = idx
            df2 = df2.sort_index()

            aqr_pan.append(df2["aqr_mom"].astype(np.float64).rename(t))
            ret_pan.append(df2["ret"].astype(np.float64).rename(t))
            close_pan.append(df2["close"].astype(np.float64).rename(t))
            sma_pan.append(df2["sma_200"].astype(np.float64).rename(t))

        aqr_df, _ = _align_panel_frames(aqr_pan)
        ret_df, _ = _align_panel_frames(ret_pan)
        close_df, _ = _align_panel_frames(close_pan)
        sma_df, _ = _align_panel_frames(sma_pan)

        # Common daily timeline across tickers.
        master_index = aqr_df.index.union(ret_df.index).union(close_df.index).union(sma_df.index).sort_values()
        aqr_df = aqr_df.reindex(master_index)
        ret_df = ret_df.reindex(master_index)
        close_df = close_df.reindex(master_index)
        sma_df = sma_df.reindex(master_index)

        # Forward-fill "state-like" levels so month-end feature values are stable
        # (NaNs due to IPO / holidays should not wipe entire panels).
        close_df = close_df.ffill()
        sma_df = sma_df.ffill()
        aqr_df = aqr_df.ffill()

        # Volatility: 20-day rolling std of daily returns.
        # Keep NaNs where insufficient history; later we turn invalid weights into 0.
        vol_20_df = ret_df.rolling(window=20, min_periods=20).std(ddof=1)

        # --- Monthly rebalance: sample features at business month-end (last trading day). ---
        bm_index = aqr_df.resample("BME").last().index
        aqr_m = aqr_df.resample("BME").last()
        close_m = close_df.resample("BME").last()
        sma_m = sma_df.resample("BME").last()
        vol_m = vol_20_df.resample("BME").last()

        aqr_arr = aqr_m.to_numpy(dtype=np.float64)  # [M, N]
        close_arr = close_m.to_numpy(dtype=np.float64)
        sma_arr = sma_m.to_numpy(dtype=np.float64)
        vol_arr = vol_m.to_numpy(dtype=np.float64)
        valid_a = np.isfinite(aqr_arr)

        M, N = aqr_arr.shape
        if M < 2 or N < 2:
            # Not enough data to rebalance meaningfully.
            # Execute weights would still be shifted, but we return zeros to keep the pipeline stable.
            out_z = pd.Series(0.0, index=master_index, name="ls_ret", dtype=np.float64)
            if diagnostics is not None:
                diagnostics.clear()
                diagnostics["engine"] = "cross_sectional_momentum"
                diagnostics["note"] = "insufficient_monthly_rows_or_tickers"
            return out_z

        if top_n > N:
            raise ValueError(f"top_n={top_n} cannot exceed universe size N={N}")

        eligible_long = valid_a & (aqr_arr > float(long_aqr_momentum_min))
        eligible_short = (close_arr < sma_arr) & valid_a

        rows = np.arange(M)[:, None]
        if sector_map is not None:
            sector_ids = np.array([str(sector_map.get(str(t), "Unknown")) for t in aqr_m.columns], dtype=object)
            long_mask, short_mask = _sector_neutral_ls_masks(
                aqr_arr, eligible_long, eligible_short, sector_ids, top_n
            )
        else:
            # --- Long leg: top_n by aqr_mom among names with momentum > long_aqr_momentum_min. ---
            aqr_for_long = np.where(eligible_long, aqr_arr, -np.inf)
            idx_long = np.argpartition(-aqr_for_long, kth=top_n - 1, axis=1)[:, :top_n]  # [M, top_n]
            long_mask = np.zeros((M, N), dtype=bool)
            long_mask[rows, idx_long] = True
            long_mask &= eligible_long

            # --- Short leg: bottom_n by aqr_mom, smart-short filter close < sma_200. ---
            aqr_for_short = np.where(eligible_short, aqr_arr, np.inf)
            idx_short = np.argpartition(aqr_for_short, kth=top_n - 1, axis=1)[:, :top_n]
            short_mask = np.zeros((M, N), dtype=bool)
            short_mask[rows, idx_short] = True
            short_mask &= eligible_short

        # Count shorts per month (for scaling behavior).
        count_short = short_mask.sum(axis=1).astype(np.float64)  # [M]

        # --- Inverse-volatility raw weights at month-end. ---
        # Convert vol to inverse vol. Invalid vol => weight 0.
        inv_vol = np.where(np.isfinite(vol_arr) & (vol_arr > 0.0), 1.0 / vol_arr, 0.0)

        # Long: raw weights only on selected longs.
        long_raw = inv_vol * long_mask  # [M,N]
        long_raw_sum = long_raw.sum(axis=1)  # [M]
        long_w = np.divide(
            long_raw,
            long_raw_sum[:, None],
            out=np.zeros_like(long_raw, dtype=np.float64),
            where=long_raw_sum[:, None] > 0.0,
        )
        # Normalize longs to +1.0 (when any long has valid vol, this sums to 1).

        # Short: normalize magnitude to -1.0 if >= top_n eligible shorts; else scale down by count_short/top_n.
        short_raw = inv_vol * short_mask  # positive magnitudes
        short_raw_sum = short_raw.sum(axis=1)  # [M]
        # Scale down if fewer than top_n eligible shorts; yields sum in [-1,0].
        scale_factor = np.minimum(1.0, count_short / float(top_n))  # [M]
        short_w = np.divide(
            short_raw,
            short_raw_sum[:, None],
            out=np.zeros_like(short_raw, dtype=np.float64),
            where=short_raw_sum[:, None] > 0.0,
        )
        short_w = -short_w * scale_factor[:, None]

        # Combine into signed target weights on month-end.
        weights_m = long_w + short_w  # [M,N]
        weights_m_df = pd.DataFrame(weights_m, index=bm_index, columns=aqr_m.columns)

        # --- Forward-fill monthly weights to daily timeline + execute next bar ---
        weights_d = weights_m_df.reindex(master_index).ffill()
        weights_d = weights_d.shift(1).fillna(0.0)

        # Apply to daily returns; if ret is NaN, treat as 0 for PnL.
        ret_filled = ret_df.fillna(0.0).astype(np.float64)

        # --- Regime-based exposure throttle (SPY trend + realized vol vs 1y average) ---
        spy_panel = _resolve_spy_frame(spy_df, equity_dict)
        regime_raw = _regime_exposure_series(spy_panel, master_index)
        regime_exec = regime_raw.shift(1).fillna(1.0).clip(0.0, 1.0)

        daily_rf = float(cash_annual_yield) / 252.0
        weights_throttled = weights_d.mul(regime_exec, axis=0)
        ls_core = (weights_throttled.to_numpy(dtype=np.float64) * ret_filled.to_numpy(dtype=np.float64)).sum(axis=1)
        cash_carry = (1.0 - regime_exec.to_numpy(dtype=np.float64)) * daily_rf
        ls_pre_vol = pd.Series(ls_core + cash_carry, index=master_index, dtype=np.float64)

        ls_daily_ret, scale_exec = _apply_ls_vol_targeting(
            ls_pre_vol,
            daily_rf=daily_rf,
            target_annual_vol=float(target_annual_vol),
            max_scale=float(vol_scale_cap),
            smooth_window=int(vol_scale_smooth_days),
            vol_window=20,
        )

        if verbose:
            gross = weights_throttled.abs().sum(axis=1).astype(np.float64)
            avg_gross_pct = float(gross.mean()) * 100.0 if len(gross) else float("nan")
            avg_g = float(regime_exec.mean()) if len(regime_exec) else float("nan")
            print(
                f"  L/S regime throttle: avg gross exposure = {avg_gross_pct:.2f}% "
                f"(full-throttle reference ≈ 200%; avg regime factor ≈ {avg_g:.3f})"
            )
            avg_scale = float(scale_exec.mean()) if len(scale_exec) else float("nan")
            print(
                f"  L/S vol targeting: target_ann_vol={float(target_annual_vol):.2f} "
                f"cap={float(vol_scale_cap):.2f} smooth={int(vol_scale_smooth_days)}d "
                f"| avg executed scale ≈ {avg_scale:.3f}"
            )

        if diagnostics is not None:
            diagnostics.clear()
            _populate_cross_sectional_ls_diagnostics(
                diagnostics,
                weights_throttled=weights_throttled,
                ret_filled=ret_filled,
                ls_daily_ret=ls_daily_ret,
                long_mask=long_mask,
                short_mask=short_mask,
                top_n=int(top_n),
            )

        return ls_daily_ret.rename("ls_ret")


@dataclass
class BuyTheDipSleeve:
    """
    **Long-only buy-the-dip** sleeve (daily signals, fixed holding period, daily PnL).

    Classic rule set (configurable): **close > SMA(200)** (long-term uptrend) and washout signal
    (RSI or pct-drop). Among new signals each day, keep up to ``top_n`` names ranked by score.
    Each entry is held exactly ``hold_trading_days`` sessions; PnL uses close-to-close returns
    (signal at close *t*, first return day *t+1*). Optional ``max_concurrent`` caps **total**
    names held at once (skips new entries when the book is full).

    At most **one active trade per ticker**; overlapping names are weighted **equal** or
    **inverse trailing vol** among names active that day. Optional **SPY > SMA200** gate applies
    only when **opening** new positions.

    Same SPY regime throttle and trailing vol targeting on the sleeve return stream as other equity
    sleeves. Universe membership is **as loaded** (e.g. current S&P list), not point-in-time
    constituents.
    """

    rsi_period: int = 5
    rsi_max: float = 20.0
    atr_period: int = 5
    hold_trading_days: int = 5
    vol_window: int = 10
    sma_trend_window: int = 200
    weighting: str = "equal"
    rank_by: str = "atr_norm"
    dip_in_uptrend: bool = True
    only_when_spy_bull: bool = False
    signal_mode: str = "rsi"
    pct_drop_min: float = 0.03
    min_atr_norm_pct: float = 3.0
    # Require same-day return to underperform SPY by at least this fraction (0 = off).
    relative_spy_min: float = 0.0
    # Subtract SPY return × gross long exposure (market-neutral excess-return sleeve).
    hedge_spy: bool = False
    # Cap total names held at once (None = unlimited; overlapping holds can exceed top_n).
    max_concurrent: int | None = None
    # ``legacy`` = close/next-bar entry, fixed hold. ``cracking_markets`` = limit entry +
    # multi-rule exits (CrackingMarkets / RealTest short mean-reversion spec).
    execution_style: str = "legacy"
    limit_atr_mult: float = 0.9
    profit_atr_mult: float = 0.5
    exit_on_prior_high: bool = True
    # When True with ``cracking_markets``, skip SPY regime throttle and vol targeting.
    cracking_markets_pure: bool = False

    def _generate_returns_cracking_markets(
        self,
        equity_dict: Dict[str, pd.DataFrame],
        top_n: int,
        *,
        spy_df: pd.DataFrame | None,
        master_index: pd.DatetimeIndex,
        close_df: pd.DataFrame,
        high_df: pd.DataFrame,
        low_df: pd.DataFrame,
        atr_df: pd.DataFrame,
        atrn_df: pd.DataFrame,
        sma_df: pd.DataFrame,
        ret_chg_df: pd.DataFrame,
        dip_sig: np.ndarray,
        hold: int,
        max_conc: int | None,
        wg: str,
        vol_arr: np.ndarray,
        spy_panel: pd.DataFrame,
        spy_ret_arr: np.ndarray,
        cash_annual_yield: float,
        target_annual_vol: float,
        vol_scale_cap: float,
        vol_scale_smooth_days: int,
        verbose: bool,
        tw: int,
        min_atr_pct: float,
        pct_drop: float,
    ) -> pd.Series:
        """
        CrackingMarkets article rules:
          * Signal: −3% day, SMA200 up, ATR%/close > 3%, rank ATR/close
          * Entry: next session limit at signal_close − limit_atr_mult × ATR(5)
          * Exit (first hit): hold > N days | close > prior high | close ≥ fill + profit_atr_mult × ATR
        """
        close_arr = close_df.to_numpy(dtype=np.float64)
        high_arr = high_df.to_numpy(dtype=np.float64)
        low_arr = low_df.to_numpy(dtype=np.float64)
        atr_arr = atr_df.to_numpy(dtype=np.float64)
        D, N = close_arr.shape
        lim_mult = float(self.limit_atr_mult)
        prof_mult = float(self.profit_atr_mult)

        pos_fill_day = np.full(N, -1, dtype=np.int32)
        fill_px = np.full(N, np.nan, dtype=np.float64)
        entry_atr = np.full(N, np.nan, dtype=np.float64)
        pending = np.zeros(N, dtype=bool)
        pending_limit = np.full(N, np.nan, dtype=np.float64)
        pending_atr = np.full(N, np.nan, dtype=np.float64)
        pending_from = np.full(N, -1, dtype=np.int32)

        port_ret = np.zeros(D, dtype=np.float64)
        peak_concurrent = 0
        peak_gross = 0.0

        for d in range(D):
            # Drop unfilled limits from older signal days.
            stale = pending & (pending_from < d - 1)
            pending[stale] = False

            # Limit fills for signals from prior session.
            if d >= 1:
                fill_today = pending & (pending_from == d - 1)
                for i in np.flatnonzero(fill_today):
                    lim = pending_limit[i]
                    if np.isfinite(lim) and np.isfinite(low_arr[d, i]) and low_arr[d, i] <= lim:
                        pos_fill_day[i] = d
                        fill_px[i] = lim
                        entry_atr[i] = pending_atr[i]
                    pending[i] = False

            held = pos_fill_day >= 0
            name_ret = np.zeros(N, dtype=np.float64)
            if np.any(held):
                for i in np.flatnonzero(held):
                    fd = int(pos_fill_day[i])
                    if fd > d:
                        continue
                    c_d = close_arr[d, i]
                    if fd == d:
                        fp = fill_px[i]
                        if fp > 0 and np.isfinite(c_d):
                            name_ret[i] = c_d / fp - 1.0
                    elif d > 0:
                        c_prev = close_arr[d - 1, i]
                        if c_prev > 0 and np.isfinite(c_d):
                            name_ret[i] = c_d / c_prev - 1.0

                idxs = np.flatnonzero(held & (pos_fill_day <= d))
                if idxs.size > 0:
                    peak_concurrent = max(peak_concurrent, int(idxs.size))
                    if wg == "equal":
                        w = np.full(idxs.size, 1.0 / float(idxs.size), dtype=np.float64)
                    else:
                        v = vol_arr[d, idxs]
                        inv = np.where(np.isfinite(v) & (v > 0.0), 1.0 / v, 0.0)
                        s = float(inv.sum())
                        w = inv / s if s > 0 else np.full(idxs.size, 1.0 / float(idxs.size))
                    port_ret[d] = float(np.dot(w, name_ret[idxs]))
                    peak_gross = max(peak_gross, float(w.sum()))

            # Exits at close(d).
            for i in np.flatnonzero(pos_fill_day >= 0):
                fd = int(pos_fill_day[i])
                if fd > d:
                    continue
                days_held = d - fd
                exit_now = days_held > hold
                if self.exit_on_prior_high and d > 0:
                    hi_prev = high_arr[d - 1, i]
                    if np.isfinite(hi_prev) and close_arr[d, i] > hi_prev:
                        exit_now = True
                fp, ea = fill_px[i], entry_atr[i]
                if np.isfinite(fp) and np.isfinite(ea):
                    if close_arr[d, i] >= fp + prof_mult * ea:
                        exit_now = True
                if exit_now:
                    pos_fill_day[i] = -1

            # Queue new limit orders for next session.
            flat = (pos_fill_day < 0) & (~pending)
            n_open = int(np.sum(pos_fill_day >= 0)) + int(np.sum(pending))
            slots = int(top_n)
            if max_conc is not None:
                slots = min(slots, max(0, max_conc - n_open))
            if slots > 0 and d + 1 < D:
                cand = np.flatnonzero(flat & dip_sig[d])
                if cand.size > 0:
                    sc = atrn_df.to_numpy(dtype=np.float64)[d, cand]
                    key = np.where(np.isfinite(sc), sc, -np.inf)
                    order = np.argsort(-key, kind="stable")
                    take = cand[order[: min(slots, cand.size)]]
                    for i in take:
                        atr_v = atr_arr[d, i]
                        cl = close_arr[d, i]
                        if not (np.isfinite(atr_v) and np.isfinite(cl)):
                            continue
                        pending[i] = True
                        pending_from[i] = d
                        pending_limit[i] = cl - lim_mult * atr_v
                        pending_atr[i] = atr_v

        ls_pre_vol = pd.Series(port_ret, index=master_index, dtype=np.float64)
        if bool(self.cracking_markets_pure):
            if verbose:
                avg_gross_pct = float(np.nanmean(port_ret)) * 100.0  # not quite gross; use peak
                print(
                    f"  CrackingMarkets dip (pct_drop>={pct_drop:.1%}, ATR%>{min_atr_pct:g}, "
                    f"limit={lim_mult:.1f}×ATR, exits:>{hold}d|prior-high|+{prof_mult:.1f}ATR, "
                    f"rank ATR/px): peak concurrent={peak_concurrent} | "
                    f"peak gross≈{peak_gross * 100:.1f}%"
                )
            return ls_pre_vol.rename("dip_ret")

        regime_raw = _regime_exposure_series(spy_panel, master_index)
        regime_exec = regime_raw.shift(1).fillna(1.0).clip(0.0, 1.0)
        daily_rf = float(cash_annual_yield) / 252.0
        gross_proxy = pd.Series(port_ret, index=master_index).abs()  # placeholder
        reg = regime_exec.to_numpy(dtype=np.float64)
        core = port_ret * reg
        cash_carry = (1.0 - reg) * daily_rf
        ls_pre_vol = pd.Series(core + cash_carry, index=master_index, dtype=np.float64)
        dip_daily_ret, scale_exec = _apply_ls_vol_targeting(
            ls_pre_vol,
            daily_rf=daily_rf,
            target_annual_vol=float(target_annual_vol),
            max_scale=float(vol_scale_cap),
            smooth_window=int(vol_scale_smooth_days),
            vol_window=20,
        )
        if verbose:
            print(
                f"  CrackingMarkets dip (+regime/vol): peak concurrent={peak_concurrent} | "
                f"avg vol scale ≈ {float(scale_exec.mean()) if len(scale_exec) else float('nan'):.3f}"
            )
        return dip_daily_ret.rename("dip_ret")

    def generate_returns(
        self,
        equity_dict: Dict[str, pd.DataFrame],
        top_n: int = 10,
        *,
        spy_df: pd.DataFrame | None = None,
        cash_annual_yield: float = 0.04,
        target_annual_vol: float = 0.10,
        vol_scale_cap: float = 2.0,
        vol_scale_smooth_days: int = 5,
        verbose: bool = True,
    ) -> pd.Series:
        if top_n <= 0:
            raise ValueError("top_n must be > 0")
        if not equity_dict:
            raise ValueError("equity_dict cannot be empty")

        rsi_p = int(self.rsi_period)
        atr_p = int(self.atr_period)
        hold = int(self.hold_trading_days)
        tw = int(self.sma_trend_window)
        rsi_thr = float(self.rsi_max)
        if rsi_p < 1 or atr_p < 1 or hold < 1 or tw < 1:
            raise ValueError("rsi_period, atr_period, hold_trading_days, sma_trend_window must be >= 1")

        rk = str(self.rank_by).lower().strip()
        if rk not in ("atr_norm", "rsi", "rel_underperf"):
            raise ValueError("rank_by must be 'atr_norm', 'rsi', or 'rel_underperf'")
        rel_spy_min = float(self.relative_spy_min)
        if rel_spy_min < 0.0:
            raise ValueError("relative_spy_min must be >= 0")

        sig_mode = str(self.signal_mode).lower().strip()
        if sig_mode not in ("rsi", "pct_drop"):
            raise ValueError("signal_mode must be 'rsi' or 'pct_drop'")
        pct_drop = float(self.pct_drop_min)
        min_atr_pct = float(self.min_atr_norm_pct)
        if sig_mode == "pct_drop" and not (0.0 < pct_drop < 1.0):
            raise ValueError("pct_drop_min must be in (0, 1) for pct_drop signal_mode")

        tickers = list(equity_dict.keys())
        ret_pan: list[pd.Series] = []
        close_pan: list[pd.Series] = []
        high_pan: list[pd.Series] = []
        low_pan: list[pd.Series] = []
        atr_pan: list[pd.Series] = []
        rsi_pan: list[pd.Series] = []
        atrn_pan: list[pd.Series] = []
        sma_pan: list[pd.Series] = []

        req = {"close", "ret"}
        for t in tickers:
            df = equity_dict[t]
            missing = req.difference(df.columns)
            if missing:
                raise KeyError(f"{t} missing required columns for buy-the-dip: {sorted(missing)}")

            idx = pd.to_datetime(df.index).tz_localize(None)
            df2 = df.copy()
            df2.index = idx
            df2 = df2.sort_index()

            close_s = df2["close"].astype(np.float64)
            hi, lo = _high_low_for_atr(df2, close_s)
            # True range needs OHLC; if High/Low missing, both collapse to close → TR = |Δclose|.
            atr_s = _wilder_atr(hi, lo, close_s, atr_p)
            atrn_s = (atr_s / close_s.replace(0.0, np.nan)).astype(np.float64)
            rsi_s = _wilder_rsi(close_s, rsi_p)
            sma_s = close_s.rolling(window=tw, min_periods=tw).mean()

            ret_pan.append(df2["ret"].astype(np.float64).rename(t))
            close_pan.append(close_s.rename(t))
            high_pan.append(hi.rename(t))
            low_pan.append(lo.rename(t))
            atr_pan.append(atr_s.rename(t))
            rsi_pan.append(rsi_s.rename(t))
            atrn_pan.append(atrn_s.rename(t))
            sma_pan.append(sma_s.rename(t))

        ret_df, _ = _align_panel_frames(ret_pan)
        close_df, _ = _align_panel_frames(close_pan)
        high_df, _ = _align_panel_frames(high_pan)
        low_df, _ = _align_panel_frames(low_pan)
        atr_df, _ = _align_panel_frames(atr_pan)
        rsi_df, _ = _align_panel_frames(rsi_pan)
        atrn_df, _ = _align_panel_frames(atrn_pan)
        sma_df, _ = _align_panel_frames(sma_pan)

        master_index = ret_df.index.union(close_df.index).sort_values()
        ret_df = ret_df.reindex(master_index)
        close_df = close_df.reindex(master_index).ffill()
        high_df = high_df.reindex(master_index).ffill()
        low_df = low_df.reindex(master_index).ffill()
        atr_df = atr_df.reindex(master_index)
        rsi_df = rsi_df.reindex(master_index)
        atrn_df = atrn_df.reindex(master_index)
        sma_df = sma_df.reindex(master_index)
        ret_chg_df = close_df.pct_change()

        vw = int(self.vol_window)
        vol_df = ret_df.rolling(window=vw, min_periods=vw).std(ddof=1)

        D, N = len(master_index), len(close_df.columns)
        if D < 2 or N < 2:
            return pd.Series(0.0, index=master_index, name="ls_short_ret", dtype=np.float64)

        if top_n > N:
            raise ValueError(f"top_n={top_n} cannot exceed universe size N={N}")
        max_conc = self.max_concurrent
        if max_conc is not None:
            max_conc = int(max_conc)
            if max_conc < 1:
                raise ValueError("max_concurrent must be >= 1 when set")

        spy_panel = _resolve_spy_frame(spy_df, equity_dict)
        spy_close = spy_panel["close"].astype(np.float64)
        spy_close.index = pd.to_datetime(spy_close.index).tz_localize(None)
        spy_close = spy_close.reindex(master_index).ffill()
        if "sma_200" in spy_panel.columns:
            spy_sma = spy_panel["sma_200"].astype(np.float64)
            spy_sma.index = pd.to_datetime(spy_sma.index).tz_localize(None)
            spy_sma = spy_sma.reindex(master_index).ffill()
        else:
            spy_sma = spy_close.rolling(window=200, min_periods=200).mean()
        spy_bull_d = (spy_close > spy_sma).fillna(False).to_numpy(dtype=bool)
        spy_ret_s = spy_close.pct_change().reindex(master_index).fillna(0.0)
        spy_ret_arr = spy_ret_s.to_numpy(dtype=np.float64)

        close_arr = close_df.to_numpy(dtype=np.float64)
        rsi_arr = rsi_df.to_numpy(dtype=np.float64)
        sma_arr = sma_df.to_numpy(dtype=np.float64)
        atrn_arr = atrn_df.to_numpy(dtype=np.float64)
        ret_chg_arr = ret_chg_df.to_numpy(dtype=np.float64)
        vol_arr = vol_df.to_numpy(dtype=np.float64)

        base_ok = np.isfinite(close_arr) & np.isfinite(atrn_arr)
        if sig_mode == "pct_drop":
            base_ok &= np.isfinite(ret_chg_arr)
            dip_sig = base_ok & (ret_chg_arr <= -pct_drop)
            dip_sig &= (atrn_arr * 100.0) > min_atr_pct
            if rel_spy_min > 0.0:
                spy_row = spy_ret_arr[:, np.newaxis]
                rel_vs_spy = ret_chg_arr - spy_row
                dip_sig &= np.isfinite(rel_vs_spy) & (rel_vs_spy <= -rel_spy_min)
        else:
            base_ok &= np.isfinite(rsi_arr)
            dip_sig = base_ok & (rsi_arr < rsi_thr)
        if bool(self.dip_in_uptrend):
            dip_sig &= np.isfinite(sma_arr) & (close_arr > sma_arr)

        exec_style = str(self.execution_style).lower().strip()
        if exec_style == "cracking_markets":
            wg = str(self.weighting).lower()
            if wg not in ("equal", "inv_vol"):
                raise ValueError(f"weighting must be 'equal' or 'inv_vol', got {self.weighting!r}")
            return self._generate_returns_cracking_markets(
                equity_dict,
                top_n,
                spy_df=spy_df,
                master_index=master_index,
                close_df=close_df,
                high_df=high_df,
                low_df=low_df,
                atr_df=atr_df,
                atrn_df=atrn_df,
                sma_df=sma_df,
                ret_chg_df=ret_chg_df,
                dip_sig=dip_sig,
                hold=hold,
                max_conc=max_conc,
                wg=wg,
                vol_arr=vol_arr,
                spy_panel=spy_panel,
                spy_ret_arr=spy_ret_arr,
                cash_annual_yield=cash_annual_yield,
                target_annual_vol=target_annual_vol,
                vol_scale_cap=vol_scale_cap,
                vol_scale_smooth_days=vol_scale_smooth_days,
                verbose=verbose,
                tw=tw,
                min_atr_pct=min_atr_pct,
                pct_drop=pct_drop,
            )

        active_entry = np.full(N, -1, dtype=np.int32)
        weights_mat = np.zeros((D, N), dtype=np.float64)

        wg = str(self.weighting).lower()
        if wg not in ("equal", "inv_vol"):
            raise ValueError(f"weighting must be 'equal' or 'inv_vol', got {self.weighting!r}")

        spy_gate = bool(self.only_when_spy_bull)
        peak_concurrent = 0

        for d in range(D):
            expired = (active_entry >= 0) & (d > active_entry + hold)
            active_entry[expired] = -1

            in_trade = (active_entry >= 0) & (active_entry < d) & (d <= active_entry + hold)
            idxs = np.flatnonzero(in_trade)
            if idxs.size > 0:
                peak_concurrent = max(peak_concurrent, int(idxs.size))
                if wg == "equal":
                    w_each = 1.0 / float(idxs.size)
                    weights_mat[d, idxs] = w_each
                else:
                    v = vol_arr[d, idxs]
                    inv = np.where(np.isfinite(v) & (v > 0.0), 1.0 / v, 0.0)
                    s = float(inv.sum())
                    if s > 0.0:
                        weights_mat[d, idxs] = inv / s

            flat = active_entry < 0
            open_ok = np.ones(D, dtype=bool)
            if spy_gate:
                open_ok = spy_bull_d
            if open_ok[d]:
                cand = np.flatnonzero(flat & dip_sig[d])
                if cand.size > 0:
                    if rk == "atr_norm":
                        sc = atrn_arr[d, cand].astype(np.float64)
                        # Descending ATR/px; non-finite scores last (treat as -inf before negation).
                        key = np.where(np.isfinite(sc), sc, -np.inf)
                        order = np.argsort(-key, kind="stable")
                    elif rk == "rel_underperf":
                        rel = ret_chg_arr[d, cand] - spy_ret_arr[d]
                        sc = (-rel).astype(np.float64)  # larger = more underperformance vs SPY
                        key = np.where(np.isfinite(sc), sc, -np.inf)
                        order = np.argsort(-key, kind="stable")
                    else:
                        sc = rsi_arr[d, cand].astype(np.float64)
                        # Ascending RSI (deeper dip first); non-finite last.
                        key = np.where(np.isfinite(sc), sc, np.inf)
                        order = np.argsort(key, kind="stable")
                    slots = int(top_n)
                    if max_conc is not None:
                        n_open = int(idxs.size) if idxs.size > 0 else 0
                        slots = min(slots, max(0, max_conc - n_open))
                    take_n = min(slots, int(cand.size))
                    if take_n > 0:
                        chosen = cand[order[:take_n]]
                        active_entry[chosen] = d

        weights_d = pd.DataFrame(weights_mat, index=master_index, columns=close_df.columns)
        ret_filled = ret_df.fillna(0.0).astype(np.float64)

        regime_raw = _regime_exposure_series(spy_panel, master_index)
        regime_exec = regime_raw.shift(1).fillna(1.0).clip(0.0, 1.0)

        daily_rf = float(cash_annual_yield) / 252.0
        w_arr = weights_d.to_numpy(dtype=np.float64)
        reg = regime_exec.to_numpy(dtype=np.float64)
        w_th = w_arr * reg[:, np.newaxis]
        gross = w_th.sum(axis=1)
        core = (w_th * ret_filled.to_numpy(dtype=np.float64)).sum(axis=1)
        if bool(self.hedge_spy):
            core = core - spy_ret_arr * gross
        cash_carry = (1.0 - reg) * daily_rf
        ls_pre_vol = pd.Series(core + cash_carry, index=master_index, dtype=np.float64)

        dip_daily_ret, scale_exec = _apply_ls_vol_targeting(
            ls_pre_vol,
            daily_rf=daily_rf,
            target_annual_vol=float(target_annual_vol),
            max_scale=float(vol_scale_cap),
            smooth_window=int(vol_scale_smooth_days),
            vol_window=20,
        )

        if verbose:
            avg_gross_pct = float(np.nanmean(w_th.sum(axis=1))) * 100.0 if w_th.size else float("nan")
            avg_g = float(regime_exec.mean()) if len(regime_exec) else float("nan")
            up = f"SMA{tw} uptrend" if bool(self.dip_in_uptrend) else "no SMA filter"
            sp = "SPY>SMA200@entry" if spy_gate else "any SPY"
            if rk == "atr_norm":
                rk_msg = "rank ATR/px"
            elif rk == "rel_underperf":
                rk_msg = "rank rel underperf vs SPY"
            else:
                rk_msg = "rank RSI"
            if sig_mode == "pct_drop":
                sig_msg = f"pct_drop>={pct_drop:.1%}, ATR%>{min_atr_pct:g}"
                if rel_spy_min > 0.0:
                    sig_msg += f", underperf SPY>={rel_spy_min:.1%}"
            else:
                sig_msg = f"RSI({rsi_p})<{rsi_thr:g}"
            hedge_msg = ", SPY-hedged" if bool(self.hedge_spy) else ""
            conc_msg = (
                f" | max concurrent holdings={peak_concurrent}"
                + (f" (cap {max_conc})" if max_conc is not None else "")
            )
            print(
                f"  Buy-the-dip ({sig_msg}, hold {hold}d, {up}, {sp}, {rk_msg}, {wg}{hedge_msg}): "
                f"avg sleeve equity ≈ {avg_gross_pct:.2f}% | regime ≈ {avg_g:.3f} | "
                f"avg vol scale ≈ {float(scale_exec.mean()) if len(scale_exec) else float('nan'):.3f}"
                f"{conc_msg}"
            )

        return dip_daily_ret.rename("ls_short_ret")


@dataclass
class SellTheRipSleeve:
    """
    **Short sell-the-rip** sleeve — mirror image of :class:`BuyTheDipSleeve`.

    Signal (configurable): asset has a **rally day** (prior close-to-close return ≥ ``pct_rise``)
    while trading in a **downtrend** (close < SMA ``sma_trend_window``).  Among candidates
    each day, up to ``top_n`` names are **shorted** for ``hold_trading_days`` sessions, ranked
    by highest normalized ATR (most volatile expected to resume the down-leg fastest).

    Optional gate ``only_when_spy_bear=True``: only open new short positions when the SPY
    itself is below its SMA-200 (i.e. broad bear market).  Default is **off** — the strategy
    targets individual stock downtrends regardless of index regime.

    Returns are **negated** (short position) so a falling price produces positive P&L.
    Regime throttle and vol targeting applied to the resulting stream, identical to the
    long dip sleeve.
    """

    atr_period: int = 5
    hold_trading_days: int = 10
    vol_window: int = 10
    sma_trend_window: int = 200
    weighting: str = "equal"
    pct_rise_min: float = 0.03
    min_atr_norm_pct: float = 3.0
    only_when_spy_bear: bool = False

    def generate_returns(
        self,
        equity_dict: Dict[str, pd.DataFrame],
        top_n: int = 10,
        *,
        spy_df: pd.DataFrame | None = None,
        cash_annual_yield: float = 0.04,
        target_annual_vol: float = 0.10,
        vol_scale_cap: float = 2.0,
        vol_scale_smooth_days: int = 5,
        verbose: bool = True,
    ) -> pd.Series:
        if top_n <= 0:
            raise ValueError("top_n must be > 0")
        if not equity_dict:
            raise ValueError("equity_dict cannot be empty")

        atr_p = int(self.atr_period)
        hold = int(self.hold_trading_days)
        tw = int(self.sma_trend_window)
        pct_rise = float(self.pct_rise_min)
        min_atr_pct = float(self.min_atr_norm_pct)
        if not (0.0 < pct_rise < 1.0):
            raise ValueError("pct_rise_min must be in (0, 1)")

        tickers = list(equity_dict.keys())
        ret_pan: list[pd.Series] = []
        close_pan: list[pd.Series] = []
        atrn_pan: list[pd.Series] = []
        sma_pan: list[pd.Series] = []

        req = {"close", "ret"}
        for t in tickers:
            df = equity_dict[t]
            missing = req.difference(df.columns)
            if missing:
                raise KeyError(f"{t} missing required columns for sell-the-rip: {sorted(missing)}")

            idx = pd.to_datetime(df.index).tz_localize(None)
            df2 = df.copy()
            df2.index = idx
            df2 = df2.sort_index()

            close_s = df2["close"].astype(np.float64)
            hi, lo = _high_low_for_atr(df2, close_s)
            atr_s = _wilder_atr(hi, lo, close_s, atr_p)
            atrn_s = (atr_s / close_s.replace(0.0, np.nan)).astype(np.float64)
            sma_s = close_s.rolling(window=tw, min_periods=tw).mean()

            ret_pan.append(df2["ret"].astype(np.float64).rename(t))
            close_pan.append(close_s.rename(t))
            atrn_pan.append(atrn_s.rename(t))
            sma_pan.append(sma_s.rename(t))

        ret_df, _ = _align_panel_frames(ret_pan)
        close_df, _ = _align_panel_frames(close_pan)
        atrn_df, _ = _align_panel_frames(atrn_pan)
        sma_df, _ = _align_panel_frames(sma_pan)

        master_index = ret_df.index.union(close_df.index).sort_values()
        ret_df = ret_df.reindex(master_index)
        close_df = close_df.reindex(master_index).ffill()
        atrn_df = atrn_df.reindex(master_index)
        sma_df = sma_df.reindex(master_index)
        ret_chg_df = close_df.pct_change()

        vw = int(self.vol_window)
        vol_df = ret_df.rolling(window=vw, min_periods=vw).std(ddof=1)

        D, N = len(master_index), len(close_df.columns)
        if D < 2 or N < 2:
            return pd.Series(0.0, index=master_index, name="sell_rip_ret", dtype=np.float64)

        if top_n > N:
            raise ValueError(f"top_n={top_n} cannot exceed universe size N={N}")

        spy_panel = _resolve_spy_frame(spy_df, equity_dict)
        spy_close = spy_panel["close"].astype(np.float64)
        spy_close.index = pd.to_datetime(spy_close.index).tz_localize(None)
        spy_close = spy_close.reindex(master_index).ffill()
        if "sma_200" in spy_panel.columns:
            spy_sma = spy_panel["sma_200"].astype(np.float64)
            spy_sma.index = pd.to_datetime(spy_sma.index).tz_localize(None)
            spy_sma = spy_sma.reindex(master_index).ffill()
        else:
            spy_sma = spy_close.rolling(window=200, min_periods=200).mean()
        # Bear market gate: SPY < SMA200
        spy_bear_d = (spy_close < spy_sma).fillna(False).to_numpy(dtype=bool)

        close_arr = close_df.to_numpy(dtype=np.float64)
        sma_arr = sma_df.to_numpy(dtype=np.float64)
        atrn_arr = atrn_df.to_numpy(dtype=np.float64)
        ret_chg_arr = ret_chg_df.to_numpy(dtype=np.float64)
        vol_arr = vol_df.to_numpy(dtype=np.float64)

        # Signal: prior-day bounce in a downtrend
        base_ok = np.isfinite(close_arr) & np.isfinite(atrn_arr) & np.isfinite(ret_chg_arr)
        rip_sig = base_ok & (ret_chg_arr >= pct_rise)  # rally day
        rip_sig &= np.isfinite(sma_arr) & (close_arr < sma_arr)  # in downtrend
        rip_sig &= (atrn_arr * 100.0) > min_atr_pct  # enough volatility to matter

        active_entry = np.full(N, -1, dtype=np.int32)
        weights_mat = np.zeros((D, N), dtype=np.float64)

        wg = str(self.weighting).lower()
        if wg not in ("equal", "inv_vol"):
            raise ValueError(f"weighting must be 'equal' or 'inv_vol', got {self.weighting!r}")

        spy_gate = bool(self.only_when_spy_bear)

        for d in range(D):
            expired = (active_entry >= 0) & (d > active_entry + hold)
            active_entry[expired] = -1

            in_trade = (active_entry >= 0) & (active_entry < d) & (d <= active_entry + hold)
            idxs = np.flatnonzero(in_trade)
            if idxs.size > 0:
                if wg == "equal":
                    weights_mat[d, idxs] = 1.0 / float(idxs.size)
                else:
                    v = vol_arr[d, idxs]
                    inv = np.where(np.isfinite(v) & (v > 0.0), 1.0 / v, 0.0)
                    s = float(inv.sum())
                    if s > 0.0:
                        weights_mat[d, idxs] = inv / s

            flat = active_entry < 0
            open_ok = spy_bear_d if spy_gate else np.ones(D, dtype=bool)
            if open_ok[d]:
                cand = np.flatnonzero(flat & rip_sig[d])
                if cand.size > 0:
                    # Rank by highest ATR/px (most volatile downtrend names first)
                    sc = atrn_arr[d, cand].astype(np.float64)
                    key = np.where(np.isfinite(sc), sc, -np.inf)
                    order = np.argsort(-key, kind="stable")
                    take_n = min(int(top_n), int(cand.size))
                    chosen = cand[order[:take_n]]
                    active_entry[chosen] = d

        weights_d = pd.DataFrame(weights_mat, index=master_index, columns=close_df.columns)
        ret_filled = ret_df.fillna(0.0).astype(np.float64)

        # Regime: use SPY regime exposure but DON'T invert — short positions work in bear;
        # we keep regime throttle to reduce gross exposure in flat/rising markets if desired.
        # NOTE: regime_exec throttles exposure to 0 when SPY is in max-bear regime; for a
        # short sleeve that is conservative but consistent with risk management convention.
        spy_panel_for_regime = spy_panel
        regime_raw = _regime_exposure_series(spy_panel_for_regime, master_index)
        regime_exec = regime_raw.shift(1).fillna(1.0).clip(0.0, 1.0)

        daily_rf = float(cash_annual_yield) / 252.0
        w_arr = weights_d.to_numpy(dtype=np.float64)
        reg = regime_exec.to_numpy(dtype=np.float64)
        w_th = w_arr * reg[:, np.newaxis]
        # SHORT: negate returns (falling price → positive P&L)
        core = -(w_th * ret_filled.to_numpy(dtype=np.float64)).sum(axis=1)
        cash_carry = (1.0 - reg) * daily_rf
        ls_pre_vol = pd.Series(core + cash_carry, index=master_index, dtype=np.float64)

        rip_daily_ret, scale_exec = _apply_ls_vol_targeting(
            ls_pre_vol,
            daily_rf=daily_rf,
            target_annual_vol=float(target_annual_vol),
            max_scale=float(vol_scale_cap),
            smooth_window=int(vol_scale_smooth_days),
            vol_window=20,
        )

        if verbose:
            avg_gross_pct = float(np.nanmean(w_th.sum(axis=1))) * 100.0 if w_th.size else float("nan")
            avg_g = float(regime_exec.mean()) if len(regime_exec) else float("nan")
            gate_msg = "SPY<SMA200@entry" if spy_gate else "any SPY (no bear gate)"
            print(
                f"  Sell-the-rip (pct_rise>={pct_rise:.1%}, ATR%>{min_atr_pct:g}, "
                f"close<SMA{tw}, hold {hold}d, {gate_msg}, rank ATR/px, {wg}): "
                f"avg sleeve equity ≈ {avg_gross_pct:.2f}% | regime ≈ {avg_g:.3f} | "
                f"avg vol scale ≈ {float(scale_exec.mean()) if len(scale_exec) else float('nan'):.3f}"
            )

        return rip_daily_ret.rename("sell_rip_ret")


@dataclass
class SectorETFRotation:
    """
    Long-only sector rotation using SPDR sector ETFs.

    Each month-end, rank ETFs by ``aqr_mom`` (12-minus-1 on ``close``), hold the top ``top_k``
    at equal weight (100% invested in those names). Weights are forward-filled and executed
    from the next bar; any slack weight earns ``cash_annual_yield`` (e.g. early sample or all-NaN month).
    """

    def generate_returns(
        self,
        etf_dict: Dict[str, pd.DataFrame],
        top_k: int = 3,
        *,
        cash_annual_yield: float = 0.04,
        verbose: bool = True,
    ) -> pd.Series:
        if top_k <= 0:
            raise ValueError("top_k must be > 0")
        if not etf_dict:
            raise ValueError("etf_dict cannot be empty")

        tickers = sorted(etf_dict.keys())
        aqr_pan: list[pd.Series] = []
        ret_pan: list[pd.Series] = []

        for t in tickers:
            df = etf_dict[t]
            required = {"close", "ret", "aqr_mom"}
            missing = required.difference(df.columns)
            if missing:
                raise KeyError(f"{t} missing required columns: {sorted(missing)}")
            idx = pd.to_datetime(df.index).tz_localize(None)
            df2 = df.copy()
            df2.index = idx
            df2 = df2.sort_index()
            aqr_pan.append(df2["aqr_mom"].astype(np.float64).rename(t))
            ret_pan.append(df2["ret"].astype(np.float64).rename(t))

        aqr_df, _ = _align_panel_frames(aqr_pan)
        ret_df, _ = _align_panel_frames(ret_pan)
        master_index = aqr_df.index.union(ret_df.index).sort_values()
        aqr_df = aqr_df.reindex(master_index).ffill()
        ret_df = ret_df.reindex(master_index)
        ret_filled = ret_df.fillna(0.0).astype(np.float64)

        bm_index = aqr_df.resample("BME").last().index
        aqr_m = aqr_df.resample("BME").last()
        aqr_arr = aqr_m.to_numpy(dtype=np.float64)
        M, N = aqr_arr.shape
        if M < 1 or N < 1:
            return pd.Series(0.0, index=master_index, name="sector_rot_ret", dtype=np.float64)

        weights_m = np.zeros((M, N), dtype=np.float64)
        for m in range(M):
            mom = aqr_arr[m]
            valid = np.isfinite(mom)
            idx_valid = np.nonzero(valid)[0]
            if idx_valid.size == 0:
                continue
            k = min(int(top_k), int(idx_valid.size))
            sub_mom = mom[idx_valid]
            pick_local = np.argpartition(-sub_mom, k - 1)[:k]
            chosen = idx_valid[pick_local]
            weights_m[m, chosen] = 1.0 / float(k)

        weights_m_df = pd.DataFrame(weights_m, index=bm_index, columns=aqr_df.columns)
        weights_d = weights_m_df.reindex(master_index).ffill().shift(1).fillna(0.0)

        daily_rf = float(cash_annual_yield) / 252.0
        w = weights_d.to_numpy(dtype=np.float64)
        r = ret_filled.to_numpy(dtype=np.float64)
        core = (w * r).sum(axis=1)
        wsum = w.sum(axis=1)
        out = core + np.maximum(0.0, 1.0 - wsum) * daily_rf

        if verbose and M > 0:
            avg_k = float(np.mean([np.sum(weights_m[i] > 0) for i in range(M)]))
            print(f"  Sector ETF sleeve: top_k={int(top_k)} | avg ETFs held/month ≈ {avg_k:.2f}")

        return pd.Series(out, index=master_index, name="sector_rot_ret", dtype=np.float64)

    def generate_rebalance_log(
        self,
        etf_dict: Dict[str, pd.DataFrame],
        top_k: int = 3,
    ) -> pd.DataFrame:
        """
        Month-end rebalance audit: one row per held sector per signal month.

        Columns include signal_date (BME), effective_date (first session weights apply),
        ticker, weight, aqr_mom, rank, and entry/exit flags vs prior month.
        """
        if top_k <= 0:
            raise ValueError("top_k must be > 0")
        if not etf_dict:
            raise ValueError("etf_dict cannot be empty")

        tickers = sorted(etf_dict.keys())
        aqr_pan: list[pd.Series] = []
        for t in tickers:
            df = etf_dict[t]
            required = {"close", "ret", "aqr_mom"}
            missing = required.difference(df.columns)
            if missing:
                raise KeyError(f"{t} missing required columns: {sorted(missing)}")
            idx = pd.to_datetime(df.index).tz_localize(None)
            dfx = df.copy()
            dfx.index = idx
            dfx = dfx.sort_index()
            aqr_pan.append(dfx["aqr_mom"].astype(np.float64).rename(t))

        aqr_df, _ = _align_panel_frames(aqr_pan)
        master_index = aqr_df.index.sort_values()
        bm_index = aqr_df.resample("BME").last().index
        aqr_m = aqr_df.resample("BME").last()
        aqr_arr = aqr_m.to_numpy(dtype=np.float64)
        M, N = aqr_arr.shape
        if M < 1:
            return pd.DataFrame()

        rows: list[dict] = []
        prev_held: set[str] = set()
        for m in range(M):
            signal_dt = pd.Timestamp(bm_index[m]).normalize()
            mom = aqr_arr[m]
            valid = np.isfinite(mom)
            idx_valid = np.nonzero(valid)[0]
            if idx_valid.size == 0:
                prev_held = set()
                continue
            k = min(int(top_k), int(idx_valid.size))
            sub_mom = mom[idx_valid]
            pick_local = np.argpartition(-sub_mom, k - 1)[:k]
            chosen = idx_valid[pick_local]
            order = np.argsort(-sub_mom[pick_local])
            chosen = chosen[order]

            held = {tickers[int(j)] for j in chosen}
            entries = held - prev_held
            exits = prev_held - held

            after = master_index[master_index > signal_dt]
            effective_dt = (
                pd.Timestamp(after[0]).normalize()
                if len(after)
                else signal_dt
            )
            if m + 1 < M:
                period_end = pd.Timestamp(bm_index[m + 1]).normalize()
            else:
                period_end = pd.Timestamp(master_index[-1]).normalize()

            for rank_i, j in enumerate(chosen, start=1):
                tkr = tickers[int(j)]
                rows.append(
                    {
                        "signal_date": signal_dt.strftime("%Y-%m-%d"),
                        "effective_date": effective_dt.strftime("%Y-%m-%d"),
                        "period_end_date": period_end.strftime("%Y-%m-%d"),
                        "ticker": tkr,
                        "weight": 1.0 / float(k),
                        "aqr_mom": float(mom[j]),
                        "rank": int(rank_i),
                        "is_new_entry": tkr in entries,
                        "is_exit_from_prior": tkr in exits,
                        "top_k": int(k),
                    }
                )
            for tkr in sorted(exits):
                rows.append(
                    {
                        "signal_date": signal_dt.strftime("%Y-%m-%d"),
                        "effective_date": effective_dt.strftime("%Y-%m-%d"),
                        "period_end_date": period_end.strftime("%Y-%m-%d"),
                        "ticker": tkr,
                        "weight": 0.0,
                        "aqr_mom": float("nan"),
                        "rank": 0,
                        "is_new_entry": False,
                        "is_exit_from_prior": True,
                        "top_k": int(k),
                    }
                )
            prev_held = held

        return pd.DataFrame(rows)


@dataclass
class SectorETFLongShort:
    """
    **Market-neutral** SPDR sector rotation: long top-k / short bottom-k by ``aqr_mom``.

    Each leg is equal-weight; long weights sum to +1.0 and short weights sum to −1.0
    (200% gross, ~zero beta vs SPY in aggregate). Monthly rebalance, weights
    forward-filled and executed from the next bar.
    """

    def generate_returns(
        self,
        etf_dict: Dict[str, pd.DataFrame],
        top_k: int = 3,
        *,
        cash_annual_yield: float = 0.04,
        verbose: bool = True,
    ) -> pd.Series:
        if top_k <= 0:
            raise ValueError("top_k must be > 0")
        if not etf_dict:
            raise ValueError("etf_dict cannot be empty")

        tickers = sorted(etf_dict.keys())
        if top_k * 2 > len(tickers):
            raise ValueError(
                f"top_k={top_k} requires at least {2 * top_k} ETFs; got {len(tickers)}"
            )

        aqr_pan: list[pd.Series] = []
        ret_pan: list[pd.Series] = []
        for t in tickers:
            df = etf_dict[t]
            required = {"close", "ret", "aqr_mom"}
            missing = required.difference(df.columns)
            if missing:
                raise KeyError(f"{t} missing required columns: {sorted(missing)}")
            idx = pd.to_datetime(df.index).tz_localize(None)
            df2 = df.copy()
            df2.index = idx
            df2 = df2.sort_index()
            aqr_pan.append(df2["aqr_mom"].astype(np.float64).rename(t))
            ret_pan.append(df2["ret"].astype(np.float64).rename(t))

        aqr_df, _ = _align_panel_frames(aqr_pan)
        ret_df, _ = _align_panel_frames(ret_pan)
        master_index = aqr_df.index.union(ret_df.index).sort_values()
        aqr_df = aqr_df.reindex(master_index).ffill()
        ret_df = ret_df.reindex(master_index)
        ret_filled = ret_df.fillna(0.0).astype(np.float64)

        bm_index = aqr_df.resample("BME").last().index
        aqr_m = aqr_df.resample("BME").last()
        aqr_arr = aqr_m.to_numpy(dtype=np.float64)
        M, N = aqr_arr.shape
        if M < 1 or N < 1:
            return pd.Series(0.0, index=master_index, name="sector_ls_ret", dtype=np.float64)

        weights_m = np.zeros((M, N), dtype=np.float64)
        for m in range(M):
            mom = aqr_arr[m]
            valid = np.isfinite(mom)
            idx_valid = np.nonzero(valid)[0]
            if idx_valid.size < 2 * top_k:
                continue
            k = int(top_k)
            sub_mom = mom[idx_valid]
            long_local = np.argpartition(-sub_mom, k - 1)[:k]
            short_local = np.argpartition(sub_mom, k - 1)[:k]
            long_idx = idx_valid[long_local]
            short_idx = idx_valid[short_local]
            weights_m[m, long_idx] = 1.0 / float(k)
            weights_m[m, short_idx] = -1.0 / float(k)

        weights_m_df = pd.DataFrame(weights_m, index=bm_index, columns=aqr_df.columns)
        weights_d = weights_m_df.reindex(master_index).ffill().shift(1).fillna(0.0)

        daily_rf = float(cash_annual_yield) / 252.0
        w = weights_d.to_numpy(dtype=np.float64)
        r = ret_filled.to_numpy(dtype=np.float64)
        core = (w * r).sum(axis=1)
        gross = np.abs(w).sum(axis=1)
        # Unallocated notional (when month has no signal) earns cash yield.
        out = core + np.maximum(0.0, 1.0 - gross) * daily_rf

        if verbose and M > 0:
            avg_gross = float(np.mean(gross[gross > 0])) if np.any(gross > 0) else float("nan")
            print(
                f"  Sector L/S sleeve: top_k={int(top_k)} per leg | "
                f"avg gross exposure ≈ {avg_gross * 100:.1f}%"
            )

        return pd.Series(out, index=master_index, name="sector_ls_ret", dtype=np.float64)

    def generate_rebalance_log(
        self,
        etf_dict: Dict[str, pd.DataFrame],
        top_k: int = 3,
    ) -> pd.DataFrame:
        """Month-end audit: long (+) and short (−) legs with weights."""
        if top_k <= 0:
            raise ValueError("top_k must be > 0")
        tickers = sorted(etf_dict.keys())
        aqr_pan: list[pd.Series] = []
        for t in tickers:
            df = etf_dict[t]
            idx = pd.to_datetime(df.index).tz_localize(None)
            dfx = df.copy()
            dfx.index = idx
            dfx = dfx.sort_index()
            aqr_pan.append(dfx["aqr_mom"].astype(np.float64).rename(t))

        aqr_df, _ = _align_panel_frames(aqr_pan)
        master_index = aqr_df.index.sort_values()
        bm_index = aqr_df.resample("BME").last().index
        aqr_m = aqr_df.resample("BME").last()
        aqr_arr = aqr_m.to_numpy(dtype=np.float64)
        M, N = aqr_arr.shape
        if M < 1:
            return pd.DataFrame()

        rows: list[dict] = []
        for m in range(M):
            signal_dt = pd.Timestamp(bm_index[m]).normalize()
            mom = aqr_arr[m]
            valid = np.isfinite(mom)
            idx_valid = np.nonzero(valid)[0]
            if idx_valid.size < 2 * int(top_k):
                continue
            k = int(top_k)
            sub_mom = mom[idx_valid]
            long_local = np.argpartition(-sub_mom, k - 1)[:k]
            short_local = np.argpartition(sub_mom, k - 1)[:k]
            long_idx = idx_valid[np.argsort(-sub_mom[long_local])]
            short_idx = idx_valid[np.argsort(sub_mom[short_local])]

            after = master_index[master_index > signal_dt]
            effective_dt = pd.Timestamp(after[0]).normalize() if len(after) else signal_dt
            period_end = (
                pd.Timestamp(bm_index[m + 1]).normalize()
                if m + 1 < M
                else pd.Timestamp(master_index[-1]).normalize()
            )

            for rank_i, j in enumerate(long_idx, start=1):
                rows.append(
                    {
                        "signal_date": signal_dt.strftime("%Y-%m-%d"),
                        "effective_date": effective_dt.strftime("%Y-%m-%d"),
                        "period_end_date": period_end.strftime("%Y-%m-%d"),
                        "ticker": tickers[int(j)],
                        "leg": "long",
                        "weight": 1.0 / float(k),
                        "aqr_mom": float(mom[j]),
                        "rank": int(rank_i),
                        "top_k": k,
                    }
                )
            for rank_i, j in enumerate(short_idx, start=1):
                rows.append(
                    {
                        "signal_date": signal_dt.strftime("%Y-%m-%d"),
                        "effective_date": effective_dt.strftime("%Y-%m-%d"),
                        "period_end_date": period_end.strftime("%Y-%m-%d"),
                        "ticker": tickers[int(j)],
                        "leg": "short",
                        "weight": -1.0 / float(k),
                        "aqr_mom": float(mom[j]),
                        "rank": int(rank_i),
                        "top_k": k,
                    }
                )
        return pd.DataFrame(rows)


def _panic_stress_state(z: pd.Series, *, entry_z: float, exit_z: float) -> pd.Series:
    """
    Hysteresis on negative return z-scores only.

    Activates (1) when ``z <= -entry_z``, clears when ``z >= -exit_z`` (with ``entry_z``,
    ``exit_z`` positive magnitudes, e.g. 2.0 and 0.5).
    """
    active = np.zeros(len(z), dtype=np.int8)
    cur = np.int8(0)
    vals = z.to_numpy(dtype=np.float64)
    idx = z.index
    neg_exit = -float(exit_z)

    for i in range(len(vals)):
        zi = vals[i]
        if not np.isfinite(zi):
            active[i] = cur
            continue
        if cur == 0:
            if zi <= -float(entry_z):
                cur = np.int8(1)
        else:
            if zi >= neg_exit:
                cur = np.int8(0)
        active[i] = cur

    return pd.Series(active, index=idx, dtype=np.int8)


@dataclass
class DefensiveMeanReversionSleeve:
    """
    Bond-tilt sleeve triggered by **mean-reversion setup** on SPY: extremely negative
    realized daily returns vs a rolling window (z-score). When active, the sleeve holds
    long-duration Treasuries (TLT); otherwise it sits in cash at the risk-free rate.

    This is not a forecast of the equity bounce; it leans on the empirical tendency for
    flight-to-quality into bonds during equity return panics. Shifted 1 bar for execution.
    """

    z_window: int = 20
    entry_z: float = 2.0
    exit_z: float = 0.5
    tlt_weight_when_active: float = 1.0

    def generate_returns(
        self,
        spy_df: pd.DataFrame,
        tlt_df: pd.DataFrame,
        *,
        cash_annual_yield: float = 0.04,
        verbose: bool = True,
    ) -> pd.Series:
        for name, df in (("SPY", spy_df), ("TLT", tlt_df)):
            if "ret" not in df.columns:
                raise KeyError(f"{name} DataFrame must include 'ret'")

        master_index = pd.to_datetime(spy_df.index).tz_localize(None).union(
            pd.to_datetime(tlt_df.index).tz_localize(None)
        ).sort_values()

        spy_r = spy_df["ret"].astype(np.float64).copy()
        spy_r.index = pd.to_datetime(spy_df.index).tz_localize(None)
        spy_ret = spy_r.reindex(master_index).fillna(0.0)

        tlt_r = tlt_df["ret"].astype(np.float64).copy()
        tlt_r.index = pd.to_datetime(tlt_df.index).tz_localize(None)
        tlt_ret = tlt_r.reindex(master_index).fillna(0.0)

        z = rolling_zscore(spy_ret, int(self.z_window))
        stress = _panic_stress_state(z, entry_z=float(self.entry_z), exit_z=float(self.exit_z))
        w_target = stress.astype(np.float64) * float(self.tlt_weight_when_active)
        w = w_target.shift(1).fillna(0.0).clip(0.0, 1.0)

        daily_rf = float(cash_annual_yield) / 252.0
        out = w * tlt_ret + (1.0 - w) * daily_rf

        if verbose and len(out) > 0:
            frac = float((w > 0.5).mean()) if len(w) else 0.0
            print(
                f"  Defensive MR sleeve: z_window={int(self.z_window)} entry_z={float(self.entry_z):.2f} "
                f"exit_z={float(self.exit_z):.2f} | share of days TLT-tilted ≈ {frac:.1%}"
            )

        return pd.Series(out, index=master_index, name="defensive_mr_ret", dtype=np.float64)


@dataclass
class EnsembleManager:
    """
    Blends Macro (Tactical All Weather) + monthly AQR L/S + optional buy-the-dip (long-only) +
    sector rotation + defensive MR sleeve.
    """

    # For a true "overlay" the macro book stays fully invested (1.0) and we add
    # a market-neutral L/S sleeve on top (0.30). This can result in >1.0 gross capital.
    weight_macro: float = 1.00
    weight_alpha: float = 0.30
    weight_ls_short: float = 0.0
    weight_sector: float = 0.0
    weight_defensive: float = 0.0

    def build_ensemble(
        self,
        macro_dict: Dict[str, pd.DataFrame],
        equity_dict: Dict[str, pd.DataFrame] | None = None,
        *,
        cash_annual_yield: float = 0.04,
        top_n: int = 10,
        ls_daily_ret: pd.Series | None = None,
        sector_map: Dict[str, str] | None = None,
        sector_daily_ret: pd.Series | None = None,
        defensive_daily_ret: pd.Series | None = None,
        ls_short_daily_ret: pd.Series | None = None,
        ls_diagnostics: dict | None = None,
    ) -> pd.DataFrame:
        """
        Parameters
        ----------
        macro_dict
            ticker -> daily DataFrame with columns close, ret, sma_200, aqr_mom.
        equity_dict
            Equity universe for the AQR / cross-sectional L/S engine (required unless
            ``ls_daily_ret`` is provided).
        cash_annual_yield
            Annualized yield applied to cash weight (macro book).
        top_n
            L/S leg size (AQR engine only).
        ls_daily_ret
            If provided, skip the built-in :class:`CrossSectionalMomentum` path and use this
            daily L/S return series (e.g. from :class:`MLMomentumEngine`).
        sector_map
            Passed to :class:`CrossSectionalMomentum` when building the AQR sleeve (ignored if
            ``ls_daily_ret`` is set).
        sector_daily_ret
            Optional :class:`SectorETFRotation` daily returns. Used only if ``weight_sector > 0``.
        defensive_daily_ret
            Optional :class:`DefensiveMeanReversionSleeve` daily returns. Used only if
            ``weight_defensive > 0``.
        ls_short_daily_ret
            Optional :class:`BuyTheDipSleeve` daily returns (long-only; column name unchanged).
            Used only if ``weight_ls_short > 0``.
        ls_diagnostics
            If ``ls_daily_ret`` is None and this dict is provided, passed through to
            :meth:`CrossSectionalMomentum.generate_ls_returns` and filled with sleeve stats.

        Returns
        -------
        pd.DataFrame
            Columns:
              * macro_cumulative_ret
              * ls_cumulative_ret
              * ls_short_cumulative_ret (zeros if buy-the-dip sleeve off)
              * sector_cumulative_ret (zeros if sector sleeve off)
              * defensive_cumulative_ret (zeros if defensive sleeve off)
              * ensemble_cumulative_ret
        """
        # Step A: Tactical All Weather beta core (keep SPY in baseline).
        tactical = TacticalAllWeatherManager(baseline_weights=BASE_WEIGHTS)
        macro_df = tactical.build_portfolio(macro_dict, cash_annual_yield=cash_annual_yield)
        macro_daily_ret = macro_df["portfolio_bar_ret"].astype(np.float64)

        # Step B: Alpha satellite L/S
        if ls_daily_ret is None:
            if not equity_dict:
                raise ValueError("equity_dict is required when ls_daily_ret is not provided.")
            spy_macro = macro_dict.get("SPY")
            if spy_macro is None or spy_macro.empty:
                raise ValueError("macro_dict must include a non-empty 'SPY' DataFrame for the L/S regime throttle.")
            ls_engine = CrossSectionalMomentum()
            ls_daily_ret = ls_engine.generate_ls_returns(
                equity_dict,
                top_n=top_n,
                spy_df=spy_macro,
                cash_annual_yield=cash_annual_yield,
                sector_map=sector_map,
                diagnostics=ls_diagnostics,
            ).astype(np.float64)
        else:
            ls_daily_ret = ls_daily_ret.astype(np.float64)

        use_sector = float(self.weight_sector) > 0.0 and sector_daily_ret is not None
        sec_al: pd.Series
        if use_sector:
            sec_al = sector_daily_ret.astype(np.float64)
        else:
            sec_al = pd.Series(0.0, index=macro_daily_ret.index, dtype=np.float64)

        use_def = float(self.weight_defensive) > 0.0 and defensive_daily_ret is not None
        def_al: pd.Series
        if use_def:
            def_al = defensive_daily_ret.astype(np.float64)
        else:
            def_al = pd.Series(0.0, index=macro_daily_ret.index, dtype=np.float64)

        use_ls_s = float(self.weight_ls_short) > 0.0 and ls_short_daily_ret is not None
        ls_s_al: pd.Series
        if use_ls_s:
            ls_s_al = ls_short_daily_ret.astype(np.float64)
        else:
            ls_s_al = pd.Series(0.0, index=macro_daily_ret.index, dtype=np.float64)

        # Step C: Capital split + alignment
        common_index = macro_daily_ret.index.intersection(ls_daily_ret.index)
        if use_sector:
            common_index = common_index.intersection(sec_al.index)
        if use_def:
            common_index = common_index.intersection(def_al.index)
        if use_ls_s:
            common_index = common_index.intersection(ls_s_al.index)

        macro_al = macro_daily_ret.reindex(common_index).fillna(0.0)
        ls_al = ls_daily_ret.reindex(common_index).fillna(0.0)
        sec_al = sec_al.reindex(common_index).fillna(0.0)
        def_al = def_al.reindex(common_index).fillna(0.0)
        ls_s_al = ls_s_al.reindex(common_index).fillna(0.0)

        ws = float(self.weight_sector) if use_sector else 0.0
        wd = float(self.weight_defensive) if use_def else 0.0
        wls = float(self.weight_ls_short) if use_ls_s else 0.0
        ensemble_daily_ret = (
            macro_al * self.weight_macro
            + ls_al * self.weight_alpha
            + ls_s_al * wls
            + sec_al * ws
            + def_al * wd
        )

        macro_cum = (1.0 + macro_al).cumprod() - 1.0
        ls_cum = (1.0 + ls_al).cumprod() - 1.0
        ls_short_cum = (1.0 + ls_s_al).cumprod() - 1.0
        sector_cum = (1.0 + sec_al).cumprod() - 1.0
        defensive_cum = (1.0 + def_al).cumprod() - 1.0
        ensemble_cum = (1.0 + ensemble_daily_ret).cumprod() - 1.0

        out = pd.DataFrame(
            {
                "macro_cumulative_ret": macro_cum,
                "ls_cumulative_ret": ls_cum,
                "ls_short_cumulative_ret": ls_short_cum,
                "sector_cumulative_ret": sector_cum,
                "defensive_cumulative_ret": defensive_cum,
                "ensemble_cumulative_ret": ensemble_cum,
            },
            index=common_index,
        )
        return out


def plot_ensemble(ensemble_df: pd.DataFrame, *, save_path: str | None = None) -> None:
    """
    Plot Macro vs L/S vs optional sector vs Ensemble equity curves.

    Expects:
      * macro_cumulative_ret
      * ls_cumulative_ret
      * ensemble_cumulative_ret
      * sector_cumulative_ret (optional; plotted if present and non-flat)
    """
    required = [
        "macro_cumulative_ret",
        "ls_cumulative_ret",
        "ensemble_cumulative_ret",
    ]
    missing = [c for c in required if c not in ensemble_df.columns]
    if missing:
        raise KeyError(f"ensemble_df missing columns: {missing}")

    try:
        import matplotlib.pyplot as plt
    except ImportError as e:  # pragma: no cover
        raise ImportError("matplotlib is required for plotting.") from e

    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(
        ensemble_df.index,
        ensemble_df["macro_cumulative_ret"].values,
        label="Macro (Tactical All Weather)",
        linewidth=2.0,
    )
    ax.plot(
        ensemble_df.index,
        ensemble_df["ls_cumulative_ret"].values,
        label="Equity L/S Momentum",
        linewidth=1.6,
    )
    if "sector_cumulative_ret" in ensemble_df.columns:
        s = ensemble_df["sector_cumulative_ret"].astype(np.float64)
        if float(np.nanmax(np.abs(s.values))) > 1e-12:
            ax.plot(
                ensemble_df.index,
                s.values,
                label="Sector ETF rotation (standalone leg)",
                linewidth=1.4,
                linestyle=":",
            )
    if "ls_short_cumulative_ret" in ensemble_df.columns:
        ls_s = ensemble_df["ls_short_cumulative_ret"].astype(np.float64)
        if float(np.nanmax(np.abs(ls_s.values))) > 1e-12:
            ax.plot(
                ensemble_df.index,
                ls_s.values,
                label="Buy the dip (long-only sleeve)",
                linewidth=1.3,
                linestyle=(0, (3, 1, 1, 1)),
            )
    if "defensive_cumulative_ret" in ensemble_df.columns:
        d = ensemble_df["defensive_cumulative_ret"].astype(np.float64)
        if float(np.nanmax(np.abs(d.values))) > 1e-12:
            ax.plot(
                ensemble_df.index,
                d.values,
                label="Defensive MR (TLT panic tilt)",
                linewidth=1.3,
                linestyle="-.",
            )
    ax.plot(
        ensemble_df.index,
        ensemble_df["ensemble_cumulative_ret"].values,
        label="Ensemble (blended)",
        linewidth=2.5,
        linestyle="--",
    )
    ax.axhline(0.0, color="gray", linewidth=0.5)
    ax.set_title("Multi-Strategy Ensemble Equity Curves")
    ax.set_ylabel("Cumulative Return")
    ax.grid(True, alpha=0.3)
    ax.legend(loc="best")
    fig.tight_layout()
    if save_path:
        fig.savefig(save_path, dpi=150)
        plt.close(fig)
    else:
        plt.show()
        plt.close(fig)


def plot_ls_standalone(ls_daily_ret: pd.Series, *, save_path: str | None = None) -> None:
    """
    Plot cumulative equity for the standalone L/S momentum sleeve (daily simple returns).
    """
    r = ls_daily_ret.fillna(0.0).astype(np.float64)
    cum = (1.0 + r).cumprod() - 1.0

    try:
        import matplotlib.pyplot as plt
    except ImportError as e:  # pragma: no cover
        raise ImportError("matplotlib is required for plotting.") from e

    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(cum.index, cum.values, label="Equity L/S Momentum (standalone)", linewidth=2.0, color="C1")
    ax.axhline(0.0, color="gray", linewidth=0.5)
    ax.set_title("Cross-Sectional Equity L/S Momentum")
    ax.set_ylabel("Cumulative Return")
    ax.grid(True, alpha=0.3)
    ax.legend(loc="best")
    fig.tight_layout()
    if save_path:
        fig.savefig(save_path, dpi=150)
        plt.close(fig)
    else:
        plt.show()
        plt.close(fig)


def plot_defensive_mr_standalone(daily_ret: pd.Series, *, save_path: str | None = None) -> None:
    """Plot cumulative equity for the standalone defensive MR (TLT panic tilt) sleeve."""
    r = daily_ret.fillna(0.0).astype(np.float64)
    cum = (1.0 + r).cumprod() - 1.0

    try:
        import matplotlib.pyplot as plt
    except ImportError as e:  # pragma: no cover
        raise ImportError("matplotlib is required for plotting.") from e

    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(cum.index, cum.values, label="Defensive MR sleeve", linewidth=2.0, color="C3")
    ax.axhline(0.0, color="gray", linewidth=0.5)
    ax.set_title("Defensive Mean Reversion (SPY return z → TLT tilt)")
    ax.set_ylabel("Cumulative Return")
    ax.grid(True, alpha=0.3)
    ax.legend(loc="best")
    fig.tight_layout()
    if save_path:
        fig.savefig(save_path, dpi=150)
        plt.close(fig)
    else:
        plt.show()
        plt.close(fig)


def plot_sector_rotation_standalone(daily_ret: pd.Series, *, save_path: str | None = None) -> None:
    """Plot cumulative equity for the standalone sector ETF rotation sleeve."""
    r = daily_ret.fillna(0.0).astype(np.float64)
    cum = (1.0 + r).cumprod() - 1.0

    try:
        import matplotlib.pyplot as plt
    except ImportError as e:  # pragma: no cover
        raise ImportError("matplotlib is required for plotting.") from e

    fig, ax = plt.subplots(figsize=(10, 5))
    ax.plot(cum.index, cum.values, label="Sector ETF momentum (top-k)", linewidth=2.0, color="C2")
    ax.axhline(0.0, color="gray", linewidth=0.5)
    ax.set_title("SPDR Sector ETF Rotation (12-1 momentum)")
    ax.set_ylabel("Cumulative Return")
    ax.grid(True, alpha=0.3)
    ax.legend(loc="best")
    fig.tight_layout()
    if save_path:
        fig.savefig(save_path, dpi=150)
        plt.close(fig)
    else:
        plt.show()
        plt.close(fig)