"""XM09 — Market-trend-gated momentum
Score = XS momentum (past_return L=60) but ACTIVE only when the equal-weight
market return trailing sum over L days is > 0; else 0 (flat).

Idea: plain cross-sectional momentum tends to fail during broad market downtrends
(all alts fall together, 'market neutral' still bleeds). Gate it off when the market
equal-weight trend is negative. Distinct from XS01 (plain XS mom) because it selectively
silences the strategy in bear regimes, producing a different return pattern.

Grid (<=5 calls): vary universe / H / k.
"""
import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/xsec")
import xslib as xs
import numpy as np


# ---------------------------------------------------------------------------
# SCORE: market-trend-gated momentum
# ---------------------------------------------------------------------------
def xm09_score(P, L=60):
    """
    score[i, a] = past_return(close, L)[i, a] * market_up[i]
    market_up[i] = 1 if trailing-L sum of equal-weight market daily returns > 0, else 0.
    Fully causal: uses close and ret up to row i only.
    """
    close = P.close  # (n, A)
    ret = P.ret      # (n, A) simple daily returns

    n, A = close.shape

    # Base momentum score (causal)
    pr = xs.past_return(close, L)  # (n, A), nan for i < L

    # Equal-weight market return per day (causal, mean across assets ignoring NaN)
    mret = xs.market_ret(ret)  # (n,) equal-weight market return

    # Trailing L-day cumulative market return (causal rolling sum)
    # roll_mean(mat, win) works on 2D; use it on a column vector
    mret_2d = mret.reshape(-1, 1)  # (n, 1)
    mkt_trail = xs.roll_mean(mret_2d, L) * L  # approximate trailing sum via roll_mean * L
    # Actually compute exact rolling sum using cumsum trick (causal)
    mret_cumsum = np.cumsum(mret)  # (n,)
    mkt_rolling_sum = np.empty(n)
    mkt_rolling_sum[:] = np.nan
    for i in range(L - 1, n):
        mkt_rolling_sum[i] = mret_cumsum[i] - (mret_cumsum[i - L] if i >= L else 0.0)

    # Market uptrend gate: 1 when trailing sum > 0, else 0
    market_up = (mkt_rolling_sum > 0).astype(float)  # (n,)
    market_up[:L - 1] = np.nan  # not enough history

    # Broadcast: score is 0 (flat) when market is down
    score = pr * market_up[:, None]  # (n, A)

    return score


# ---------------------------------------------------------------------------
# GRID (<=5 calls)
# ---------------------------------------------------------------------------
results = []

# 1. Base: majors, L=60, H=10, k=5, long_short
rep1 = xs.study_xs(
    "XM09_majors_H10_k5_ls",
    lambda P: xm09_score(P, 60),
    universe="majors", H=10, k=5, long_short=True, target_vol=0.20
)
print(xs.fmt(rep1))
print("JSON:", xs.as_json(rep1))
results.append(rep1)

# 2. All assets, L=60, H=10, k=5, long_short
rep2 = xs.study_xs(
    "XM09_all_H10_k5_ls",
    lambda P: xm09_score(P, 60),
    universe="all", H=10, k=5, long_short=True, target_vol=0.20
)
print(xs.fmt(rep2))
print("JSON:", xs.as_json(rep2))
results.append(rep2)

# 3. Majors, H=10, k=5, long-only (when market is up, just go long top-k)
rep3 = xs.study_xs(
    "XM09_majors_H10_k5_lo",
    lambda P: xm09_score(P, 60),
    universe="majors", H=10, k=5, long_short=False, target_vol=0.20
)
print(xs.fmt(rep3))
print("JSON:", xs.as_json(rep3))
results.append(rep3)

# 4. All assets, H=20, k=5, long_short (slower rebalance)
rep4 = xs.study_xs(
    "XM09_all_H20_k5_ls",
    lambda P: xm09_score(P, 60),
    universe="all", H=20, k=5, long_short=True, target_vol=0.20
)
print(xs.fmt(rep4))
print("JSON:", xs.as_json(rep4))
results.append(rep4)

# 5. Majors, H=10, k=7, long_short (wider buckets on smaller universe)
rep5 = xs.study_xs(
    "XM09_majors_H10_k7_ls",
    lambda P: xm09_score(P, 60),
    universe="majors", H=10, k=7, long_short=True, target_vol=0.20
)
print(xs.fmt(rep5))
print("JSON:", xs.as_json(rep5))
results.append(rep5)

# ---------------------------------------------------------------------------
# PICK BEST
# ---------------------------------------------------------------------------
def score_result(r):
    """Prefer earns_slot, then hold-out sharpe, then distinctness from XS01."""
    earns = r.get("earns_slot", False)
    ho = r.get("holdout", {}).get("sharpe", -999)
    corr = abs(r.get("corr_xs01", 1.0))
    return (int(earns), ho, -corr)

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