"""XS07b — Trend-quality (R^2) ranking.

IDEA: Score each asset by the R^2 of a linear fit of log-price over the last W bars,
      signed by the direction of the trend (positive slope = long candidate).
      Score = sign(slope) * R^2

High R^2 + upward slope -> strong smooth uptrend -> long.
High R^2 + downward slope -> strong smooth downtrend -> short.
Low R^2 -> noisy / not trending -> near-zero score.

W=60 canonical, but we try W=30 and W=90 too. The score is CAUSAL: for row i,
we fit on close[i-W+1 .. i] (inclusive), using only past data.
"""
import sys
sys.path.insert(0, "/opt/docker/PythagorasGoal/scripts/research/xsec")
import xslib as xs
import numpy as np


def r2_trend_score(close, W=60):
    """
    Per-asset, rolling R^2 of linear fit on log(close), signed by slope direction.
    Returns (n_days x n_assets) matrix. Causal: row i uses close[i-W+1..i].
    """
    n, A = close.shape
    out = np.full((n, A), np.nan)
    x = np.arange(W, dtype=float)
    x -= x.mean()   # center x for numerical stability
    xss = (x ** 2).sum()
    for i in range(W - 1, n):
        log_p = np.log(close[i - W + 1: i + 1, :])  # (W, A)
        # For each asset: fit log_p = a + b*x
        # b = cov(x, log_p) / var(x)
        mean_y = log_p.mean(axis=0)           # (A,)
        b = (x[:, None] * (log_p - mean_y)).sum(axis=0) / xss  # (A,)
        y_hat = x[:, None] * b + mean_y       # (W, A)
        ss_res = ((log_p - y_hat) ** 2).sum(axis=0)
        ss_tot = ((log_p - mean_y) ** 2).sum(axis=0)
        r2 = np.where(ss_tot > 0, 1.0 - ss_res / ss_tot, 0.0)
        # Score = sign(slope) * R^2. Ranges in [-1, 1].
        out[i] = np.sign(b) * r2
    return out


def make_score_fn(W=60):
    def score_fn(P):
        return r2_trend_score(P.close, W=W)
    return score_fn


if __name__ == "__main__":
    # Grid: (W, universe, H, k, long_short)
    # Keep <= 5 backtests total
    configs = [
        # Canonical: W=60, all assets, H=10, k=5, LS
        dict(name="XS07b_W60_all_H10_k5_LS",  W=60, universe="all",    H=10, k=5, long_short=True),
        # Shorter trend window
        dict(name="XS07b_W30_all_H10_k5_LS",  W=30, universe="all",    H=10, k=5, long_short=True),
        # Longer trend window
        dict(name="XS07b_W90_all_H10_k5_LS",  W=90, universe="all",    H=10, k=5, long_short=True),
        # Majors only (less noisy universe)
        dict(name="XS07b_W60_maj_H10_k5_LS",  W=60, universe="majors", H=10, k=5, long_short=True),
        # Long-only variant (majors)
        dict(name="XS07b_W60_maj_H10_k3_LO",  W=60, universe="majors", H=10, k=3, long_short=False),
    ]

    best_rep = None
    best_key = (-999, -999, False)  # (earns_slot, hold_sharpe, robust_oos)

    for cfg in configs:
        W = cfg.pop("W")
        name = cfg.pop("name")
        rep = xs.study_xs(name, make_score_fn(W=W), **cfg)
        print(xs.fmt(rep))
        hold_sh = rep["holdout"].get("sharpe", -999)
        earns = int(rep["earns_slot"])
        robust = int(rep["marginal"].get("robust_oos", False))
        key = (earns, robust, hold_sh)
        if key > best_key:
            best_key = key
            best_rep = rep

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