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_v2_experiments.html"
JSON_PATH = MODEL_DIR / "attack_prediction_v2_experiments.json"
BEST_MODEL_PATH = MODEL_DIR / "attack_prediction_v2_best_bundle.pt"
BEST_PREDICTIONS_PATH = MODEL_DIR / "attack_prediction_v2_best_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 = 220
PATIENCE = 28


@dataclass
class SplitData:
    x: np.ndarray
    y_attack: np.ndarray
    y_pv: np.ndarray
    baseline_attack: np.ndarray
    baseline_pv_scaled: np.ndarray
    metadata: pd.DataFrame


class BaseModel(nn.Module):
    def forward(
        self,
        x: torch.Tensor,
        baseline_attack: torch.Tensor,
        baseline_pv_scaled: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        raise NotImplementedError


class FFNModel(BaseModel):
    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,
        baseline_attack: torch.Tensor,
        baseline_pv_scaled: 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 ResidualModel(BaseModel):
    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_delta = nn.Linear(prev, attack_dim)
        self.pv_delta = nn.Linear(prev, pv_dim)

    def forward(
        self,
        x: torch.Tensor,
        baseline_attack: torch.Tensor,
        baseline_pv_scaled: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        h = self.backbone(x)
        attack_logits = torch.log(torch.clamp(baseline_attack, min=1e-6)) + self.attack_delta(h)
        attack = torch.softmax(attack_logits, dim=1)
        pv = baseline_pv_scaled + self.pv_delta(h)
        return attack, pv


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


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


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)
    return attack_loss + pv_loss, attack_loss, pv_loss


def _evaluate_loader(model: BaseModel, 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, y_attack, y_pv, baseline_attack, baseline_pv_scaled in loader:
            x = x.to(device)
            y_attack = y_attack.to(device)
            y_pv = y_pv.to(device)
            baseline_attack = baseline_attack.to(device)
            baseline_pv_scaled = baseline_pv_scaled.to(device)
            pred_attack, pred_pv = model(x, baseline_attack, baseline_pv_scaled)
            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: BaseModel, train: SplitData, val: SplitData, lr: float, weight_decay: float) -> tuple[BaseModel, 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=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")
    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, y_attack, y_pv, baseline_attack, baseline_pv_scaled in train_loader:
            x = x.to(device)
            y_attack = y_attack.to(device)
            y_pv = y_pv.to(device)
            baseline_attack = baseline_attack.to(device)
            baseline_pv_scaled = baseline_pv_scaled.to(device)

            optimizer.zero_grad()
            pred_attack, pred_pv = model(x, baseline_attack, baseline_pv_scaled)
            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: BaseModel, 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, _, _, baseline_attack, baseline_pv_scaled in loader:
            x = x.to(device)
            baseline_attack = baseline_attack.to(device)
            baseline_pv_scaled = baseline_pv_scaled.to(device)
            attack, pv = model(x, baseline_attack, baseline_pv_scaled)
            attack_preds.append(attack.cpu().numpy())
            pv_preds.append(pv.cpu().numpy())
    return np.vstack(attack_preds), np.vstack(pv_preds)


def _prepare_data() -> tuple[pd.DataFrame, pd.DataFrame, list[str], list[str], np.ndarray, np.ndarray, dict[str, np.ndarray], np.ndarray, np.ndarray, np.ndarray]:
    df = base._load_dataset()
    df = df[df["usable_for_model"]].copy().reset_index(drop=True)

    features, _, attack_targets, pv_targets, _ = base._feature_matrix(df)
    train_df = df[df["split"] == "train"].copy()
    train_idx_pd, val_idx_pd, _ = base._build_temporal_validation(train_df)
    x_scaled, _ = base._standardize_features(features, train_idx_pd)

    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_pd.to_numpy()]
    pv_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 - pv_scaler["mean"]) / pv_scaler["std"]).astype(np.float32)

    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(df[baseline_attack_cols].to_numpy(dtype=float)).astype(np.float32)
    baseline_pv_scaled = ((df[baseline_pv_cols].to_numpy(dtype=np.float32) - pv_scaler["mean"]) / pv_scaler["std"]).astype(np.float32)

    return (
        df,
        x_scaled,
        attack_targets,
        pv_targets,
        y_attack,
        y_pv,
        pv_scaler,
        baseline_attack,
        baseline_pv_scaled,
        train_idx_pd.to_numpy(),
        val_idx_pd.to_numpy(),
        df.index[df["split"] == "test"].to_numpy(),
    )


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


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_experiment(
    df: pd.DataFrame,
    x_scaled_df: pd.DataFrame,
    attack_targets: list[str],
    pv_targets: list[str],
    y_attack: np.ndarray,
    y_pv: np.ndarray,
    pv_scaler: dict[str, np.ndarray],
    baseline_attack: np.ndarray,
    baseline_pv_scaled: np.ndarray,
    train_idx: np.ndarray,
    val_idx: np.ndarray,
    test_idx: np.ndarray,
    name: str,
    hidden_dims: list[int],
    dropouts: list[float],
    lr: float,
    weight_decay: float,
    residual: bool,
) -> dict:
    x = x_scaled_df.to_numpy(dtype=np.float32)
    train_split = _pack_split(df, x, y_attack, y_pv, baseline_attack, baseline_pv_scaled, train_idx)
    val_split = _pack_split(df, x, y_attack, y_pv, baseline_attack, baseline_pv_scaled, val_idx)
    test_split = _pack_split(df, x, y_attack, y_pv, baseline_attack, baseline_pv_scaled, test_idx)

    if residual:
        model = ResidualModel(len(x_scaled_df.columns), len(attack_targets), len(pv_targets), hidden_dims, dropouts)
    else:
        model = FFNModel(len(x_scaled_df.columns), len(attack_targets), len(pv_targets), hidden_dims, dropouts)

    model, history = _train_model(model, train_split, val_split, lr=lr, weight_decay=weight_decay)
    val_attack_pred, val_pv_scaled = _predict(model, val_split)
    test_attack_pred, test_pv_scaled = _predict(model, test_split)
    val_pv_pred = (val_pv_scaled * pv_scaler["std"]) + pv_scaler["mean"]
    test_pv_pred = (test_pv_scaled * pv_scaler["std"]) + pv_scaler["mean"]

    return {
        "name": name,
        "residual": residual,
        "hidden_dims": hidden_dims,
        "dropouts": dropouts,
        "lr": lr,
        "weight_decay": weight_decay,
        "history": history,
        "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),
        "model_state_dict": {k: v.detach().cpu() for k, v in model.state_dict().items()},
        "test_attack_pred": test_attack_pred,
        "test_pv_pred": test_pv_pred,
    }


def _results_table(results: list[dict], split_key: str) -> pd.DataFrame:
    rows = []
    for r in results:
        m = r[f"{split_key}_metrics"]
        rows.append(
            {
                "modelo": r["name"],
                "residual": r["residual"],
                "attack_mae_model": m["attack_mae_model"],
                "attack_mae_baseline": m["attack_mae_baseline"],
                "attack_delta": m["attack_mae_baseline"] - m["attack_mae_model"],
                "attack_jsd_model": m["attack_jsd_model"],
                "attack_jsd_baseline": m["attack_jsd_baseline"],
                "pv_mae_model": m["pv_mae_model"],
                "pv_mae_baseline": m["pv_mae_baseline"],
                "pv_delta": m["pv_mae_baseline"] - m["pv_mae_model"],
                "pv_r2_model": m["pv_r2_model"],
                "pv_r2_baseline": m["pv_r2_baseline"],
            }
        )
    return pd.DataFrame(rows).sort_values(["attack_mae_model", "pv_mae_model"]).reset_index(drop=True)


def _render_table(title: str, df: pd.DataFrame) -> str:
    rows = []
    for _, row in df.iterrows():
        rows.append(
            "<tr>"
            f"<td>{html.escape(str(row['modelo']))}</td>"
            f"<td>{'si' if row['residual'] else 'no'}</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['attack_jsd_model']:.4f}</td>"
            f"<td>{row['attack_jsd_baseline']:.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>Residual</th><th>Attack MAE</th><th>Baseline</th><th>Mejora</th>
          <th>Attack JSD</th><th>Baseline</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_racing_table(best_test_meta: pd.DataFrame, attack_targets: list[str], pv_targets: list[str]) -> str:
    racing = best_test_meta[best_test_meta["teamId"] == base.RACING_TEAM_ID].sort_values(["fecha", "matchId"]).tail(3)
    rows = []
    for _, row in racing.iterrows():
        att_true = row[attack_targets].to_numpy(dtype=float)
        att_model = row[[f"pred_model__{c}" for c in attack_targets]].to_numpy(dtype=float)
        att_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(att_true - att_model)):.4f}</td>"
            f"<td>{np.mean(np.abs(att_true - att_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, val_table: pd.DataFrame, test_table: pd.DataFrame, racing_table: str) -> str:
    return f"""<!DOCTYPE html>
<html lang="es">
<head>
  <meta charset="utf-8" />
  <title>Experimentos v2 ataque/PV</title>
  <style>
    body {{ margin: 0; background: #f1f4ef; color: #14342B; font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", sans-serif; }}
    .wrap {{ max-width: 1380px; 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; }}
    .card {{ margin-bottom: 18px; }}
    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) {{ .hero {{ grid-template-columns: 1fr; }} }}
  </style>
</head>
<body>
  <div class="wrap">
    <h1>Experimentos v2: hiperparametros, red mas grande y residual</h1>
    <p class="lead">Segunda tanda enfocada en mejorar ataque: nuevas configuraciones, una red bastante mas grande y un modelo residual que parte del baseline de temporada y aprende una correccion.</p>
    <section class="hero">
      <div class="stat"><h3>Mejor modelo</h3><p>{html.escape(summary['best_model'])}</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("Comparacion por validacion", val_table)}
    {_render_table("Comparacion por test", test_table)}
    {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,
        x_scaled_df,
        attack_targets,
        pv_targets,
        y_attack,
        y_pv,
        pv_scaler,
        baseline_attack,
        baseline_pv_scaled,
        train_idx,
        val_idx,
        test_idx,
    ) = _prepare_data()

    configs = [
        {
            "name": "large_ffn_prev_like",
            "hidden_dims": [768, 384, 192],
            "dropouts": [0.25, 0.20, 0.15],
            "lr": 1e-3,
            "weight_decay": 1e-4,
            "residual": False,
        },
        {
            "name": "xlarge_ffn_tuned",
            "hidden_dims": [1536, 768, 384, 192],
            "dropouts": [0.18, 0.12, 0.08, 0.05],
            "lr": 5e-4,
            "weight_decay": 5e-5,
            "residual": False,
        },
        {
            "name": "xlarge_ffn_lowdrop",
            "hidden_dims": [2048, 1024, 512, 256],
            "dropouts": [0.10, 0.08, 0.05, 0.03],
            "lr": 3e-4,
            "weight_decay": 1e-5,
            "residual": False,
        },
        {
            "name": "residual_xlarge",
            "hidden_dims": [1024, 512, 256, 128],
            "dropouts": [0.12, 0.08, 0.05, 0.03],
            "lr": 4e-4,
            "weight_decay": 1e-5,
            "residual": True,
        },
    ]

    results = []
    for cfg in configs:
        results.append(
            _run_experiment(
                df=df,
                x_scaled_df=x_scaled_df,
                attack_targets=attack_targets,
                pv_targets=pv_targets,
                y_attack=y_attack,
                y_pv=y_pv,
                pv_scaler=pv_scaler,
                baseline_attack=baseline_attack,
                baseline_pv_scaled=baseline_pv_scaled,
                train_idx=train_idx,
                val_idx=val_idx,
                test_idx=test_idx,
                **cfg,
            )
        )

    val_table = _results_table(results, "val")
    test_table = _results_table(results, "test")
    best_name = str(val_table.sort_values(["attack_mae_model", "pv_mae_model"]).iloc[0]["modelo"])
    best_result = next(r for r in results if r["name"] == best_name)

    best_test_meta = df.iloc[test_idx].copy().reset_index(drop=True)
    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]
    best_test_meta["goals_for"] = best_test_meta.get("goals_for")
    best_test_meta["goals_against"] = best_test_meta.get("goals_against")
    best_attack_pred = best_result["test_attack_pred"]
    best_pv_pred = best_result["test_pv_pred"]
    baseline_attack_test = base._normalize_rows(best_test_meta[baseline_attack_cols].to_numpy(dtype=float))
    baseline_pv_test = 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_test[:, 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_test[:, 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)

    summary = {
        "best_model": best_name,
        "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)),
        "experiments": [
            {
                "name": r["name"],
                "residual": r["residual"],
                "hidden_dims": r["hidden_dims"],
                "dropouts": r["dropouts"],
                "lr": r["lr"],
                "weight_decay": r["weight_decay"],
                "val_metrics": r["val_metrics"],
                "test_metrics": r["test_metrics"],
            }
            for r in results
        ],
    }

    torch.save(
        {
            "best_model": best_name,
            "model_state_dict": best_result["model_state_dict"],
            "attack_targets": attack_targets,
            "pv_targets": pv_targets,
            "feature_columns": list(x_scaled_df.columns),
            "pv_target_scaler": {
                "mean": pv_scaler["mean"].tolist(),
                "std": pv_scaler["std"].tolist(),
            },
            "summary": summary,
        },
        BEST_MODEL_PATH,
    )
    JSON_PATH.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    REPORT_PATH.write_text(_build_report(summary, val_table, test_table, _render_racing_table(best_test_meta, attack_targets, pv_targets)), encoding="utf-8")

    print(f"Resumen guardado en: {JSON_PATH}")
    print(f"Mejor modelo guardado en: {BEST_MODEL_PATH}")
    print(f"Predicciones 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()