#!/usr/bin/env python3
"""The quant-fund method: combine MANY weak cross-sectional signals via a combiner ladder.

Assemble ~13 weak signals on the liquid US-equity universe, then run a ladder of combiners
from dumb to fancy and ask the disciplined question: does COMBINING beat the BEST SINGLE signal
OOS, and does the nonlinear ML beat the LINEAR combiner OOS (or just overfit the 3.2y)?
  best-single -> equal-weight -> IC-weighted -> ridge(IS-fit) -> gradient-boost(IS-fit)
Strict IS/OOS split, realistic illiquidity-scaled cost, weekly rebalance. Reject if the
combination doesn't beat best-single + equal-weight OOS after deflation.
"""
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)
    # ~13 weak cross-sectional signals (each z-scored per day inside the universe)
    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, "skew63": -roll(lambda a, axis: ((a - a.mean(axis))**3).mean(axis) / (a.std(axis)**3 + 1e-9), R, 63),
    }
    names = list(raw)
    sig = np.stack([np.where(univ, v, np.nan) for v in raw.values()], axis=2)  # [T,N,K]
    sigz = np.stack([zc(np.where(univ, v, np.nan)) for v in raw.values()], axis=2)
    K = len(names)
    fwd = np.full((T, N), np.nan); fwd[:-1] = R[1:]                            # next-day return
    split = int(0.65 * T)

    # pooled IS dataset for fitted combiners
    ist = np.zeros(T, bool); ist[:split] = True
    Xall = sigz.reshape(T * N, K)
    yall = fwd.reshape(T * N)
    rowok = np.isfinite(yall) & np.isfinite(Xall).all(1) & np.repeat(univ.reshape(T * N), 1)
    isrow = rowok & np.repeat(ist[:, None], N, 1).reshape(T * N)
    Xis, yis = Xall[isrow], yall[isrow]
    print(f"signals K={K}, universe/day~{int(univ.sum(1).mean())}, IS obs={len(yis):,}")

    ridge = Ridge(alpha=10.0).fit(Xis, yis)
    gb = HistGradientBoostingRegressor(max_depth=3, max_iter=120, learning_rate=0.05,
                                       l2_regularization=1.0, min_samples_leaf=200).fit(Xis, yis)
    comp_ridge = (sigz.reshape(T * N, K) @ ridge.coef_).reshape(T, N)
    Xpred = np.nan_to_num(sigz.reshape(T * N, K))
    comp_gb = gb.predict(Xpred).reshape(T, N)
    # IC-weighted (trailing 60d IC per signal, causal)
    comp_ic = np.zeros((T, N))
    for t in range(60, T):
        w = []
        for k in range(K):
            a = sigz[t - 60:t, :, k].reshape(-1); b = fwd[t - 60:t].reshape(-1)
            m = np.isfinite(a) & np.isfinite(b)
            w.append(np.corrcoef(a[m], b[m])[0, 1] if m.sum() > 200 else 0.0)
        comp_ic[t] = np.nan_to_num(sigz[t] @ np.nan_to_num(np.array(w)))
    comp_eq = np.nansum(sigz, axis=2)

    rt_cost = np.clip(40.0 / np.sqrt(np.maximum(dv30, 1.0) / 1e6), 3.0, 60.0) / 1e4

    def book(comp):
        s = comp.copy(); s[~univ] = np.nan
        w = xs_weights(s)
        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]
        g = np.sum(wh[:-1] * R[1:], axis=1) - np.sum(np.abs(wh[1:] - wh[:-1]) * rt_cost[1:], axis=1)
        return g

    T_ = lambda x: torch.tensor(np.asarray(x)[np.isfinite(np.asarray(x))], device=DEV, dtype=torch.float64)
    NT = K + 5
    print(f"\n===== MULTI-SIGNAL COMBINATION (liquid equities, net cost, deflate N={NT}) =====")
    # best single
    singles = {nm: validate(book(sigz[:, :, k]), days, NT) for k, nm in enumerate(names)}
    bs = max(singles.items(), key=lambda kv: kv[1]["oos"])
    print(f"BEST SINGLE signal: {bs[0]}  OOS {bs[1]['oos']:+.2f}  full {bs[1]['full']:+.2f}  DSR {bs[1]['dsr']:.2f}")
    print(f"{'combiner':>14} {'full':>6} {'IS':>6} {'OOS':>6} {'CPCVmed':>8} {'DSR':>5}")
    for nm, comp in [("equal", comp_eq), ("IC-weighted", comp_ic), ("ridge(IS)", comp_ridge), ("gradboost(IS)", comp_gb)]:
        v = validate(book(comp), days, NT)
        print(f"{nm:>14} {v['full']:>+6.2f} {v['is_']:>+6.2f} {v['oos']:>+6.2f} {v['med']:>+8.2f} {v['dsr']:>5.2f}")
    print("\nVERDICT: combination OOS > best-single OOS => combining works. gradboost OOS > ridge OOS => ML adds;")
    print("gradboost < ridge => ML overfits 3.2y (linear is the real-quant answer). DSR>0.5 = deploy-grade.")


if __name__ == "__main__":
    main()