scripts/train_models.py

"""Entrena los dos modelos (bloque defensivo y pasillos) sobre el dataset RICO
matchteam_dataset.parquet. Arma features SIN leakage:
  - ELO propio por (liga,temporada), pre-partido (de los resultados).
  - Rolling "hasta la fecha" (promedio expandido de partidos PREVIOS) del equipo y del
    rival, y splits LOCAL / VISITA (cómo juega/resulta de local vs visitante).
  - Resultados rolling (win%, goles a favor/en contra) all/home/away.
  - One-hot de la formación del partido (info pre-partido).
LightGBM por componente, normalizado a simplex; baseline = promedio propio.

Uso: python scripts/train_models.py
"""

from __future__ import annotations

from pathlib import Path
import sys

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

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

DATA = Path(__file__).resolve().parents[1] / "vendor" / "data" / "modeling" / "matchteam_dataset.parquet"
OUTDIR = DATA.parent
DEF = ["def_H", "def_M", "def_L"]
ATK = ["atk_Der", "atk_Centro", "atk_Izq"]
OFFXT = [f"offxt_{b}_{l}" for b in "HML" for l in ["Der", "Centro", "Izq"]]
DEFXT = [f"defxt_{b}_{l}" for b in "HML" for l in ["Der", "Centro", "Izq"]]
XTLANE = ["xt_lane_Der", "xt_lane_Centro", "xt_lane_Izq"]  # xT absoluto por pasillo (objetivo Modelo 2-b)
DIST_GROUPS = {"def": DEF, "atk": ATK, "offxt": OFFXT, "defxt": DEFXT}
SCALARS = ["sh_buildup", "sh_counter", "sh_direct", "sh_setpiece", "sh_recovery", "sh_progr",
           "cross_pct", "pct_long", "pct_short", "pct_pass_ok", "n_attacks", "xt_total",
           "xg_total", "xt_per_attack", "n_seq", "n_pass", "n_shots", "avg_x", "avg_y",
           "share_opp_half", "recovery_height", "n_progr", "goals_for", "goals_against", "win"] + XTLANE
RNG = np.random.default_rng(7)


def _norm_df(df, cols):
    s = df[cols].sum(axis=1)
    out = df[cols].div(s.where(s > 0), axis=0)
    return out, s


def _elo(df: pd.DataFrame) -> pd.Series:
    """ELO pre-partido por (liga,temporada). Devuelve serie alineada al índice de df."""
    K, HOME = 24.0, 60.0
    elo = pd.Series(np.nan, index=df.index)
    for _, g in df.groupby(["Competencia", "Temporada"], sort=False):
        rating: dict = {}
        # ordenar por fecha y partido; procesar cada partido (2 filas) una vez
        g = g.sort_values(["fecha", "matchId"])
        for mid, mm in g.groupby("matchId", sort=False):
            if len(mm) != 2:
                for i in mm.index:
                    elo.loc[i] = rating.get(mm.loc[i, "teamId"], 1500.0)
                continue
            i1, i2 = mm.index
            t1, t2 = mm.loc[i1, "teamId"], mm.loc[i2, "teamId"]
            r1, r2 = rating.get(t1, 1500.0), rating.get(t2, 1500.0)
            elo.loc[i1], elo.loc[i2] = r1, r2  # pre-partido
            h1 = HOME if mm.loc[i1, "is_home"] else 0.0
            h2 = HOME if mm.loc[i2, "is_home"] else 0.0
            e1 = 1.0 / (1.0 + 10 ** ((r2 - r1 - h1 + h2) / 400.0))
            gf1, ga1 = mm.loc[i1, "goals_for"], mm.loc[i1, "goals_against"]
            s1 = 0.5 if pd.isna(gf1) or pd.isna(ga1) or gf1 == ga1 else (1.0 if gf1 > ga1 else 0.0)
            gd = 1.0 if pd.isna(gf1) or pd.isna(ga1) else max(1.0, abs(gf1 - ga1)) ** 0.5
            rating[t1] = r1 + K * gd * (s1 - e1)
            rating[t2] = r2 + K * gd * ((1 - s1) - (1 - e1))
    return elo


def _prep(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df["win"] = np.where(df["goals_for"] > df["goals_against"], 1.0,
                         np.where(df["goals_for"] < df["goals_against"], 0.0, 0.5))
    # xT absoluto por pasillo = suma del xT por bloque en cada pasillo
    for ln in ["Der", "Centro", "Izq"]:
        df[f"xt_lane_{ln}"] = df[[f"offxt_{b}_{ln}" for b in "HML"]].sum(axis=1)
    # normalizar grupos distribucionales a shares (para features de estilo/peligro)
    feat = {}
    for name, cols in DIST_GROUPS.items():
        nd, _ = _norm_df(df, cols)
        for c in cols:
            feat["f_" + c] = nd[c]
    for c in SCALARS:
        feat["f_" + c] = pd.to_numeric(df[c], errors="coerce")
    F = pd.DataFrame(feat, index=df.index)
    df = pd.concat([df, F], axis=1)
    df["elo"] = _elo(df)
    return df


def _rolling(df: pd.DataFrame) -> pd.DataFrame:
    fcols = [c for c in df.columns if c.startswith("f_")]
    df = df.sort_values(["Competencia", "Temporada", "teamId", "fecha", "matchId"]).copy()
    keys = ["Competencia", "Temporada", "teamId"]

    def _prior(frame):  # expanding mean de partidos previos
        return frame.groupby([df[k] for k in keys], group_keys=False).apply(lambda x: x.shift(1).expanding().mean())

    roll_all = _prior(df[fcols]); roll_all.columns = ["self_" + c[2:] for c in fcols]
    # splits local/visita: expanding sobre el subconjunto, luego ffill al resto
    def _split(is_home_val):
        sub = df[fcols].where(df["is_home"] == is_home_val)
        r = sub.groupby([df[k] for k in keys], group_keys=False).apply(lambda x: x.shift(1).expanding().mean())
        r = r.groupby([df[k] for k in keys], group_keys=False).ffill()
        return r
    rh = _split(True); rh.columns = ["selfH_" + c[2:] for c in fcols]
    ra = _split(False); ra.columns = ["selfA_" + c[2:] for c in fcols]
    eloprior = df.groupby(keys, group_keys=False)["elo"].apply(lambda x: x)  # ya es pre-partido
    out = pd.concat([df[keys + ["matchId", "rival_teamId", "is_home", "formation"] + DEF + ATK + XTLANE],
                     df["elo"].rename("self_elo"), roll_all, rh, ra], axis=1)
    return out


def _assemble(df: pd.DataFrame):
    df = _prep(df)
    roll = _rolling(df)
    selfcols = [c for c in roll.columns if c.startswith(("self_", "selfH_", "selfA_"))]
    oppmap = roll[["matchId", "teamId"] + selfcols].rename(
        columns={"teamId": "rival_teamId", **{c: "opp_" + c for c in selfcols}})
    m = roll.merge(oppmap, on=["matchId", "rival_teamId"], how="left")
    m["is_home"] = m["is_home"].astype(float)
    m["elo_diff"] = m["self_elo"] - m["opp_self_elo"]
    # one-hot formación (top 12 + otras)
    top = m["formation"].astype(str).value_counts().head(12).index
    for f in top:
        m[f"form_{f}"] = (m["formation"].astype(str) == f).astype(float)
    feat_cols = (selfcols + [c for c in m.columns if c.startswith("opp_")]
                 + ["is_home", "elo_diff"] + [f"form_{f}" for f in top])
    return m, feat_cols


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 _train_one(m, feat_cols, target_cols, name):
    sub = m.dropna(subset=["self_" + c for c in target_cols]).copy()  # con historial rolling
    tgt, tot = _norm_df(sub, target_cols)
    keep = tot.values > 0
    sub = sub[keep]; Y = tgt[keep].to_numpy(dtype=np.float32)
    base = sub[["self_" + c for c in target_cols]].to_numpy(dtype=np.float32)
    base = base / np.clip(base.sum(1, keepdims=True), 1e-9, None)
    n = len(sub); idx = RNG.permutation(n); cut = int(n * 0.8); tr, te = idx[:cut], idx[cut:]
    Xtr, Xte = sub[feat_cols].iloc[tr], sub[feat_cols].iloc[te]
    vm = np.zeros(len(tr), bool); vm[RNG.permutation(len(tr))[:int(len(tr) * 0.12)]] = True
    print(f"\n=== Modelo {name}: {n} (test {len(te)}) | {len(feat_cols)} features ===")
    preds = []
    for j in range(Y.shape[1]):
        mdl = 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)
        mdl.fit(Xtr[~vm], Y[tr][~vm, j], eval_set=[(Xtr[vm], Y[tr][vm, j])],
                callbacks=[lgb.early_stopping(50, verbose=False)])
        preds.append(np.clip(mdl.predict(Xte), 0, None))
        mdl.booster_.save_model(str(OUTDIR / f"model_{name}__{target_cols[j]}.txt"))
    P = np.vstack(preds).T; P = P / np.clip(P.sum(1, keepdims=True), 1e-9, None)
    mae = lambda a, b: float(np.mean(np.abs(a - b)))
    print(f"  modelo  : MAE={mae(Y[te], P):.4f}  KL={_kl(Y[te], P):.4f}")
    print(f"  baseline: MAE={mae(Y[te], base[te]):.4f}  KL={_kl(Y[te], base[te]):.4f}")


def _train_abs(m, feat_cols, target_cols, name):
    """Regresión de VALORES ABSOLUTOS (no shares). Eval MAE/RMSE vs promedio propio."""
    sub = m.dropna(subset=["self_" + c for c in target_cols]).copy()
    Y = sub[target_cols].to_numpy(dtype=np.float32)
    base = sub[["self_" + c for c in target_cols]].to_numpy(dtype=np.float32)
    n = len(sub); idx = RNG.permutation(n); cut = int(n * 0.8); tr, te = idx[:cut], idx[cut:]
    Xtr, Xte = sub[feat_cols].iloc[tr], sub[feat_cols].iloc[te]
    vm = np.zeros(len(tr), bool); vm[RNG.permutation(len(tr))[:int(len(tr) * 0.12)]] = True
    print(f"\n=== Modelo {name} (xT absoluto): {n} (test {len(te)}) ===")
    preds = []
    for j in range(Y.shape[1]):
        mdl = 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)
        mdl.fit(Xtr[~vm], Y[tr][~vm, j], eval_set=[(Xtr[vm], Y[tr][vm, j])],
                callbacks=[lgb.early_stopping(50, verbose=False)])
        preds.append(np.clip(mdl.predict(Xte), 0, None))
        mdl.booster_.save_model(str(OUTDIR / f"model_{name}__{target_cols[j]}.txt"))
    P = np.vstack(preds).T
    mae = lambda a, b: float(np.mean(np.abs(a - b)))
    rmse = lambda a, b: float(np.sqrt(np.mean((a - b) ** 2)))
    print(f"  modelo  : MAE={mae(Y[te], P):.4f}  RMSE={rmse(Y[te], P):.4f}")
    print(f"  baseline: MAE={mae(Y[te], base[te]):.4f}  RMSE={rmse(Y[te], base[te]):.4f}")
    print(f"  (xT medio real por pasillo: {np.round(Y[te].mean(0), 3)})")


def main():
    import argparse
    ap = argparse.ArgumentParser(); ap.add_argument("--leagues", nargs="*", default=None); a = ap.parse_args()
    df = pd.read_parquet(DATA)
    if a.leagues:
        df = df[df["Competencia"].isin(a.leagues)].copy()
    blk = (df.assign(_b=df[DEF].sum(axis=1) > 0).groupby("Competencia")["_b"].mean())
    df = df[df["Competencia"].isin(blk[blk > 0.5].index)].copy()
    print("dataset:", len(df), "filas |", df["Competencia"].nunique(), "ligas |", df["matchId"].nunique(), "partidos")
    m, feat = _assemble(df)
    _train_one(m, feat, DEF, "bloques")
    _train_one(m, feat, ATK, "pasillos")
    _train_abs(m, feat, XTLANE, "peligro_pasillo")


if __name__ == "__main__":
    main()