"""XM03 — Vol-Scaled (Risk-Adjusted) Momentum
MECHANISM: Score = past_return(close, L) / roll_std(ret, L)
This is a Sharpe-like signal: normalises raw momentum by the volatility of that asset
over the same window. Should favour assets that moved up *smoothly* (high Sharpe trend)
over those that had large one-off jumps (noisy high return).
Grid: L in {30, 60, 90}; universe in {all, majors}; long_short True/False.
Goal: test if risk-adjusted scoring is DISTINCT from plain XS01 momentum and ADDS to the
live TP01+XS01+VRP01 portfolio.
"""
import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/xsec")
import xslib as xs
import numpy as np


def vol_adj_momentum(P, L: int) -> np.ndarray:
    """Causal Sharpe-like score: past_return / roll_std(ret, L).
    Higher = long. Returns (n_days x n_assets).
    Avoid divide-by-zero by replacing 0-vol rows with NaN -> harness treats NaN as neutral.
    """
    pr = xs.past_return(P.close, L)          # causal past return over L days
    rv = xs.roll_std(P.ret, L)               # causal rolling std of daily returns
    # Replace zeros/near-zeros with NaN to avoid Inf
    rv_safe = np.where(rv < 1e-8, np.nan, rv)
    score = pr / rv_safe
    return score


print("XM03 — Vol-Scaled (Risk-Adjusted) Momentum")
print("=" * 60)

# 1) All universe, L=30 (short horizon vol-adj)
rep1 = xs.study_xs(
    "XM03_ALL_L30",
    lambda P: vol_adj_momentum(P, 30),
    universe="all", H=10, k=5, long_short=True
)
print(xs.fmt(rep1))
print("JSON:", xs.as_json(rep1))
print()

# 2) All universe, L=60 (medium horizon)
rep2 = xs.study_xs(
    "XM03_ALL_L60",
    lambda P: vol_adj_momentum(P, 60),
    universe="all", H=10, k=5, long_short=True
)
print(xs.fmt(rep2))
print("JSON:", xs.as_json(rep2))
print()

# 3) All universe, L=90 (long horizon)
rep3 = xs.study_xs(
    "XM03_ALL_L90",
    lambda P: vol_adj_momentum(P, 90),
    universe="all", H=10, k=5, long_short=True
)
print(xs.fmt(rep3))
print("JSON:", xs.as_json(rep3))
print()

# 4) Majors only (same universe as XS01), L=60 — can vol-adj beat plain MOM on XS01 turf?
rep4 = xs.study_xs(
    "XM03_MAJORS_L60",
    lambda P: vol_adj_momentum(P, 60),
    universe="majors", H=10, k=5, long_short=True
)
print(xs.fmt(rep4))
print("JSON:", xs.as_json(rep4))
print()

# 5) All universe, L=60, long-only — does vol-adj work as selection filter?
rep5 = xs.study_xs(
    "XM03_ALL_L60_LO",
    lambda P: vol_adj_momentum(P, 60),
    universe="all", H=10, k=5, long_short=False
)
print(xs.fmt(rep5))
print("JSON:", xs.as_json(rep5))
print()


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


all_reps = [rep1, rep2, rep3, rep4, rep5]
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))