vendor/scripts/experiment_team_embedding_gnn.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
import html
import json
import random
import sys

import numpy as np
import pandas as pd
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset


SCRIPT_DIR = Path(__file__).resolve().parent
PROJECT_ROOT = SCRIPT_DIR.parent
MODEL_DIR = PROJECT_ROOT / "data" / "modeling"
REPORTS_DIR = PROJECT_ROOT / "reports"
REPORT_PATH = REPORTS_DIR / "team_embedding_gnn_report.html"
JSON_PATH = MODEL_DIR / "team_embedding_gnn_metrics.json"
MODEL_PATH = MODEL_DIR / "team_embedding_gnn_bundle.pt"
ATTACK_PRED_PATH = MODEL_DIR / "attack_team_embedding_gnn_test_predictions.parquet"
PV_PRED_PATH = MODEL_DIR / "pv_team_embedding_gnn_test_predictions.parquet"

if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

import experiment_attack_distribution_gnn as attack_base  # noqa: E402
import experiment_pv_distribution_gnn as pv_base  # noqa: E402
import train_attack_prediction_ffn as base  # noqa: E402


RANDOM_SEED = 42
EMBED_DIM = 16
ATTACK_LR = 4e-4
PV_LR = 4e-4
WEIGHT_DECAY = 1e-5
MAX_EPOCHS = 260
PATIENCE = 32
BATCH_SIZE = 256


@dataclass
class AttackSplit:
    global_x: np.ndarray
    node_x: np.ndarray
    y: np.ndarray
    baseline_long: np.ndarray
    baseline_short: np.ndarray
    team_idx: np.ndarray
    opp_idx: np.ndarray
    metadata: pd.DataFrame


@dataclass
class PVSplit:
    global_x: np.ndarray
    node_x: np.ndarray
    y: np.ndarray
    baseline_long: np.ndarray
    baseline_short: np.ndarray
    team_idx: np.ndarray
    opp_idx: np.ndarray
    metadata: pd.DataFrame


def _set_seed(seed: int = RANDOM_SEED) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def _build_team_index(df: pd.DataFrame) -> tuple[pd.Series, pd.Series, dict[str, int]]:
    ids = pd.concat(
        [
            df["teamId"].astype(str),
            df["opponent_team_id"].astype(str),
        ],
        ignore_index=True,
    ).dropna().drop_duplicates().sort_values()
    mapping = {team_id: idx for idx, team_id in enumerate(ids.tolist())}
    team_idx = df["teamId"].astype(str).map(mapping).astype(int)
    opp_idx = df["opponent_team_id"].astype(str).map(mapping).astype(int)
    return team_idx, opp_idx, mapping


class TeamEmbeddingAttackGNN(nn.Module):
    def __init__(self, node_dim: int, global_dim: int, n_teams: int, hidden_dim: int = 96, global_hidden: int = 96) -> None:
        super().__init__()
        self.team_emb = nn.Embedding(n_teams, EMBED_DIM)
        self.opp_emb = nn.Embedding(n_teams, EMBED_DIM)
        self.global_encoder = nn.Sequential(
            nn.Linear(global_dim + (EMBED_DIM * 2), 192),
            nn.ReLU(),
            nn.Dropout(0.10),
            nn.Linear(192, global_hidden),
            nn.ReLU(),
        )
        self.node_encoder = nn.Sequential(
            nn.Linear(node_dim + global_hidden, hidden_dim),
            nn.ReLU(),
        )
        self.graph_blocks = nn.ModuleList([attack_base.GraphBlock(hidden_dim, hidden_dim, dropout=0.06) for _ in range(3)])
        self.gate_head = nn.Linear(hidden_dim, 1)
        self.delta_head = nn.Linear(hidden_dim, 1)

    def forward(
        self,
        node_x: torch.Tensor,
        global_x: torch.Tensor,
        baseline_long: torch.Tensor,
        baseline_short: torch.Tensor,
        team_idx: torch.Tensor,
        opp_idx: torch.Tensor,
        adj: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        team_emb = self.team_emb(team_idx)
        opp_emb = self.opp_emb(opp_idx)
        g = self.global_encoder(torch.cat([global_x, team_emb, opp_emb], dim=1))
        g_rep = g.unsqueeze(1).expand(-1, node_x.size(1), -1)
        h = self.node_encoder(torch.cat([node_x, g_rep], dim=-1))
        gate = torch.sigmoid(self.gate_head(h)).squeeze(-1)
        mixed = gate * baseline_short + (1.0 - gate) * baseline_long
        for block in self.graph_blocks:
            h = h + block(h, adj)
        delta = self.delta_head(h).squeeze(-1)
        logits = torch.log(torch.clamp(mixed, min=1e-6)) + delta
        pred = torch.softmax(logits, dim=1)
        return pred, delta, gate, mixed


class TeamEmbeddingPVGNN(nn.Module):
    def __init__(self, node_dim: int, global_dim: int, n_teams: int, hidden_dim: int = 96, global_hidden: int = 96) -> None:
        super().__init__()
        self.team_emb = nn.Embedding(n_teams, EMBED_DIM)
        self.opp_emb = nn.Embedding(n_teams, EMBED_DIM)
        self.global_encoder = nn.Sequential(
            nn.Linear(global_dim + (EMBED_DIM * 2), 192),
            nn.ReLU(),
            nn.Dropout(0.10),
            nn.Linear(192, global_hidden),
            nn.ReLU(),
        )
        self.node_encoder = nn.Sequential(
            nn.Linear(node_dim + global_hidden, hidden_dim),
            nn.ReLU(),
        )
        self.graph_blocks = nn.ModuleList([pv_base.GraphBlock(hidden_dim, hidden_dim, dropout=0.06) for _ in range(3)])
        self.gate_head = nn.Linear(hidden_dim, 1)
        self.delta_head = nn.Linear(hidden_dim, 1)

    def forward(
        self,
        node_x: torch.Tensor,
        global_x: torch.Tensor,
        baseline_long: torch.Tensor,
        baseline_short: torch.Tensor,
        team_idx: torch.Tensor,
        opp_idx: torch.Tensor,
        adj: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        team_emb = self.team_emb(team_idx)
        opp_emb = self.opp_emb(opp_idx)
        g = self.global_encoder(torch.cat([global_x, team_emb, opp_emb], dim=1))
        g_rep = g.unsqueeze(1).expand(-1, node_x.size(1), -1)
        h = self.node_encoder(torch.cat([node_x, g_rep], dim=-1))
        gate = torch.sigmoid(self.gate_head(h)).squeeze(-1)
        mixed = gate * baseline_short + (1.0 - gate) * baseline_long
        for block in self.graph_blocks:
            h = h + block(h, adj)
        delta = self.delta_head(h).squeeze(-1)
        logits = torch.log(torch.clamp(mixed, min=1e-6)) + delta
        pred = torch.softmax(logits, dim=1)
        return pred, delta, gate, mixed


def _make_attack_split(
    df: pd.DataFrame,
    idx: np.ndarray,
    global_x: np.ndarray,
    node_x: np.ndarray,
    y: np.ndarray,
    baseline_long: np.ndarray,
    baseline_short: np.ndarray,
    team_idx: np.ndarray,
    opp_idx: np.ndarray,
) -> AttackSplit:
    return AttackSplit(
        global_x=global_x[idx].astype(np.float32),
        node_x=node_x[idx].astype(np.float32),
        y=y[idx].astype(np.float32),
        baseline_long=baseline_long[idx].astype(np.float32),
        baseline_short=baseline_short[idx].astype(np.float32),
        team_idx=team_idx[idx].astype(np.int64),
        opp_idx=opp_idx[idx].astype(np.int64),
        metadata=df.iloc[idx].copy().reset_index(drop=True),
    )


def _make_pv_split(
    df: pd.DataFrame,
    idx: np.ndarray,
    global_x: np.ndarray,
    node_x: np.ndarray,
    y: np.ndarray,
    baseline_long: np.ndarray,
    baseline_short: np.ndarray,
    team_idx: np.ndarray,
    opp_idx: np.ndarray,
) -> PVSplit:
    return PVSplit(
        global_x=global_x[idx].astype(np.float32),
        node_x=node_x[idx].astype(np.float32),
        y=y[idx].astype(np.float32),
        baseline_long=baseline_long[idx].astype(np.float32),
        baseline_short=baseline_short[idx].astype(np.float32),
        team_idx=team_idx[idx].astype(np.int64),
        opp_idx=opp_idx[idx].astype(np.int64),
        metadata=df.iloc[idx].copy().reset_index(drop=True),
    )


def _make_loader(split: AttackSplit | PVSplit, shuffle: bool) -> DataLoader:
    dataset = TensorDataset(
        torch.from_numpy(split.global_x),
        torch.from_numpy(split.node_x),
        torch.from_numpy(split.y),
        torch.from_numpy(split.baseline_long),
        torch.from_numpy(split.baseline_short),
        torch.from_numpy(split.team_idx),
        torch.from_numpy(split.opp_idx),
    )
    return DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=shuffle)


def _train_attack(train: AttackSplit, val: AttackSplit, n_teams: int) -> tuple[TeamEmbeddingAttackGNN, list[dict[str, float]]]:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = TeamEmbeddingAttackGNN(train.node_x.shape[2], train.global_x.shape[1], n_teams=n_teams).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=ATTACK_LR, weight_decay=WEIGHT_DECAY)
    train_loader = _make_loader(train, shuffle=True)
    val_loader = _make_loader(val, shuffle=False)
    best_state = None
    best_val = float("inf")
    patience_left = PATIENCE
    history: list[dict[str, float]] = []
    adj = attack_base.ADJ_MATRIX.to(device)

    for epoch in range(1, MAX_EPOCHS + 1):
        model.train()
        running = 0.0
        n_batches = 0
        for global_x, node_x, y, baseline_long, baseline_short, team_idx, opp_idx in train_loader:
            global_x = global_x.to(device)
            node_x = node_x.to(device)
            y = y.to(device)
            baseline_long = baseline_long.to(device)
            baseline_short = baseline_short.to(device)
            team_idx = team_idx.to(device)
            opp_idx = opp_idx.to(device)
            optimizer.zero_grad()
            pred, delta, gate, _mixed = model(node_x, global_x, baseline_long, baseline_short, team_idx, opp_idx, adj)
            loss, _parts = attack_base._distribution_loss(pred, y, delta, gate)
            loss.backward()
            optimizer.step()
            running += float(loss.item())
            n_batches += 1

        model.eval()
        total_val = 0.0
        n_val = 0
        with torch.no_grad():
            for global_x, node_x, y, baseline_long, baseline_short, team_idx, opp_idx in val_loader:
                global_x = global_x.to(device)
                node_x = node_x.to(device)
                y = y.to(device)
                baseline_long = baseline_long.to(device)
                baseline_short = baseline_short.to(device)
                team_idx = team_idx.to(device)
                opp_idx = opp_idx.to(device)
                pred, delta, gate, _mixed = model(node_x, global_x, baseline_long, baseline_short, team_idx, opp_idx, adj)
                loss, parts = attack_base._distribution_loss(pred, y, delta, gate)
                total_val += float(loss.item())
                val_kl = parts["kl"]
                n_val += 1
        val_loss = total_val / max(n_val, 1)
        history.append({"epoch": epoch, "train_loss": running / max(n_batches, 1), "val_loss": val_loss, "val_kl": val_kl})
        if val_loss < best_val - 1e-6:
            best_val = val_loss
            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
            patience_left = PATIENCE
        else:
            patience_left -= 1
        if patience_left <= 0:
            break

    if best_state is None:
        best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
    model.load_state_dict(best_state)
    return model, history


def _train_pv(train: PVSplit, val: PVSplit, n_teams: int) -> tuple[TeamEmbeddingPVGNN, list[dict[str, float]]]:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = TeamEmbeddingPVGNN(train.node_x.shape[2], train.global_x.shape[1], n_teams=n_teams).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=PV_LR, weight_decay=WEIGHT_DECAY)
    train_loader = _make_loader(train, shuffle=True)
    val_loader = _make_loader(val, shuffle=False)
    best_state = None
    best_val = float("inf")
    patience_left = PATIENCE
    history: list[dict[str, float]] = []
    adj = pv_base.ADJ_MATRIX.to(device)

    for epoch in range(1, MAX_EPOCHS + 1):
        model.train()
        running = 0.0
        n_batches = 0
        for global_x, node_x, y, baseline_long, baseline_short, team_idx, opp_idx in train_loader:
            global_x = global_x.to(device)
            node_x = node_x.to(device)
            y = y.to(device)
            baseline_long = baseline_long.to(device)
            baseline_short = baseline_short.to(device)
            team_idx = team_idx.to(device)
            opp_idx = opp_idx.to(device)
            optimizer.zero_grad()
            pred, delta, gate, _mixed = model(node_x, global_x, baseline_long, baseline_short, team_idx, opp_idx, adj)
            loss, _parts = pv_base._dist_loss(pred, y, delta, gate)
            loss.backward()
            optimizer.step()
            running += float(loss.item())
            n_batches += 1

        model.eval()
        total_val = 0.0
        n_val = 0
        with torch.no_grad():
            for global_x, node_x, y, baseline_long, baseline_short, team_idx, opp_idx in val_loader:
                global_x = global_x.to(device)
                node_x = node_x.to(device)
                y = y.to(device)
                baseline_long = baseline_long.to(device)
                baseline_short = baseline_short.to(device)
                team_idx = team_idx.to(device)
                opp_idx = opp_idx.to(device)
                pred, delta, gate, _mixed = model(node_x, global_x, baseline_long, baseline_short, team_idx, opp_idx, adj)
                loss, parts = pv_base._dist_loss(pred, y, delta, gate)
                total_val += float(loss.item())
                val_kl = parts["kl"]
                n_val += 1
        val_loss = total_val / max(n_val, 1)
        history.append({"epoch": epoch, "train_loss": running / max(n_batches, 1), "val_loss": val_loss, "val_kl": val_kl})
        if val_loss < best_val - 1e-6:
            best_val = val_loss
            best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
            patience_left = PATIENCE
        else:
            patience_left -= 1
        if patience_left <= 0:
            break

    if best_state is None:
        best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
    model.load_state_dict(best_state)
    return model, history


def _predict_attack(model: TeamEmbeddingAttackGNN, split: AttackSplit) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.eval()
    preds, gates, mixeds = [], [], []
    loader = _make_loader(split, shuffle=False)
    adj = attack_base.ADJ_MATRIX.to(device)
    with torch.no_grad():
        for global_x, node_x, _y, baseline_long, baseline_short, team_idx, opp_idx in loader:
            pred, _delta, gate, mixed = model(
                node_x.to(device),
                global_x.to(device),
                baseline_long.to(device),
                baseline_short.to(device),
                team_idx.to(device),
                opp_idx.to(device),
                adj,
            )
            preds.append(pred.cpu().numpy())
            gates.append(gate.cpu().numpy())
            mixeds.append(mixed.cpu().numpy())
    return np.vstack(preds), np.vstack(gates), np.vstack(mixeds)


def _predict_pv(model: TeamEmbeddingPVGNN, split: PVSplit) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.eval()
    preds, gates, mixeds = [], [], []
    loader = _make_loader(split, shuffle=False)
    adj = pv_base.ADJ_MATRIX.to(device)
    with torch.no_grad():
        for global_x, node_x, _y, baseline_long, baseline_short, team_idx, opp_idx in loader:
            pred, _delta, gate, mixed = model(
                node_x.to(device),
                global_x.to(device),
                baseline_long.to(device),
                baseline_short.to(device),
                team_idx.to(device),
                opp_idx.to(device),
                adj,
            )
            preds.append(pred.cpu().numpy())
            gates.append(gate.cpu().numpy())
            mixeds.append(mixed.cpu().numpy())
    return np.vstack(preds), np.vstack(gates), np.vstack(mixeds)


def _mae_rows(df: pd.DataFrame, target_prefix: str) -> pd.DataFrame:
    target_cols = [c for c in df.columns if c.startswith(target_prefix)]
    model_cols = [c.replace(target_prefix, f"pred_model__{target_prefix}") for c in target_cols]
    season_cols = [c.replace(target_prefix, f"pred_season__{target_prefix}") for c in target_cols]
    short8_cols = [c.replace(target_prefix, f"pred_short8__{target_prefix}") for c in target_cols]
    out = df[["fecha", "team_name", "opponent_name"]].copy()
    out["model_mae"] = np.abs(df[target_cols].to_numpy() - df[model_cols].to_numpy()).mean(axis=1)
    out["season_mae"] = np.abs(df[target_cols].to_numpy() - df[season_cols].to_numpy()).mean(axis=1)
    out["short8_mae"] = np.abs(df[target_cols].to_numpy() - df[short8_cols].to_numpy()).mean(axis=1)
    return out


def _last3_rows(df: pd.DataFrame) -> list[dict[str, float | str]]:
    sub = df[df["team_name"].eq("Racing de Santander")].sort_values("fecha").tail(3)
    rows = sub.to_dict(orient="records")
    for row in rows:
        if "fecha" in row and pd.notna(row["fecha"]):
            row["fecha"] = pd.Timestamp(row["fecha"]).strftime("%Y-%m-%d")
    return rows


def _report_html(summary: dict) -> str:
    attack_last3 = summary["attack"]["last3_racing"]
    pv_last3 = summary["pv"]["last3_racing"]

    def rows_html(rows: list[dict[str, float | str]]) -> str:
        return "".join(
            f"<tr><td>{html.escape(str(r['fecha'])[:10])}</td><td>{html.escape(str(r['opponent_name']))}</td>"
            f"<td>{r['model_mae']:.4f}</td><td>{r['season_mae']:.4f}</td><td>{r['short8_mae']:.4f}</td><td>{r['previous_model_mae']:.4f}</td></tr>"
            for r in rows
        )

    return f"""<!DOCTYPE html>
<html lang="es">
<head>
  <meta charset="utf-8" />
  <title>GNN con embeddings de equipo</title>
  <style>
    body {{ margin: 0; background: #f1f4ef; color: #14342B; font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; }}
    .wrap {{ max-width: 1200px; margin: 0 auto; padding: 30px 24px 48px; }}
    h1 {{ margin: 0 0 10px; font-size: 38px; }}
    .lead {{ margin: 0 0 22px; font-size: 18px; color: #35574D; }}
    .grid2 {{ display: grid; grid-template-columns: repeat(2, minmax(0, 1fr)); gap: 16px; margin-bottom: 18px; }}
    .card {{ background: white; border-radius: 20px; padding: 18px 20px; box-shadow: 0 8px 24px rgba(12, 36, 28, 0.08); }}
    table {{ width: 100%; border-collapse: collapse; font-size: 14px; }}
    th, td {{ padding: 10px 8px; border-bottom: 1px solid #E5ECE6; text-align: left; }}
    th {{ color: #587468; text-transform: uppercase; font-size: 12px; letter-spacing: .06em; }}
    @media (max-width: 980px) {{ .grid2 {{ grid-template-columns: 1fr; }} }}
  </style>
</head>
<body>
  <div class="wrap">
    <h1>GNN con embeddings de equipo y rival</h1>
    <p class="lead">Mismo GNN espacial anterior, pero agregando embeddings aprendidos para equipo y rival dentro del encoder global. La comparacion se hace contra baseline temporada, baseline ultimos 8 y el GNN simple previo.</p>

    <div class="grid2">
      <section class="card">
        <h2>Ataque - test</h2>
        <table>
          <tr><th>Metrica</th><th>Embeddings</th><th>Temporada</th><th>Ultimos 8</th><th>GNN previo</th></tr>
          <tr><td>MAE</td><td>{summary['attack']['test_metrics']['model_mae']:.4f}</td><td>{summary['attack']['test_metrics']['season_baseline_mae']:.4f}</td><td>{summary['attack']['test_metrics']['short8_baseline_mae']:.4f}</td><td>{summary['attack']['previous_test_metrics']['model_mae']:.4f}</td></tr>
          <tr><td>JSD</td><td>{summary['attack']['test_metrics']['model_jsd']:.4f}</td><td>{summary['attack']['test_metrics']['season_baseline_jsd']:.4f}</td><td>{summary['attack']['test_metrics']['short8_baseline_jsd']:.4f}</td><td>{summary['attack']['previous_test_metrics']['model_jsd']:.4f}</td></tr>
          <tr><td>KL</td><td>{summary['attack']['test_metrics']['model_kl_proxy']:.4f}</td><td>{summary['attack']['test_metrics']['season_baseline_kl_proxy']:.4f}</td><td>{summary['attack']['test_metrics']['short8_baseline_kl_proxy']:.4f}</td><td>{summary['attack']['previous_test_metrics']['model_kl_proxy']:.4f}</td></tr>
        </table>
      </section>
      <section class="card">
        <h2>PV - test</h2>
        <table>
          <tr><th>Metrica</th><th>Embeddings</th><th>Temporada</th><th>Ultimos 8</th><th>GNN previo</th></tr>
          <tr><td>MAE</td><td>{summary['pv']['test_metrics']['model_mae']:.4f}</td><td>{summary['pv']['test_metrics']['season_baseline_mae']:.4f}</td><td>{summary['pv']['test_metrics']['short8_baseline_mae']:.4f}</td><td>{summary['pv']['previous_test_metrics']['model_mae']:.4f}</td></tr>
          <tr><td>JSD</td><td>{summary['pv']['test_metrics']['model_jsd']:.4f}</td><td>{summary['pv']['test_metrics']['season_baseline_jsd']:.4f}</td><td>{summary['pv']['test_metrics']['short8_baseline_jsd']:.4f}</td><td>{summary['pv']['previous_test_metrics']['model_jsd']:.4f}</td></tr>
          <tr><td>KL</td><td>{summary['pv']['test_metrics']['model_kl_proxy']:.4f}</td><td>{summary['pv']['test_metrics']['season_baseline_kl_proxy']:.4f}</td><td>{summary['pv']['test_metrics']['short8_baseline_kl_proxy']:.4f}</td><td>{summary['pv']['previous_test_metrics']['model_kl_proxy']:.4f}</td></tr>
        </table>
      </section>
    </div>

    <div class="grid2">
      <section class="card">
        <h2>Ataque - ultimos 3 de Racing</h2>
        <table>
          <tr><th>Fecha</th><th>Rival</th><th>Embeddings</th><th>Temporada</th><th>Ultimos 8</th><th>GNN previo</th></tr>
          {rows_html(attack_last3)}
        </table>
      </section>
      <section class="card">
        <h2>PV - ultimos 3 de Racing</h2>
        <table>
          <tr><th>Fecha</th><th>Rival</th><th>Embeddings</th><th>Temporada</th><th>Ultimos 8</th><th>GNN previo</th></tr>
          {rows_html(pv_last3)}
        </table>
      </section>
    </div>
  </div>
</body>
</html>"""


def main() -> None:
    _set_seed()
    MODEL_DIR.mkdir(parents=True, exist_ok=True)
    REPORTS_DIR.mkdir(parents=True, exist_ok=True)

    attack_df, attack_targets = attack_base._load_data()
    team_idx_attack, opp_idx_attack, team_mapping = _build_team_index(attack_df)
    train_idx, val_idx, test_idx, val_start_date = attack_base._train_val_test_indices(attack_df)
    attack_global_features, attack_node_tensor, _zone_order_array, attack_global_cols, attack_node_feat_names = attack_base._build_feature_matrices(attack_df, attack_targets)
    attack_global_x, attack_global_bundle = attack_base._standardize_global(train_idx, attack_global_features)
    attack_node_x, attack_node_bundle = attack_base._standardize_node(train_idx, attack_node_tensor)
    attack_y = attack_df[attack_targets].to_numpy(dtype=np.float32)
    attack_baseline_long = base._normalize_rows(attack_df[[f"long_mean__actual_attack_share__{zone}" for zone in attack_base.ZONE_ORDER]].to_numpy(dtype=float)).astype(np.float32)
    attack_baseline_short = base._normalize_rows(attack_df[[f"short_mean__actual_attack_share__{zone}" for zone in attack_base.ZONE_ORDER]].to_numpy(dtype=float)).astype(np.float32)

    attack_train = _make_attack_split(attack_df, train_idx, attack_global_x, attack_node_x, attack_y, attack_baseline_long, attack_baseline_short, team_idx_attack.to_numpy(), opp_idx_attack.to_numpy())
    attack_val = _make_attack_split(attack_df, val_idx, attack_global_x, attack_node_x, attack_y, attack_baseline_long, attack_baseline_short, team_idx_attack.to_numpy(), opp_idx_attack.to_numpy())
    attack_test = _make_attack_split(attack_df, test_idx, attack_global_x, attack_node_x, attack_y, attack_baseline_long, attack_baseline_short, team_idx_attack.to_numpy(), opp_idx_attack.to_numpy())

    attack_model, attack_history = _train_attack(attack_train, attack_val, n_teams=len(team_mapping))
    attack_val_pred, _attack_val_gates, _attack_val_mixed = _predict_attack(attack_model, attack_val)
    attack_test_pred, attack_test_gates, attack_test_mixed = _predict_attack(attack_model, attack_test)
    attack_val_metrics = attack_base._metrics_against_baselines(attack_val.y.astype(float), attack_val_pred, attack_val.baseline_long.astype(float), attack_val.baseline_short.astype(float))
    attack_test_metrics = attack_base._metrics_against_baselines(attack_test.y.astype(float), attack_test_pred, attack_test.baseline_long.astype(float), attack_test.baseline_short.astype(float))

    attack_pred_df = attack_test.metadata.copy().reset_index(drop=True)
    for i, zone in enumerate(attack_base.ZONE_ORDER):
        attack_pred_df[f"target_attack_share__{zone}"] = attack_test.y[:, i]
        attack_pred_df[f"pred_model__target_attack_share__{zone}"] = attack_test_pred[:, i]
        attack_pred_df[f"pred_season__target_attack_share__{zone}"] = attack_test.baseline_long[:, i]
        attack_pred_df[f"pred_short8__target_attack_share__{zone}"] = attack_test.baseline_short[:, i]
        attack_pred_df[f"gate__{zone}"] = attack_test_gates[:, i]
        attack_pred_df[f"mixed_base__{zone}"] = attack_test_mixed[:, i]
    attack_keep = ["matchId", "fecha", "league", "season", "teamId", "team_name", "opponent_name", "is_home", "goals_for", "goals_against", "n_prior_matches", "opp_n_prior_matches"]
    attack_keep += [f"target_attack_share__{zone}" for zone in attack_base.ZONE_ORDER]
    attack_keep += [f"pred_model__target_attack_share__{zone}" for zone in attack_base.ZONE_ORDER]
    attack_keep += [f"pred_season__target_attack_share__{zone}" for zone in attack_base.ZONE_ORDER]
    attack_keep += [f"pred_short8__target_attack_share__{zone}" for zone in attack_base.ZONE_ORDER]
    attack_keep += [f"gate__{zone}" for zone in attack_base.ZONE_ORDER]
    attack_keep += [f"mixed_base__{zone}" for zone in attack_base.ZONE_ORDER]
    attack_pred_df[attack_keep].to_parquet(ATTACK_PRED_PATH, index=False)

    pv_df = pv_base._load_data()
    team_idx_pv, opp_idx_pv, team_mapping_pv = _build_team_index(pv_df)
    pv_train_idx, pv_val_idx, pv_test_idx, pv_val_start_date = pv_base._train_val_test_indices(pv_df)
    pv_y, pv_baseline_long, pv_baseline_short = pv_base._build_distributions(pv_df)
    pv_global_features, pv_node_tensor, pv_node_feature_names = pv_base._build_feature_matrices(pv_df)
    pv_global_x, pv_global_bundle = pv_base._standardize_global(pv_train_idx, pv_global_features)
    pv_node_x, pv_node_bundle = pv_base._standardize_node(pv_train_idx, pv_node_tensor)

    pv_train = _make_pv_split(pv_df, pv_train_idx, pv_global_x, pv_node_x, pv_y, pv_baseline_long, pv_baseline_short, team_idx_pv.to_numpy(), opp_idx_pv.to_numpy())
    pv_val = _make_pv_split(pv_df, pv_val_idx, pv_global_x, pv_node_x, pv_y, pv_baseline_long, pv_baseline_short, team_idx_pv.to_numpy(), opp_idx_pv.to_numpy())
    pv_test = _make_pv_split(pv_df, pv_test_idx, pv_global_x, pv_node_x, pv_y, pv_baseline_long, pv_baseline_short, team_idx_pv.to_numpy(), opp_idx_pv.to_numpy())

    pv_model, pv_history = _train_pv(pv_train, pv_val, n_teams=len(team_mapping_pv))
    pv_val_pred, _pv_val_gates, _pv_val_mixed = _predict_pv(pv_model, pv_val)
    pv_test_pred, pv_test_gates, pv_test_mixed = _predict_pv(pv_model, pv_test)
    pv_val_metrics = pv_base._metrics(pv_val.y.astype(float), pv_val_pred, pv_val.baseline_long.astype(float), pv_val.baseline_short.astype(float))
    pv_test_metrics = pv_base._metrics(pv_test.y.astype(float), pv_test_pred, pv_test.baseline_long.astype(float), pv_test.baseline_short.astype(float))

    pv_pred_df = pv_test.metadata.copy().reset_index(drop=True)
    for i, zone in enumerate(pv_base.ZONE_ORDER):
        pv_pred_df[f"target_pv_dist__{zone}"] = pv_test.y[:, i]
        pv_pred_df[f"pred_model__target_pv_dist__{zone}"] = pv_test_pred[:, i]
        pv_pred_df[f"pred_season__target_pv_dist__{zone}"] = pv_test.baseline_long[:, i]
        pv_pred_df[f"pred_short8__target_pv_dist__{zone}"] = pv_test.baseline_short[:, i]
        pv_pred_df[f"gate__{zone}"] = pv_test_gates[:, i]
        pv_pred_df[f"mixed_base__{zone}"] = pv_test_mixed[:, i]
    pv_keep = ["matchId", "fecha", "league", "season", "teamId", "team_name", "opponent_name", "is_home", "goals_for", "goals_against", "n_prior_matches", "opp_n_prior_matches"]
    pv_keep += [f"target_pv_dist__{zone}" for zone in pv_base.ZONE_ORDER]
    pv_keep += [f"pred_model__target_pv_dist__{zone}" for zone in pv_base.ZONE_ORDER]
    pv_keep += [f"pred_season__target_pv_dist__{zone}" for zone in pv_base.ZONE_ORDER]
    pv_keep += [f"pred_short8__target_pv_dist__{zone}" for zone in pv_base.ZONE_ORDER]
    pv_keep += [f"gate__{zone}" for zone in pv_base.ZONE_ORDER]
    pv_keep += [f"mixed_base__{zone}" for zone in pv_base.ZONE_ORDER]
    pv_pred_df[pv_keep].to_parquet(PV_PRED_PATH, index=False)

    attack_prev_metrics = json.loads((MODEL_DIR / "attack_distribution_gnn_metrics.json").read_text(encoding="utf-8"))["test_metrics"]
    pv_prev_metrics = json.loads((MODEL_DIR / "pv_distribution_gnn_metrics.json").read_text(encoding="utf-8"))["test_metrics"]

    attack_last3 = _mae_rows(attack_pred_df, "target_attack_share__")
    attack_prev_last3 = _mae_rows(pd.read_parquet(MODEL_DIR / "attack_distribution_gnn_test_predictions.parquet"), "target_attack_share__")
    attack_last3 = attack_last3.merge(
        attack_prev_last3.rename(columns={"model_mae": "previous_model_mae", "season_mae": "previous_season_mae", "short8_mae": "previous_short8_mae"}),
        on=["fecha", "team_name", "opponent_name"],
        how="left",
    )

    pv_last3 = _mae_rows(pv_pred_df, "target_pv_dist__")
    pv_prev_last3 = _mae_rows(pd.read_parquet(MODEL_DIR / "pv_distribution_gnn_test_predictions.parquet"), "target_pv_dist__")
    pv_last3 = pv_last3.merge(
        pv_prev_last3.rename(columns={"model_mae": "previous_model_mae", "season_mae": "previous_season_mae", "short8_mae": "previous_short8_mae"}),
        on=["fecha", "team_name", "opponent_name"],
        how="left",
    )

    summary = {
        "embedding_dim": EMBED_DIM,
        "attack": {
            "train_rows": int(len(attack_train.metadata)),
            "val_rows": int(len(attack_val.metadata)),
            "test_rows": int(len(attack_test.metadata)),
            "val_start_date": val_start_date,
            "test_metrics": attack_test_metrics,
            "val_metrics": attack_val_metrics,
            "previous_test_metrics": attack_prev_metrics,
            "last3_racing": _last3_rows(attack_last3),
            "n_teams": len(team_mapping),
        },
        "pv": {
            "train_rows": int(len(pv_train.metadata)),
            "val_rows": int(len(pv_val.metadata)),
            "test_rows": int(len(pv_test.metadata)),
            "val_start_date": pv_val_start_date,
            "test_metrics": pv_test_metrics,
            "val_metrics": pv_val_metrics,
            "previous_test_metrics": pv_prev_metrics,
            "last3_racing": _last3_rows(pv_last3),
            "n_teams": len(team_mapping_pv),
        },
        "artifacts": {
            "attack_predictions": str(ATTACK_PRED_PATH),
            "pv_predictions": str(PV_PRED_PATH),
        },
        "training": {
            "attack_epochs": len(attack_history),
            "pv_epochs": len(pv_history),
        },
    }

    torch.save(
        {
            "embedding_dim": EMBED_DIM,
            "team_mapping_attack": team_mapping,
            "team_mapping_pv": team_mapping_pv,
            "attack_state_dict": attack_model.state_dict(),
            "pv_state_dict": pv_model.state_dict(),
            "attack_global_bundle": attack_global_bundle,
            "attack_node_bundle": attack_node_bundle,
            "attack_global_cols": attack_global_cols,
            "attack_node_feat_names": attack_node_feat_names,
            "pv_global_bundle": pv_global_bundle,
            "pv_node_bundle": pv_node_bundle,
            "pv_node_feature_names": pv_node_feature_names,
            "zone_order": attack_base.ZONE_ORDER,
        },
        MODEL_PATH,
    )
    JSON_PATH.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    REPORT_PATH.write_text(_report_html(summary), encoding="utf-8")

    print(f"Modelo guardado en: {MODEL_PATH}")
    print(f"Predicciones ataque guardadas en: {ATTACK_PRED_PATH}")
    print(f"Predicciones PV guardadas en: {PV_PRED_PATH}")
    print(f"Metricas guardadas en: {JSON_PATH}")
    print(f"Reporte guardado en: {REPORT_PATH}")
    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()