vendor/scripts/experiment_attack_prediction_variants.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 / "attack_prediction_variant_experiments.html"
JSON_PATH = MODEL_DIR / "attack_prediction_variant_experiments.json"
BEST_MODEL_PATH = MODEL_DIR / "attack_prediction_best_variant_bundle.pt"
BEST_PREDICTIONS_PATH = MODEL_DIR / "attack_prediction_best_variant_test_predictions.parquet"

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

import train_attack_prediction_ffn as base  # noqa: E402


RANDOM_SEED = 42
BATCH_SIZE = 256
MAX_EPOCHS = 180
PATIENCE = 22
LEARNING_RATE = 1e-3
WEIGHT_DECAY = 1e-4


@dataclass
class SplitData:
    x: np.ndarray
    team_idx: np.ndarray
    opp_idx: np.ndarray
    y_attack: np.ndarray
    y_pv: np.ndarray
    metadata: pd.DataFrame


class BaseVariantModel(nn.Module):
    def forward(self, x: torch.Tensor, team_idx: torch.Tensor, opp_idx: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        raise NotImplementedError


class FFNVariant(BaseVariantModel):
    def __init__(self, input_dim: int, attack_dim: int, pv_dim: int, hidden_dims: list[int], dropouts: list[float]) -> None:
        super().__init__()
        layers: list[nn.Module] = []
        prev = input_dim
        for width, drop in zip(hidden_dims, dropouts):
            layers.extend([nn.Linear(prev, width), nn.ReLU(), nn.Dropout(drop)])
            prev = width
        self.backbone = nn.Sequential(*layers)
        self.attack_head = nn.Linear(prev, attack_dim)
        self.pv_head = nn.Linear(prev, pv_dim)

    def forward(self, x: torch.Tensor, team_idx: torch.Tensor, opp_idx: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        h = self.backbone(x)
        attack = torch.softmax(self.attack_head(h), dim=1)
        pv = self.pv_head(h)
        return attack, pv


class EmbeddingVariant(BaseVariantModel):
    def __init__(
        self,
        input_dim: int,
        attack_dim: int,
        pv_dim: int,
        num_teams: int,
        team_embed_dim: int = 24,
        hidden_dims: list[int] | None = None,
        dropouts: list[float] | None = None,
        use_embeddings: bool = True,
    ) -> None:
        super().__init__()
        self.use_embeddings = use_embeddings
        hidden_dims = hidden_dims or [384, 192, 96]
        dropouts = dropouts or [0.20, 0.15, 0.10]
        embed_extra = 0
        if use_embeddings:
            self.team_embedding = nn.Embedding(num_teams + 1, team_embed_dim)
            self.opp_embedding = nn.Embedding(num_teams + 1, team_embed_dim)
            embed_extra = team_embed_dim * 2

        layers: list[nn.Module] = []
        prev = input_dim + embed_extra
        for width, drop in zip(hidden_dims, dropouts):
            layers.extend([nn.Linear(prev, width), nn.ReLU(), nn.Dropout(drop)])
            prev = width
        self.backbone = nn.Sequential(*layers)
        self.attack_head = nn.Linear(prev, attack_dim)
        self.pv_head = nn.Linear(prev, pv_dim)

    def forward(self, x: torch.Tensor, team_idx: torch.Tensor, opp_idx: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        if self.use_embeddings:
            team_emb = self.team_embedding(team_idx)
            opp_emb = self.opp_embedding(opp_idx)
            x = torch.cat([x, team_emb, opp_emb], dim=1)
        h = self.backbone(x)
        attack = torch.softmax(self.attack_head(h), dim=1)
        pv = self.pv_head(h)
        return attack, pv


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


def _loss_fn(pred_attack: torch.Tensor, pred_pv: torch.Tensor, y_attack: torch.Tensor, y_pv: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    attack_loss = torch.mean((pred_attack - y_attack) ** 2)
    pv_loss = torch.nn.functional.smooth_l1_loss(pred_pv, y_pv)
    total = attack_loss + pv_loss
    return total, attack_loss, pv_loss


def _make_loader(split: SplitData, shuffle: bool) -> DataLoader:
    dataset = TensorDataset(
        torch.from_numpy(split.x),
        torch.from_numpy(split.team_idx),
        torch.from_numpy(split.opp_idx),
        torch.from_numpy(split.y_attack),
        torch.from_numpy(split.y_pv),
    )
    return DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=shuffle)


def _evaluate_loader(model: BaseVariantModel, loader: DataLoader, device: torch.device) -> dict[str, float]:
    model.eval()
    total_loss = 0.0
    attack_loss = 0.0
    pv_loss = 0.0
    n_batches = 0
    with torch.no_grad():
        for x, team_idx, opp_idx, y_attack, y_pv in loader:
            x = x.to(device)
            team_idx = team_idx.to(device)
            opp_idx = opp_idx.to(device)
            y_attack = y_attack.to(device)
            y_pv = y_pv.to(device)
            pred_attack, pred_pv = model(x, team_idx, opp_idx)
            loss, la, lp = _loss_fn(pred_attack, pred_pv, y_attack, y_pv)
            total_loss += float(loss.item())
            attack_loss += float(la.item())
            pv_loss += float(lp.item())
            n_batches += 1
    return {
        "loss": total_loss / max(n_batches, 1),
        "attack_loss": attack_loss / max(n_batches, 1),
        "pv_loss": pv_loss / max(n_batches, 1),
    }


def _train_model(model: BaseVariantModel, train: SplitData, val: SplitData) -> tuple[BaseVariantModel, list[dict[str, float]]]:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, 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")
    history: list[dict[str, float]] = []
    patience_left = PATIENCE

    for epoch in range(1, MAX_EPOCHS + 1):
        model.train()
        running_total = 0.0
        n_batches = 0
        for x, team_idx, opp_idx, y_attack, y_pv in train_loader:
            x = x.to(device)
            team_idx = team_idx.to(device)
            opp_idx = opp_idx.to(device)
            y_attack = y_attack.to(device)
            y_pv = y_pv.to(device)

            optimizer.zero_grad()
            pred_attack, pred_pv = model(x, team_idx, opp_idx)
            loss, _, _ = _loss_fn(pred_attack, pred_pv, y_attack, y_pv)
            loss.backward()
            optimizer.step()

            running_total += float(loss.item())
            n_batches += 1

        val_metrics = _evaluate_loader(model, val_loader, device)
        history.append(
            {
                "epoch": epoch,
                "train_loss": running_total / max(n_batches, 1),
                "val_loss": val_metrics["loss"],
                "val_attack_loss": val_metrics["attack_loss"],
                "val_pv_loss": val_metrics["pv_loss"],
            }
        )
        if val_metrics["loss"] < best_val - 1e-6:
            best_val = val_metrics["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(model: BaseVariantModel, split: SplitData) -> tuple[np.ndarray, np.ndarray]:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.eval()
    attack_preds: list[np.ndarray] = []
    pv_preds: list[np.ndarray] = []
    loader = _make_loader(split, shuffle=False)
    with torch.no_grad():
        for x, team_idx, opp_idx, _, _ in loader:
            x = x.to(device)
            team_idx = team_idx.to(device)
            opp_idx = opp_idx.to(device)
            attack, pv = model(x, team_idx, opp_idx)
            attack_preds.append(attack.cpu().numpy())
            pv_preds.append(pv.cpu().numpy())
    return np.vstack(attack_preds), np.vstack(pv_preds)


def _team_index_arrays(df: pd.DataFrame) -> tuple[np.ndarray, np.ndarray, dict[str, int]]:
    team_ids = pd.Index(sorted(set(df["teamId"].astype(str)) | set(df["opponent_team_id"].astype(str))))
    mapping = {team_id: i + 1 for i, team_id in enumerate(team_ids)}
    team_idx = df["teamId"].astype(str).map(mapping).fillna(0).astype(np.int64).to_numpy()
    opp_idx = df["opponent_team_id"].astype(str).map(mapping).fillna(0).astype(np.int64).to_numpy()
    return team_idx, opp_idx, mapping


def _prepare_targets(df: pd.DataFrame, attack_targets: list[str], pv_targets: list[str], train_idx: pd.Index) -> tuple[np.ndarray, np.ndarray, dict[str, np.ndarray]]:
    y_attack = df[attack_targets].to_numpy(dtype=np.float32)
    y_pv_raw = df[pv_targets].to_numpy(dtype=np.float32)
    pv_train_raw = y_pv_raw[train_idx]
    scaler = {
        "mean": pv_train_raw.mean(axis=0, keepdims=True).astype(np.float32),
        "std": np.where(pv_train_raw.std(axis=0, keepdims=True) > 0, pv_train_raw.std(axis=0, keepdims=True), 1.0).astype(np.float32),
    }
    y_pv = ((y_pv_raw - scaler["mean"]) / scaler["std"]).astype(np.float32)
    return y_attack, y_pv, scaler


def _pack_split(
    df: pd.DataFrame,
    x_values: np.ndarray,
    team_idx: np.ndarray,
    opp_idx: np.ndarray,
    y_attack: np.ndarray,
    y_pv: np.ndarray,
    idx: np.ndarray,
) -> SplitData:
    return SplitData(
        x=x_values[idx].astype(np.float32),
        team_idx=team_idx[idx].astype(np.int64),
        opp_idx=opp_idx[idx].astype(np.int64),
        y_attack=y_attack[idx].astype(np.float32),
        y_pv=y_pv[idx].astype(np.float32),
        metadata=df.iloc[idx].copy().reset_index(drop=True),
    )


def _select_feature_columns(feature_columns: list[str], mode: str) -> list[str]:
    if mode == "full":
        return feature_columns
    if mode == "no_opp":
        return [c for c in feature_columns if not c.startswith("opp__")]
    if mode == "no_short":
        return [c for c in feature_columns if not c.startswith("short_mean__")]
    if mode == "no_long":
        return [c for c in feature_columns if not (c.startswith("long_mean__") or c.startswith("long_std__"))]
    if mode == "no_pressrec":
        keys = ("press", "pressure", "recov", "recovery", "interception")
        return [c for c in feature_columns if not any(key in c.lower() for key in keys)]
    raise ValueError(f"Ablation no reconocida: {mode}")


def _evaluate_predictions(
    metadata: pd.DataFrame,
    attack_targets: list[str],
    pv_targets: list[str],
    attack_pred: np.ndarray,
    pv_pred: np.ndarray,
) -> dict[str, float]:
    baseline_attack_cols = [base._baseline_col_from_target(c) for c in attack_targets]
    baseline_pv_cols = [base._baseline_col_from_target(c) for c in pv_targets]

    true_attack = metadata[attack_targets].to_numpy(dtype=float)
    true_pv = metadata[pv_targets].to_numpy(dtype=float)
    baseline_attack = base._normalize_rows(metadata[baseline_attack_cols].to_numpy(dtype=float))
    baseline_pv = metadata[baseline_pv_cols].to_numpy(dtype=float)

    return {
        "attack_mae_model": base._mean_abs_error(true_attack, attack_pred),
        "attack_mae_baseline": base._mean_abs_error(true_attack, baseline_attack),
        "attack_jsd_model": base._jsd_mean(true_attack, attack_pred),
        "attack_jsd_baseline": base._jsd_mean(true_attack, baseline_attack),
        "pv_mae_model": base._mean_abs_error(true_pv, pv_pred),
        "pv_mae_baseline": base._mean_abs_error(true_pv, baseline_pv),
        "pv_r2_model": base._r2_score_mean(true_pv, pv_pred),
        "pv_r2_baseline": base._r2_score_mean(true_pv, baseline_pv),
    }


def _run_single_experiment(
    df: pd.DataFrame,
    x_scaled: pd.DataFrame,
    y_attack: np.ndarray,
    y_pv: np.ndarray,
    team_idx: np.ndarray,
    opp_idx: np.ndarray,
    attack_targets: list[str],
    pv_targets: list[str],
    pv_scaler: dict[str, np.ndarray],
    train_idx: np.ndarray,
    val_idx: np.ndarray,
    test_idx: np.ndarray,
    variant_name: str,
    feature_mode: str,
    model_factory,
) -> dict:
    selected_cols = _select_feature_columns(list(x_scaled.columns), feature_mode)
    x_values = x_scaled[selected_cols].to_numpy(dtype=np.float32)

    train_split = _pack_split(df, x_values, team_idx, opp_idx, y_attack, y_pv, train_idx)
    val_split = _pack_split(df, x_values, team_idx, opp_idx, y_attack, y_pv, val_idx)
    test_split = _pack_split(df, x_values, team_idx, opp_idx, y_attack, y_pv, test_idx)

    model = model_factory(input_dim=len(selected_cols), attack_dim=len(attack_targets), pv_dim=len(pv_targets))
    model, history = _train_model(model, train_split, val_split)

    val_attack_pred, val_pv_pred_scaled = _predict(model, val_split)
    test_attack_pred, test_pv_pred_scaled = _predict(model, test_split)
    val_pv_pred = (val_pv_pred_scaled * pv_scaler["std"]) + pv_scaler["mean"]
    test_pv_pred = (test_pv_pred_scaled * pv_scaler["std"]) + pv_scaler["mean"]

    val_metrics = _evaluate_predictions(val_split.metadata, attack_targets, pv_targets, val_attack_pred, val_pv_pred)
    test_metrics = _evaluate_predictions(test_split.metadata, attack_targets, pv_targets, test_attack_pred, test_pv_pred)

    return {
        "variant_name": variant_name,
        "feature_mode": feature_mode,
        "selected_feature_count": len(selected_cols),
        "history": history,
        "val_metrics": val_metrics,
        "test_metrics": test_metrics,
        "model_state_dict": {k: v.detach().cpu() for k, v in model.state_dict().items()},
        "selected_columns": selected_cols,
        "test_attack_pred": test_attack_pred,
        "test_pv_pred": test_pv_pred,
    }


def _variant_factories(num_teams: int) -> dict[str, object]:
    return {
        "small_ffn": lambda input_dim, attack_dim, pv_dim: FFNVariant(
            input_dim=input_dim,
            attack_dim=attack_dim,
            pv_dim=pv_dim,
            hidden_dims=[128, 64],
            dropouts=[0.10, 0.05],
        ),
        "large_ffn": lambda input_dim, attack_dim, pv_dim: FFNVariant(
            input_dim=input_dim,
            attack_dim=attack_dim,
            pv_dim=pv_dim,
            hidden_dims=[768, 384, 192],
            dropouts=[0.25, 0.20, 0.15],
        ),
        "team_embeddings": lambda input_dim, attack_dim, pv_dim: EmbeddingVariant(
            input_dim=input_dim,
            attack_dim=attack_dim,
            pv_dim=pv_dim,
            num_teams=num_teams,
            team_embed_dim=24,
            hidden_dims=[384, 192, 96],
            dropouts=[0.20, 0.15, 0.10],
            use_embeddings=True,
        ),
        "team_embeddings_no_embed": lambda input_dim, attack_dim, pv_dim: EmbeddingVariant(
            input_dim=input_dim,
            attack_dim=attack_dim,
            pv_dim=pv_dim,
            num_teams=num_teams,
            team_embed_dim=24,
            hidden_dims=[384, 192, 96],
            dropouts=[0.20, 0.15, 0.10],
            use_embeddings=False,
        ),
    }


def _comparison_rows(results: list[dict], split_key: str) -> pd.DataFrame:
    rows = []
    for result in results:
        metrics = result[f"{split_key}_metrics"]
        rows.append(
            {
                "modelo": result["variant_name"],
                "features": result["feature_mode"],
                "attack_mae_model": metrics["attack_mae_model"],
                "attack_mae_baseline": metrics["attack_mae_baseline"],
                "attack_delta": metrics["attack_mae_baseline"] - metrics["attack_mae_model"],
                "pv_mae_model": metrics["pv_mae_model"],
                "pv_mae_baseline": metrics["pv_mae_baseline"],
                "pv_delta": metrics["pv_mae_baseline"] - metrics["pv_mae_model"],
                "pv_r2_model": metrics["pv_r2_model"],
                "pv_r2_baseline": metrics["pv_r2_baseline"],
            }
        )
    return pd.DataFrame(rows).sort_values(["attack_mae_model", "pv_mae_model"]).reset_index(drop=True)


def _render_table(df: pd.DataFrame, title: str) -> str:
    rows = []
    for _, row in df.iterrows():
        rows.append(
            "<tr>"
            f"<td>{html.escape(str(row['modelo']))}</td>"
            f"<td>{html.escape(str(row['features']))}</td>"
            f"<td>{row['attack_mae_model']:.4f}</td>"
            f"<td>{row['attack_mae_baseline']:.4f}</td>"
            f"<td>{row['attack_delta']:+.4f}</td>"
            f"<td>{row['pv_mae_model']:.4f}</td>"
            f"<td>{row['pv_mae_baseline']:.4f}</td>"
            f"<td>{row['pv_delta']:+.4f}</td>"
            f"<td>{row['pv_r2_model']:.4f}</td>"
            f"<td>{row['pv_r2_baseline']:.4f}</td>"
            "</tr>"
        )
    return f"""
    <section class="card">
      <h2>{html.escape(title)}</h2>
      <table>
        <tr>
          <th>Modelo</th><th>Features</th><th>Attack MAE</th><th>Baseline</th><th>Mejora</th>
          <th>PV MAE</th><th>Baseline</th><th>Mejora</th><th>PV R2</th><th>Baseline R2</th>
        </tr>
        {''.join(rows)}
      </table>
    </section>
    """


def _render_match_table(best_predictions: pd.DataFrame, attack_targets: list[str], pv_targets: list[str]) -> str:
    rows = []
    racing = best_predictions[best_predictions["teamId"] == base.RACING_TEAM_ID].sort_values(["fecha", "matchId"]).tail(3)
    for _, row in racing.iterrows():
        attack_true = row[attack_targets].to_numpy(dtype=float)
        attack_model = row[[f"pred_model__{c}" for c in attack_targets]].to_numpy(dtype=float)
        attack_base = row[[f"pred_baseline__{c}" for c in attack_targets]].to_numpy(dtype=float)
        pv_true = row[pv_targets].to_numpy(dtype=float)
        pv_model = row[[f"pred_model__{c}" for c in pv_targets]].to_numpy(dtype=float)
        pv_base = row[[f"pred_baseline__{c}" for c in pv_targets]].to_numpy(dtype=float)
        rows.append(
            "<tr>"
            f"<td>{row['fecha'].strftime('%Y-%m-%d')}</td>"
            f"<td>{html.escape(str(row.get('opponent_name', '')))}</td>"
            f"<td>{np.mean(np.abs(attack_true - attack_model)):.4f}</td>"
            f"<td>{np.mean(np.abs(attack_true - attack_base)):.4f}</td>"
            f"<td>{np.mean(np.abs(pv_true - pv_model)):.4f}</td>"
            f"<td>{np.mean(np.abs(pv_true - pv_base)):.4f}</td>"
            "</tr>"
        )
    return f"""
    <section class="card">
      <h2>Ultimos 3 partidos de Racing en test</h2>
      <table>
        <tr><th>Fecha</th><th>Rival</th><th>Attack MAE modelo</th><th>Attack MAE baseline</th><th>PV MAE modelo</th><th>PV MAE baseline</th></tr>
        {''.join(rows)}
      </table>
    </section>
    """


def _build_report(summary: dict, variant_table_val: pd.DataFrame, variant_table_test: pd.DataFrame, ablation_table: pd.DataFrame, racing_table: str) -> str:
    return f"""<!DOCTYPE html>
<html lang="es">
<head>
  <meta charset="utf-8" />
  <title>Experimentos de variantes del modelo</title>
  <style>
    body {{ margin: 0; background: #f1f4ef; color: #14342B; font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; }}
    .wrap {{ max-width: 1360px; margin: 0 auto; padding: 30px 24px 48px; }}
    h1 {{ margin: 0 0 10px; font-size: 40px; }}
    .lead {{ margin: 0 0 22px; font-size: 18px; color: #35574D; }}
    .hero {{ display: grid; grid-template-columns: repeat(4, minmax(0, 1fr)); gap: 14px; margin-bottom: 22px; }}
    .stat, .card {{ background: white; border-radius: 20px; padding: 18px 20px; box-shadow: 0 8px 24px rgba(12, 36, 28, 0.08); }}
    .stat h3 {{ margin: 0 0 8px; font-size: 12px; text-transform: uppercase; letter-spacing: .08em; color: #587468; }}
    .stat p {{ margin: 0; font-size: 28px; font-weight: 800; }}
    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; }}
    .card {{ margin-bottom: 18px; }}
    @media (max-width: 980px) {{ .hero {{ grid-template-columns: 1fr; }} }}
  </style>
</head>
<body>
  <div class="wrap">
    <h1>Experimentos de variantes del modelo</h1>
    <p class="lead">Se compararon tres variantes de arquitectura sobre el mismo split temporal. La mejor se eligio por `attack MAE` en validacion, usando `PV MAE` como desempate.</p>
    <section class="hero">
      <div class="stat"><h3>Mejor variante</h3><p>{html.escape(summary['best_variant'])}</p></div>
      <div class="stat"><h3>Train</h3><p>{summary['train_rows']}</p></div>
      <div class="stat"><h3>Val</h3><p>{summary['val_rows']}</p></div>
      <div class="stat"><h3>Test</h3><p>{summary['test_rows']}</p></div>
    </section>
    {_render_table(variant_table_val, "Comparacion de variantes - validacion")}
    {_render_table(variant_table_test, "Comparacion de variantes - test")}
    {_render_table(ablation_table, f"Ablation test - {summary['best_variant']}")}
    {racing_table}
  </div>
</body>
</html>"""


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

    df = base._load_dataset()
    df = df[df["usable_for_model"]].copy().reset_index(drop=True)

    x_scaled, numeric_cols, attack_targets, pv_targets, dummy_cols = None, None, None, None, None
    features, numeric_cols, attack_targets, pv_targets, dummy_cols = base._feature_matrix(df)
    train_df = df[df["split"] == "train"].copy()
    train_idx_pd, val_idx_pd, val_start_date = base._build_temporal_validation(train_df)
    x_scaled_df, scaler_bundle = base._standardize_features(features, train_idx_pd)

    team_idx, opp_idx, team_mapping = _team_index_arrays(df)
    y_attack, y_pv, pv_scaler = _prepare_targets(df, attack_targets, pv_targets, train_idx_pd.to_numpy())
    train_idx = train_idx_pd.to_numpy()
    val_idx = val_idx_pd.to_numpy()
    test_idx = df.index[df["split"] == "test"].to_numpy()

    factories = _variant_factories(num_teams=len(team_mapping))
    variant_order = ["small_ffn", "large_ffn", "team_embeddings"]
    variant_results: list[dict] = []

    for variant_name in variant_order:
        result = _run_single_experiment(
            df=df,
            x_scaled=x_scaled_df,
            y_attack=y_attack,
            y_pv=y_pv,
            team_idx=team_idx,
            opp_idx=opp_idx,
            attack_targets=attack_targets,
            pv_targets=pv_targets,
            pv_scaler=pv_scaler,
            train_idx=train_idx,
            val_idx=val_idx,
            test_idx=test_idx,
            variant_name=variant_name,
            feature_mode="full",
            model_factory=factories[variant_name],
        )
        variant_results.append(result)

    variant_table_val = _comparison_rows(variant_results, "val")
    variant_table_test = _comparison_rows(variant_results, "test")
    best_row = variant_table_val.sort_values(["attack_mae_model", "pv_mae_model"]).iloc[0]
    best_variant = str(best_row["modelo"])
    best_result = next(r for r in variant_results if r["variant_name"] == best_variant)

    ablation_modes = ["full", "no_opp", "no_short", "no_long", "no_pressrec"]
    ablation_results: list[dict] = []
    for feature_mode in ablation_modes:
        ablation_results.append(
            _run_single_experiment(
                df=df,
                x_scaled=x_scaled_df,
                y_attack=y_attack,
                y_pv=y_pv,
                team_idx=team_idx,
                opp_idx=opp_idx,
                attack_targets=attack_targets,
                pv_targets=pv_targets,
                pv_scaler=pv_scaler,
                train_idx=train_idx,
                val_idx=val_idx,
                test_idx=test_idx,
                variant_name=best_variant,
                feature_mode=feature_mode,
                model_factory=factories[best_variant],
            )
        )

    if best_variant == "team_embeddings":
        ablation_results.append(
            _run_single_experiment(
                df=df,
                x_scaled=x_scaled_df,
                y_attack=y_attack,
                y_pv=y_pv,
                team_idx=team_idx,
                opp_idx=opp_idx,
                attack_targets=attack_targets,
                pv_targets=pv_targets,
                pv_scaler=pv_scaler,
                train_idx=train_idx,
                val_idx=val_idx,
                test_idx=test_idx,
                variant_name="team_embeddings_no_embed",
                feature_mode="full",
                model_factory=factories["team_embeddings_no_embed"],
            )
        )

    ablation_table = _comparison_rows(ablation_results, "test")

    best_test_meta = _pack_split(
        df,
        x_scaled_df[best_result["selected_columns"]].to_numpy(dtype=np.float32),
        team_idx,
        opp_idx,
        y_attack,
        y_pv,
        test_idx,
    ).metadata
    best_attack_pred = best_result["test_attack_pred"]
    best_pv_pred = best_result["test_pv_pred"]
    baseline_attack_cols = [base._baseline_col_from_target(c) for c in attack_targets]
    baseline_pv_cols = [base._baseline_col_from_target(c) for c in pv_targets]
    baseline_attack = base._normalize_rows(best_test_meta[baseline_attack_cols].to_numpy(dtype=float))
    baseline_pv = best_test_meta[baseline_pv_cols].to_numpy(dtype=float)
    for i, col in enumerate(attack_targets):
        best_test_meta[f"pred_model__{col}"] = best_attack_pred[:, i]
        best_test_meta[f"pred_baseline__{col}"] = baseline_attack[:, i]
    for i, col in enumerate(pv_targets):
        best_test_meta[f"pred_model__{col}"] = best_pv_pred[:, i]
        best_test_meta[f"pred_baseline__{col}"] = baseline_pv[:, i]

    keep_cols = [
        "matchId", "fecha", "league", "season", "teamId", "team_name", "opponent_name", "is_home",
        "goals_for", "goals_against", "n_prior_matches", "opp_n_prior_matches",
    ]
    pred_cols = keep_cols + attack_targets + pv_targets
    pred_cols += [f"pred_model__{c}" for c in attack_targets + pv_targets]
    pred_cols += [f"pred_baseline__{c}" for c in attack_targets + pv_targets]
    best_test_meta[pred_cols].to_parquet(BEST_PREDICTIONS_PATH, index=False)

    racing_table = _render_match_table(best_test_meta, attack_targets, pv_targets)

    summary = {
        "best_variant": best_variant,
        "selection_rule": "min val attack_mae_model, tie-break by val pv_mae_model",
        "train_rows": int(len(train_idx)),
        "val_rows": int(len(val_idx)),
        "test_rows": int(len(test_idx)),
        "val_start_date": val_start_date,
        "feature_count_full": int(x_scaled_df.shape[1]),
        "attack_target_count": len(attack_targets),
        "pv_target_count": len(pv_targets),
        "variant_results": [
            {
                "variant_name": r["variant_name"],
                "feature_mode": r["feature_mode"],
                "selected_feature_count": r["selected_feature_count"],
                "val_metrics": r["val_metrics"],
                "test_metrics": r["test_metrics"],
            }
            for r in variant_results
        ],
        "ablation_results": [
            {
                "variant_name": r["variant_name"],
                "feature_mode": r["feature_mode"],
                "selected_feature_count": r["selected_feature_count"],
                "val_metrics": r["val_metrics"],
                "test_metrics": r["test_metrics"],
            }
            for r in ablation_results
        ],
    }

    best_bundle = {
        "variant_name": best_variant,
        "feature_mode": best_result["feature_mode"],
        "selected_columns": best_result["selected_columns"],
        "model_state_dict": best_result["model_state_dict"],
        "feature_bundle": scaler_bundle,
        "pv_target_scaler": {
            "mean": pv_scaler["mean"].tolist(),
            "std": pv_scaler["std"].tolist(),
        },
        "team_mapping": team_mapping,
        "attack_targets": attack_targets,
        "pv_targets": pv_targets,
        "summary": summary,
    }
    torch.save(best_bundle, BEST_MODEL_PATH)
    JSON_PATH.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")

    report_html = _build_report(summary, variant_table_val, variant_table_test, ablation_table, racing_table)
    REPORT_PATH.write_text(report_html, encoding="utf-8")

    print(f"Resumen guardado en: {JSON_PATH}")
    print(f"Mejor modelo guardado en: {BEST_MODEL_PATH}")
    print(f"Predicciones test guardadas en: {BEST_PREDICTIONS_PATH}")
    print(f"Reporte guardado en: {REPORT_PATH}")
    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()