#!/usr/bin/env python3
"""Leak-free WALK-FORWARD validation of the multi-signal combiner — the decisive test.

The single IS/OOS gradboost result (OOS +1.14, DSR 0.62) is overfit-suspect (one short OOS
window, high-capacity model, only fitted combiners work). Real test: re-fit the combiner on an
EXPANDING window and predict ONLY the next block forward (never peeking). Concatenate the
out-of-sample predictions, build the book on the OOS period only, and compare gradboost vs ridge
vs equal vs best-single — all leak-free. If gradboost still wins OOS, it's real; if it collapses,
it was memorization.
"""
import math
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from equity_factor_gate import load, roll, trailing  # noqa: E402
from signal_sweep import xs_weights, validate, sharpe_t  # noqa: E402
from sklearn.linear_model import Ridge  # noqa: E402
from sklearn.ensemble import HistGradientBoostingRegressor  # noqa: E402
import torch  # noqa: E402

DEV = "cuda" if torch.cuda.is_available() else "cpu"
DV_FLOOR, TOPK = 5e6, 1000


def zc(x):
    mu = np.nanmean(x, axis=1, keepdims=True); sd = np.nanstd(x, axis=1, keepdims=True)
    return np.nan_to_num((x - mu) / np.where(sd > 0, sd, 1))


def main():
    insts, days, close, dvol = load()
    T, N = close.shape
    lc = np.log(close)
    R = np.zeros((T, N)); R[1:] = lc[1:] - lc[:-1]; R = np.where(np.isfinite(R), R, 0.0)
    dv30 = roll(np.mean, np.nan_to_num(dvol), 30)
    vol63 = roll(np.std, R, 63)
    year = (1970 + days / 365.25).astype(int)
    univ = np.zeros((T, N), bool)
    for t in range(T):
        e = np.where((dv30[t] > DV_FLOOR) & np.isfinite(close[t]))[0]
        if len(e) > 50:
            univ[t, e[np.argsort(-dv30[t, e])[:TOPK]]] = True
    rmax252 = roll(np.max, lc, 252)
    amih = roll(np.mean, np.abs(R) / np.maximum(np.nan_to_num(dvol), 1.0), 21)
    raw = {"mom_21": trailing(lc, 21), "mom_63": trailing(lc, 63), "mom_126": trailing(lc, 126),
           "mom_252": trailing(lc, 252), "rev_5": -trailing(lc, 5), "rev_10": -trailing(lc, 10),
           "lowvol": -vol63, "amihud": amih, "max20": -roll(np.max, R, 20),
           "accel": trailing(lc, 10) - trailing(lc, 63), "hi52": lc - rmax252,
           "vol_chg": roll(np.std, R, 21) - vol63}
    names = list(raw); K = len(names)
    sigz = np.stack([zc(np.where(univ, v, np.nan)) for v in raw.values()], axis=2)  # [T,N,K] causal
    fwd = np.full((T, N), np.nan); fwd[:-1] = R[1:]
    rt_cost = np.clip(40.0 / np.sqrt(np.maximum(dv30, 1.0) / 1e6), 3.0, 60.0) / 1e4

    INIT, STEP = int(0.45 * T), 42                                  # ~1.4y initial, refit every ~2mo
    comp_gb = np.full((T, N), np.nan); comp_ri = np.full((T, N), np.nan)
    Xall = sigz.reshape(T * N, K); yall = fwd.reshape(T * N); uflat = univ.reshape(T * N)
    for s in range(INIT, T, STEP):
        e = min(s + STEP, T)
        tr = np.zeros(T, bool); tr[:s] = True
        rows = np.repeat(tr[:, None], N, 1).reshape(T * N) & uflat & np.isfinite(yall) & np.isfinite(Xall).all(1)
        Xtr, ytr = Xall[rows], yall[rows]
        if len(ytr) < 5000:
            continue
        ri = Ridge(alpha=10.0).fit(Xtr, ytr)
        gb = HistGradientBoostingRegressor(max_depth=3, max_iter=120, learning_rate=0.05,
                                           l2_regularization=1.0, min_samples_leaf=200).fit(Xtr, ytr)
        blk = sigz[s:e].reshape((e - s) * N, K)
        comp_ri[s:e] = (np.nan_to_num(blk) @ ri.coef_).reshape(e - s, N)
        comp_gb[s:e] = gb.predict(np.nan_to_num(blk)).reshape(e - s, N)

    def book(comp):
        c = comp.copy(); c[~univ] = np.nan
        w = xs_weights(c)
        a = 2.0 / 6
        for t in range(1, T):
            w[t] = a * w[t] + (1 - a) * w[t - 1]
        wh = w.copy(); last = 0
        for t in range(T):
            if t % 5 == 0:
                last = t
            wh[t] = w[last]
        return np.sum(wh[:-1] * R[1:], axis=1) - np.sum(np.abs(wh[1:] - wh[:-1]) * rt_cost[1:], axis=1)

    # OOS period = [INIT:] only
    oos = np.zeros(T - 1, bool); oos[INIT:] = True
    T_ = lambda x: torch.tensor(np.asarray(x)[INIT:][np.isfinite(np.asarray(x)[INIT:])], device=DEV, dtype=torch.float64)
    yr = year[1:]
    eqp = book(np.nansum(sigz, axis=2))
    best = max(range(K), key=lambda k: sharpe_t(T_(book(sigz[:, :, k]))))
    print(f"\n===== WALK-FORWARD (leak-free) MULTI-SIGNAL COMBINE — OOS only ({int(oos.sum())} days) =====")
    print(f"K={K} signals, refit every {STEP}d on expanding window")
    def line(nm, p):
        po = np.asarray(p)
        v = validate(po[INIT:], days, K + 4)
        print(f"{nm:>18} OOS Sharpe {sharpe_t(T_(p)):+.2f}  CPCVmed {v['med']:+.2f}  DSR {v['dsr']:.2f}")
    line(f"best-single({names[best]})", book(sigz[:, :, best]))
    line("equal-weight", eqp)
    line("ridge WALK-FWD", comp_ri)
    line("gradboost WALK-FWD", comp_gb)
    print("\nVERDICT: gradboost WF OOS > ridge WF and > best-single => the ML combination is REAL (leak-free).")
    print("If gradboost WF collapses to ~ridge or below => the single-split +1.14 was memorization.")


if __name__ == "__main__":
    main()