"""XU01 — Momentum Universe Sweep
MECHANISM: Best momentum z-blend (blend of past_return z-scores at L=30 and L=90),
run on different universe sizes: majors (19), top20, top30, all (49).

Goal: map where cross-sectional momentum alpha lives — does expanding to top20/top30/all
help or hurt vs the tight 19-major universe of XS01?

Grid (<=5 backtests):
  1. majors  (19) — baseline reference, should approach XS01
  2. top20        — add one more liquid alt
  3. top30        — mid-tier liquidity
  4. all     (49) — known to dilute (confirm)
  5. top30, long-only (best mid-tier config variant)

Signal: xs_zscore(past_return(close,30)) + xs_zscore(past_return(close,90)) — same blend as XS01.
H=10, k=5, long_short=True (except run 5 long-only).
"""
import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/xsec")
import xslib as xs
import numpy as np

print("XU01 — Momentum Universe Sweep")
print("=" * 60)


def blend_score(P):
    """Z-blend of 30d and 90d momentum — same signal as XS01 but on any universe."""
    z30 = xs.xs_zscore(xs.past_return(P.close, 30))
    z90 = xs.xs_zscore(xs.past_return(P.close, 90))
    return np.nanmean(np.stack([z30, z90], axis=0), axis=0)


# 1) Majors (19) — baseline
rep1 = xs.study_xs(
    "XU01_MAJORS",
    blend_score,
    universe="majors", H=10, k=5, long_short=True
)
print(xs.fmt(rep1))
print("JSON:", xs.as_json(rep1))
print()

# 2) Top-20 by $-volume
rep2 = xs.study_xs(
    "XU01_TOP20",
    blend_score,
    universe=20, H=10, k=5, long_short=True
)
print(xs.fmt(rep2))
print("JSON:", xs.as_json(rep2))
print()

# 3) Top-30 by $-volume
rep3 = xs.study_xs(
    "XU01_TOP30",
    blend_score,
    universe=30, H=10, k=5, long_short=True
)
print(xs.fmt(rep3))
print("JSON:", xs.as_json(rep3))
print()

# 4) All (49) — expected dilution
rep4 = xs.study_xs(
    "XU01_ALL",
    blend_score,
    universe="all", H=10, k=5, long_short=True
)
print(xs.fmt(rep4))
print("JSON:", xs.as_json(rep4))
print()

# 5) Top-30, long-only — does dropping the short leg help with mid-tier names?
rep5 = xs.study_xs(
    "XU01_TOP30_LO",
    blend_score,
    universe=30, H=10, k=5, long_short=False
)
print(xs.fmt(rep5))
print("JSON:", xs.as_json(rep5))
print()


# Pick best: prefer earns_slot, then hold-out sharpe, then full sharpe, then distinctness
all_reps = [rep1, rep2, rep3, rep4, rep5]


def score_rep(r):
    earns = int(r.get("earns_slot", False))
    hold_sh = (r.get("holdout") or {}).get("sharpe", -9)
    full_sh = (r.get("full") or {}).get("sharpe", -9)
    corr_xs01 = r.get("corr_xs01") or 1.0
    distinctness = 1 - abs(corr_xs01)
    return (earns, hold_sh, full_sh, distinctness)


best = max(all_reps, key=score_rep)
print("=" * 60)
print(f"BEST CONFIG: {best['name']}")
print(xs.fmt(best))
print("JSON:", xs.as_json(best))