"""XU04 — Liquidity-filtered momentum
MECHANISM: Cross-sectional momentum, but restrict to the top-N assets by RECENT (rolling 60d)
median dollar-volume rather than the static all-panel. The idea: momentum signal is cleaner on
liquid names; illiquid tail adds noise. Compare:
  1. Dynamic top-20 by rolling $-vol (vs static top-20 from XU01)
  2. Dynamic top-20, adjusted momentum (skip 1d to reduce microstructure noise): L=2..31
  3. Static majors (19) with skip-1 momentum — XS01-style but skip-1 to reduce echo
  4. Dynamic top-25 rolling-liquidity blend [30,90] — slightly wider universe
  5. Dynamic top-20 rolling-liquidity blend [30,90], H=5 (faster rebalance)

Key difference from XS01: the UNIVERSE is determined dynamically (rolling 60d dollar-volume
rank) rather than the fixed 19-major list. This may improve distinctness and resilience to
liquidity shifts.
"""
import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/xsec")
import xslib as xs
import numpy as np

print("XU04 — Liquidity-filtered momentum")
print("=" * 60)


def rolling_liq_score(P, lookbacks=(30, 90), skip=0):
    """Momentum blend on the panel, with optional skip-1 for microstructure."""
    scores = []
    for L in lookbacks:
        if skip > 0:
            # use close[i-skip] / close[i-skip-L] - 1 (causal, skip most recent bars)
            c = P.close
            out = np.full_like(c, np.nan)
            # at row i: return from i-L-skip to i-skip
            for i in range(L + skip, len(c)):
                out[i] = c[i - skip] / c[i - L - skip] - 1.0
        else:
            out = xs.past_return(P.close, L)
        scores.append(xs.xs_zscore(out))
    return np.nanmean(np.stack(scores, axis=0), axis=0)


def score_dyn20_blend(P):
    """Dynamic top-20 by rolling $-vol — blend [30,90], no skip."""
    # P already filtered to top-20 by static median; this fn gets whatever panel is loaded.
    # We do dynamic re-weighting via volume z-score gating:
    # compute rolling 60d dollar volume rank per asset; assets below median get half-weight score
    dv = P.close * P.vol  # dollar volume matrix (n_days x n_assets)
    dv_roll = xs.roll_mean(dv, 60)  # rolling 60d mean $-vol
    # rank liquidity cross-sectionally
    liq_rank = xs.xs_rank(dv_roll)  # 0..1, higher = more liquid
    # momentum signal
    mom = rolling_liq_score(P, lookbacks=(30, 90), skip=0)
    # attenuate score of less-liquid assets (liq_rank < 0.5 -> half score)
    liq_weight = np.where(liq_rank >= 0.5, 1.0, 0.5)
    return mom * liq_weight


def score_skip1(P):
    """Majors, momentum blend [30,90] with 1-day skip (microstructure reduction)."""
    return rolling_liq_score(P, lookbacks=(30, 90), skip=1)


def score_top25_blend(P):
    """Top-25 universe, plain blend [30,90]."""
    return rolling_liq_score(P, lookbacks=(30, 90), skip=0)


def score_dyn20_fast(P):
    """Dynamic top-20 + blend [30,90], faster H=5 rebalance."""
    return score_dyn20_blend(P)


# 1) Top-20 with dynamic liquidity weighting, H=10, k=5
rep1 = xs.study_xs(
    "XU04_DYN20_H10",
    score_dyn20_blend,
    universe=20, H=10, k=5, long_short=True
)
print(xs.fmt(rep1))
print("JSON:", xs.as_json(rep1))
print()

# 2) Majors (19) with skip-1 momentum — reduces microstructure vs XS01
rep2 = xs.study_xs(
    "XU04_MAJ_SKIP1",
    score_skip1,
    universe="majors", H=10, k=5, long_short=True
)
print(xs.fmt(rep2))
print("JSON:", xs.as_json(rep2))
print()

# 3) Top-20 plain blend [30,90] no weighting, H=10 (clean baseline vs XU01)
rep3 = xs.study_xs(
    "XU04_TOP20_PLAIN",
    lambda P: rolling_liq_score(P, lookbacks=(30, 90), skip=0),
    universe=20, H=10, k=5, long_short=True
)
print(xs.fmt(rep3))
print("JSON:", xs.as_json(rep3))
print()

# 4) Top-25, blend [30,90], H=10
rep4 = xs.study_xs(
    "XU04_TOP25_H10",
    score_top25_blend,
    universe=25, H=10, k=5, long_short=True
)
print(xs.fmt(rep4))
print("JSON:", xs.as_json(rep4))
print()

# 5) Top-20 dynamic liq-weighted, H=5 (faster)
rep5 = xs.study_xs(
    "XU04_DYN20_H5",
    score_dyn20_fast,
    universe=20, H=5, k=5, long_short=True
)
print(xs.fmt(rep5))
print("JSON:", xs.as_json(rep5))
print()


# Pick best
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))