scripts/cmp_prep.py

"""Etapa 1 de la comparación: arma el split TEMPORAL, entrena LightGBM y guarda
TODO lo necesario (arrays + predicciones LGB + baseline) en /tmp/cmp.npz para que
la red neuronal corra en un proceso APARTE (evita el deadlock/segfault de OpenMP
entre lightgbm y torch en el mismo proceso).

Uso: python scripts/cmp_prep.py
"""

from __future__ import annotations

import importlib.util
from pathlib import Path

import lightgbm as lgb
import numpy as np
import pandas as pd

ROOT = Path(__file__).resolve().parents[1]
spec = importlib.util.spec_from_file_location("tm", ROOT / "scripts" / "train_models.py")
tm = importlib.util.module_from_spec(spec); spec.loader.exec_module(tm)
RNG = np.random.default_rng(7)
NPZ = Path("/tmp/cmp.npz")


def _lgb_dist(Xtr, Ytr, Xva, Yva, Xte):
    preds = []
    for j in range(Ytr.shape[1]):
        m = lgb.LGBMRegressor(n_estimators=700, learning_rate=0.025, num_leaves=31,
                              subsample=0.8, colsample_bytree=0.5, min_child_samples=40, verbose=-1)
        m.fit(Xtr, Ytr[:, j], eval_set=[(Xva, Yva[:, j])], callbacks=[lgb.early_stopping(50, verbose=False)])
        preds.append(np.clip(m.predict(Xte), 0, None))
    return np.vstack(preds).T


def main():
    df = pd.read_parquet(tm.DATA)
    blk = df.assign(_b=df[tm.DEF].sum(axis=1) > 0).groupby("Competencia")["_b"].mean()
    df = df[df["Competencia"].isin(blk[blk > 0.5].index)].copy()
    m, feat = tm._assemble(df)
    m = m.merge(df[["matchId", "teamId", "fecha"]].drop_duplicates(), on=["matchId", "teamId"], how="left")
    m["fecha"] = pd.to_datetime(m["fecha"].astype(str).str[:10], errors="coerce")  # quita sufijo 'Z'
    thr = m.groupby(["Competencia", "Temporada"])["fecha"].transform(lambda s: s.quantile(0.8))
    m["is_test"] = (m["fecha"] > thr).fillna(False)

    m = m.dropna(subset=["self_" + c for c in tm.DEF]).copy()
    valid = (m[tm.DEF].sum(1) > 0) & (m[tm.ATK].sum(1) > 0)
    m = m[valid].copy()

    Yb, _ = tm._norm_df(m, tm.DEF); Yb = Yb.to_numpy(np.float32)
    Yp, _ = tm._norm_df(m, tm.ATK); Yp = Yp.to_numpy(np.float32)
    Yx = m[tm.XTLANE].to_numpy(np.float32)
    Bb = m[["self_" + c for c in tm.DEF]].to_numpy(np.float32); Bb = Bb / np.clip(Bb.sum(1, keepdims=True), 1e-9, None)
    Bp = m[["self_" + c for c in tm.ATK]].to_numpy(np.float32); Bp = Bp / np.clip(Bp.sum(1, keepdims=True), 1e-9, None)
    Bx = m[["self_" + c for c in tm.XTLANE]].to_numpy(np.float32)

    X = m[feat].to_numpy(np.float32); X = np.nan_to_num(X)
    te = m["is_test"].to_numpy()
    tr_all = ~te
    mu, sd = X[tr_all].mean(0), X[tr_all].std(0) + 1e-6
    Xn = ((X - mu) / sd).astype(np.float32)
    tr_idx = np.where(tr_all)[0]
    va = np.zeros(len(m), bool); va[RNG.choice(tr_idx, int(len(tr_idx) * 0.12), replace=False)] = True
    tr = tr_all & ~va
    print(f"filas: {len(m)} | train {tr.sum()} val {va.sum()} test {te.sum()} | features {len(feat)}", flush=True)

    print("--- LightGBM (split temporal) ---", flush=True)
    pb = _lgb_dist(Xn[tr], Yb[tr], Xn[va], Yb[va], Xn[te]); pb = pb / np.clip(pb.sum(1, keepdims=True), 1e-9, None)
    pp = _lgb_dist(Xn[tr], Yp[tr], Xn[va], Yp[va], Xn[te]); pp = pp / np.clip(pp.sum(1, keepdims=True), 1e-9, None)
    px = _lgb_dist(Xn[tr], Yx[tr], Xn[va], Yx[va], Xn[te])
    print("LightGBM listo, guardando arrays...", flush=True)

    # embeddings para la red
    tcodes = pd.factorize(pd.concat([m["teamId"], m["rival_teamId"]]).astype(str))[1]
    tmap = {v: i + 1 for i, v in enumerate(tcodes)}
    ts = m["teamId"].astype(str).map(tmap).fillna(0).astype(int).to_numpy()
    to = m["rival_teamId"].astype(str).map(tmap).fillna(0).astype(int).to_numpy()
    lg = pd.factorize(m["Competencia"])[0] + 1
    fm = pd.factorize(m["formation"].astype(str))[0] + 1

    np.savez(NPZ, Xn=Xn, Yb=Yb, Yp=Yp, Yx=Yx, Bb=Bb, Bp=Bp, Bx=Bx,
             tr=tr, va=va, te=te, pb=pb, pp=pp, px=px,
             ts=ts, to=to, lg=lg, fm=fm, nfeat=len(feat),
             n_team=len(tmap) + 1, n_lg=int(lg.max()) + 1, n_form=int(fm.max()) + 1)
    print("guardado →", NPZ, flush=True)


if __name__ == "__main__":
    main()