scripts/train_compare.py

"""Compara LightGBM vs una red neuronal multi-tarea (embeddings + softmax/KL) en el
MISMO split TEMPORAL (test = último ~20% de fechas de cada liga-temporada), sobre los
3 objetivos: bloque defensivo, pasillos (share) y xT absoluto por pasillo.

Uso: python scripts/train_compare.py
"""

from __future__ import annotations

import importlib.util
from pathlib import Path

import lightgbm as lgb
import numpy as np
import pandas as pd
import torch
import torch.nn as nn

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)
torch.manual_seed(7)


def _kl(a, b):
    a = np.clip(a, 1e-9, None); b = np.clip(b, 1e-9, None)
    return float(np.mean(np.sum(a * (np.log(a) - np.log(b)), 1)))


def _metrics(y, p, dist=True):
    mae = float(np.mean(np.abs(y - p)))
    if dist:
        return {"MAE": round(mae, 4), "KL": round(_kl(y, p), 4)}
    return {"MAE": round(mae, 4), "RMSE": round(float(np.sqrt(np.mean((y - p) ** 2))), 4)}


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


class MTNet(nn.Module):
    def __init__(self, n_num, n_team, n_lg, n_form):
        super().__init__()
        self.te = nn.Embedding(n_team, 16); self.le = nn.Embedding(n_lg, 4); self.fe = nn.Embedding(n_form, 4)
        d = n_num + 16 + 16 + 4 + 4
        self.trunk = nn.Sequential(nn.Linear(d, 128), nn.ReLU(), nn.Dropout(0.3),
                                   nn.Linear(128, 64), nn.ReLU(), nn.Dropout(0.2))
        self.h_blo = nn.Linear(64, 3); self.h_pas = nn.Linear(64, 3); self.h_xt = nn.Linear(64, 3)

    def forward(self, x, ts, to, lg, fm):
        z = torch.cat([x, self.te(ts), self.te(to), self.le(lg), self.fe(fm)], 1)
        z = self.trunk(z)
        return self.h_blo(z), self.h_pas(z), self.h_xt(z)


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)

    # filas con historial rolling + bloque válido
    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()

    # targets
    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)
    # baseline = promedio propio rolling
    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
    # val carve del train para early stopping
    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)}")

    # ===== LightGBM =====
    print("\n--- LightGBM (split temporal) ---")
    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])
    lgb_res = {"bloques": _metrics(Yb[te], pb), "pasillos": _metrics(Yp[te], pp), "xt": _metrics(Yx[te], px, dist=False)}

    # ===== NN multi-tarea =====
    print("--- Red neuronal (embeddings + multi-tarea) ---")
    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
    T = lambda a: torch.tensor(a)
    net = MTNet(len(feat), len(tmap) + 1, lg.max() + 1, fm.max() + 1)
    opt = torch.optim.Adam(net.parameters(), lr=2e-3, weight_decay=1e-4)
    klf = nn.KLDivLoss(reduction="batchmean"); msef = nn.MSELoss()
    Xt = T(Xn); tsT, toT, lgT, fmT = T(ts), T(to), T(lg), T(fm)
    Yb_t, Yp_t, Yx_t = T(Yb), T(Yp), T(Yx)
    tr_t, va_t, te_t = T(np.where(tr)[0]), T(np.where(va)[0]), T(np.where(te)[0])
    best = (1e9, None)
    for ep in range(300):
        net.train(); opt.zero_grad()
        b, p, x = net(Xt[tr_t], tsT[tr_t], toT[tr_t], lgT[tr_t], fmT[tr_t])
        loss = (klf(torch.log_softmax(b, 1), Yb_t[tr_t]) + klf(torch.log_softmax(p, 1), Yp_t[tr_t])
                + 0.5 * msef(x, Yx_t[tr_t]))
        loss.backward(); opt.step()
        if ep % 10 == 0:
            net.eval()
            with torch.no_grad():
                b, p, x = net(Xt[va_t], tsT[va_t], toT[va_t], lgT[va_t], fmT[va_t])
                v = (klf(torch.log_softmax(b, 1), Yb_t[va_t]) + klf(torch.log_softmax(p, 1), Yp_t[va_t])).item()
            if v < best[0]:
                best = (v, {k: val.clone() for k, val in net.state_dict().items()})
    net.load_state_dict(best[1]); net.eval()
    with torch.no_grad():
        b, p, x = net(Xt[te_t], tsT[te_t], toT[te_t], lgT[te_t], fmT[te_t])
        nb = torch.softmax(b, 1).numpy(); npp = torch.softmax(p, 1).numpy(); nx = np.clip(x.numpy(), 0, None)
    nn_res = {"bloques": _metrics(Yb[te], nb), "pasillos": _metrics(Yp[te], npp), "xt": _metrics(Yx[te], nx, dist=False)}

    base = {"bloques": _metrics(Yb[te], Bb[te]), "pasillos": _metrics(Yp[te], Bp[te]), "xt": _metrics(Yx[te], Bx[te], dist=False)}
    print("\n================ COMPARACIÓN (test temporal) ================")
    for k in ("bloques", "pasillos", "xt"):
        print(f"\n{k}:")
        print(f"  baseline : {base[k]}")
        print(f"  LightGBM : {lgb_res[k]}")
        print(f"  NN       : {nn_res[k]}")


if __name__ == "__main__":
    main()