"""XS03b — Beta-hedged momentum.

IDEA: Instead of plain cross-sectional momentum (XS01), use RESIDUAL momentum:
  score = cumulative idiosyncratic return over lookback L.
  The residual is ret - beta*mkt_ret (rolling beta vs equal-weight panel),
  so each asset's score reflects ONLY its idiosyncratic drift, stripping
  out shared market moves. The resulting book is already dollar-neutral
  (long-short) but also implicitly market-beta-neutral because the signal
  itself filters out mkt co-movement.

WHY DISTINCT FROM XS01: Plain XS01 ranks on raw momentum; the top assets
in a bull market are often the highest-beta assets (not idiosyncratic winners).
Beta-hedged momentum ranks on WHAT IS LEFT after removing mkt factor:
  - In bull: avoids accidental overweight of market beta
  - In bear: avoids accidental short of low-beta (defensive) assets
  - Net: the book is more idiosyncratic and less correlated to raw XS momentum.

GRID (5 backtests max):
  1. majors, L=30, beta_win=90, H=10, k=5, LS
  2. majors, L=60, beta_win=90, H=10, k=5, LS
  3. all,    L=30, beta_win=90, H=10, k=5, LS
  4. all,    L=60, beta_win=90, H=10, k=5, LS
  5. all,    blend L=[30,60] residuals, beta_win=90, H=10, k=5, LS  (like XS01 blend)

Pick best by: earns_slot > holdout_sharpe > distinctness from XS01.
"""

import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/xsec")
import xslib as xs
import numpy as np


def residual_momentum(P, L, beta_win=90):
    """Cumulative idiosyncratic return over L days (causal).

    Residual daily ret = ret - rolling_beta * market_ret.
    Cumulate over L days to get momentum score on idiosyncratic drift.
    """
    resid = xs.residual_return(P.ret, beta_win)   # (n_days x n_assets)
    # Cumulate residuals over L days (causal: sum of past L residual daily rets)
    n, A = resid.shape
    cum = np.full((n, A), np.nan)
    for i in range(L, n):
        cum[i] = np.nansum(resid[i - L:i], axis=0)
    return cum


def blend_residual_mom(P, Ls=(30, 60), beta_win=90):
    """Cross-sectional z-score blend of multiple lookback residual momentums."""
    scores = []
    for L in Ls:
        s = residual_momentum(P, L, beta_win)
        scores.append(xs.xs_zscore(s))
    return np.nanmean(scores, axis=0)


print("=== XS03b: Beta-hedged Momentum ===\n")

results = []

# 1. majors, L=30
print("Run 1/5: majors, L=30, beta_win=90, H=10, k=5 LS")
r1 = xs.study_xs(
    "XS03b-MAJ-L30",
    lambda P: residual_momentum(P, 30, 90),
    universe="majors", H=10, k=5, long_short=True
)
print(xs.fmt(r1))
results.append(r1)

# 2. majors, L=60
print("\nRun 2/5: majors, L=60, beta_win=90, H=10, k=5 LS")
r2 = xs.study_xs(
    "XS03b-MAJ-L60",
    lambda P: residual_momentum(P, 60, 90),
    universe="majors", H=10, k=5, long_short=True
)
print(xs.fmt(r2))
results.append(r2)

# 3. all, L=30
print("\nRun 3/5: all, L=30, beta_win=90, H=10, k=5 LS")
r3 = xs.study_xs(
    "XS03b-ALL-L30",
    lambda P: residual_momentum(P, 30, 90),
    universe="all", H=10, k=5, long_short=True
)
print(xs.fmt(r3))
results.append(r3)

# 4. all, L=60
print("\nRun 4/5: all, L=60, beta_win=90, H=10, k=5 LS")
r4 = xs.study_xs(
    "XS03b-ALL-L60",
    lambda P: residual_momentum(P, 60, 90),
    universe="all", H=10, k=5, long_short=True
)
print(xs.fmt(r4))
results.append(r4)

# 5. all, blend [30,60]
print("\nRun 5/5: all, blend L=[30,60], beta_win=90, H=10, k=5 LS")
r5 = xs.study_xs(
    "XS03b-ALL-BLEND",
    lambda P: blend_residual_mom(P, (30, 60), 90),
    universe="all", H=10, k=5, long_short=True
)
print(xs.fmt(r5))
results.append(r5)

# --- Pick best ---
def score_config(r):
    """Priority: earns_slot > holdout_sharpe > full_sharpe > distinctness."""
    slot = 1 if r.get("earns_slot") else 0
    hs = r["holdout"].get("sharpe", -9)
    fs = r["full"]["sharpe"]
    corr = r.get("corr_xs01") or 1.0
    distinct = 1.0 - abs(corr)
    return (slot, hs, fs, distinct)

best = max(results, key=score_config)

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