"""XM02 — Multi-L z-blend momentum
Score = mean of xs_zscore(past_return(close, L)) over a set of lookback windows L.
Compare two window sets: {30,90} (XS01-like) vs {20,60,120} (extended).
Grid: 5 study_xs calls total — vary universe / windows / H.
"""
import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/xsec")
import xslib as xs
import numpy as np

# ── score helpers ─────────────────────────────────────────────────────────────

def blend_mom(close, lookbacks):
    """Mean of xs_zscore(past_return(close, L)) for each L in lookbacks."""
    scores = [xs.xs_zscore(xs.past_return(close, L)) for L in lookbacks]
    stacked = np.stack(scores, axis=2)          # (n_days, n_assets, n_L)
    return np.nanmean(stacked, axis=2)           # (n_days, n_assets)


L_SHORT  = [30, 90]            # mirrors XS01 blend
L_LONG   = [20, 60, 120]       # extended set
L_WIDE   = [20, 60, 90, 120]   # even wider blend

# ── 5 backtests ───────────────────────────────────────────────────────────────

results = []

# 1. XS01-equivalent blend {30,90} on ALL universe — baseline reference
rep1 = xs.study_xs(
    "XM02-3090-all",
    lambda P: blend_mom(P.close, [30, 90]),
    universe="all", H=10, k=5, long_short=True,
)
print(xs.fmt(rep1))
results.append(rep1)

# 2. Extended blend {20,60,120} on ALL universe
rep2 = xs.study_xs(
    "XM02-206012-all",
    lambda P: blend_mom(P.close, [20, 60, 120]),
    universe="all", H=10, k=5, long_short=True,
)
print(xs.fmt(rep2))
results.append(rep2)

# 3. Extended blend {20,60,120} on MAJORS (19 alts — XS01 universe)
rep3 = xs.study_xs(
    "XM02-206012-majors",
    lambda P: blend_mom(P.close, [20, 60, 120]),
    universe="majors", H=10, k=5, long_short=True,
)
print(xs.fmt(rep3))
results.append(rep3)

# 4. Wide blend {20,60,90,120} on ALL, shorter rebalance H=5
rep4 = xs.study_xs(
    "XM02-wide-H5-all",
    lambda P: blend_mom(P.close, [20, 60, 90, 120]),
    universe="all", H=5, k=5, long_short=True,
)
print(xs.fmt(rep4))
results.append(rep4)

# 5. Wide blend on ALL, longer H=20 (less turnover)
rep5 = xs.study_xs(
    "XM02-wide-H20-all",
    lambda P: blend_mom(P.close, [20, 60, 90, 120]),
    universe="all", H=20, k=5, long_short=True,
)
print(xs.fmt(rep5))
results.append(rep5)

# ── pick BEST by: earns_slot > hold-out sharpe > distinctness ────────────────

def _score(r):
    earns   = 1 if r["earns_slot"] else 0
    verdict = 1 if r["marginal"].get("verdict") == "ADDS" else 0
    hold_sh = r["holdout"].get("sharpe", -99)
    full_sh = r["full"]["sharpe"]
    dist    = 1 if (r["corr_xs01"] or 1.0) < 0.6 else 0
    return (earns, verdict, hold_sh, full_sh, dist)

best = max(results, key=_score)

print("\n" + "=" * 60)
print("BEST CONFIG:")
print(xs.fmt(best))
print("JSON:", xs.as_json(best))