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.distributions import Dirichlet
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 / "bayesian_dirichlet_map_report.html"
JSON_PATH = MODEL_DIR / "bayesian_dirichlet_map_metrics.json"
MODEL_PATH = MODEL_DIR / "bayesian_dirichlet_map_bundle.pt"
ATTACK_PRED_PATH = MODEL_DIR / "attack_bayesian_dirichlet_map_test_predictions.parquet"
PV_PRED_PATH = MODEL_DIR / "pv_bayesian_dirichlet_map_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_mod  # noqa: E402
import experiment_pv_distribution_gnn as pv_base  # noqa: E402
import train_attack_prediction_ffn as base  # noqa: E402
import train_attack_prediction_ffn as train_model_mod  # noqa: E402


RANDOM_SEED = 42
BATCH_SIZE = 256
MAX_EPOCHS = 220
PATIENCE = 28
TARGET_EPS = 1e-4


@dataclass
class SplitData:
    x: np.ndarray
    y: np.ndarray
    prior: np.ndarray
    prior_short: np.ndarray
    metadata: pd.DataFrame


@dataclass
class Config:
    name: str
    lr: float
    prior_lambda: float
    init_concentration: float


class BayesianDirichletMAP(nn.Module):
    def __init__(self, input_dim: int, out_dim: int, init_concentration: float) -> None:
        super().__init__()
        self.linear = nn.Linear(input_dim, out_dim)
        init_raw = np.log(np.exp(max(init_concentration - 1.0, 1e-3)) - 1.0)
        self.raw_concentration = nn.Parameter(torch.tensor(float(init_raw), dtype=torch.float32))

    def forward(self, x: torch.Tensor, prior: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        delta = self.linear(x)
        logits = torch.log(torch.clamp(prior, min=1e-6)) + delta
        mean = torch.softmax(logits, dim=1)
        concentration = torch.nn.functional.softplus(self.raw_concentration) + 1.0
        alpha = mean * concentration + 1e-4
        return mean, alpha, concentration


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


def _smooth_target(y: np.ndarray) -> np.ndarray:
    k = y.shape[1]
    smoothed = ((y + TARGET_EPS) / (1.0 + (TARGET_EPS * k))).astype(np.float32)
    return base._normalize_rows(smoothed).astype(np.float32)


def _load_attack_data() -> tuple[pd.DataFrame, pd.DataFrame, list[str]]:
    df = base._load_dataset()
    df = df[df["usable_for_model"]].copy().reset_index(drop=True)
    features, _numeric_cols, _attack_targets_sorted, _pv_targets, _dummy_cols = base._feature_matrix(df)
    attack_targets = [f"target_attack_share__{zone}" for zone in attack_mod.ZONE_ORDER]
    return df, features, attack_targets


def _load_pv_data() -> tuple[pd.DataFrame, pd.DataFrame, np.ndarray]:
    df = base._load_dataset()
    df = df[df["usable_for_model"]].copy().reset_index(drop=True)
    y_dist, baseline_long, baseline_short = pv_base._build_distributions(df)
    keep_mask = baseline_long.sum(axis=1) > 0
    df = df.loc[keep_mask].copy().reset_index(drop=True)
    features, _numeric_cols, _attack_targets, _pv_targets, _dummy_cols = base._feature_matrix(df)
    y_dist, baseline_long, baseline_short = pv_base._build_distributions(df)
    return df, features, y_dist


def _make_split(
    df: pd.DataFrame,
    idx: pd.Index | np.ndarray,
    x_scaled: pd.DataFrame,
    y: np.ndarray,
    prior: np.ndarray,
    prior_short: np.ndarray,
) -> SplitData:
    arr_idx = np.asarray(idx)
    return SplitData(
        x=x_scaled.iloc[arr_idx].to_numpy(dtype=np.float32),
        y=y[arr_idx].astype(np.float32),
        prior=prior[arr_idx].astype(np.float32),
        prior_short=prior_short[arr_idx].astype(np.float32),
        metadata=df.iloc[arr_idx].copy().reset_index(drop=True),
    )


def _loader(split: SplitData, shuffle: bool) -> DataLoader:
    ds = TensorDataset(
        torch.from_numpy(split.x),
        torch.from_numpy(split.y),
        torch.from_numpy(split.prior),
        torch.from_numpy(split.prior_short),
    )
    return DataLoader(ds, batch_size=BATCH_SIZE, shuffle=shuffle)


def _l2_penalty(model: BayesianDirichletMAP) -> torch.Tensor:
    total = torch.zeros((), device=model.linear.weight.device)
    for param in model.parameters():
        total = total + torch.sum(param ** 2)
    return total


def _loss(
    model: BayesianDirichletMAP,
    mean: torch.Tensor,
    alpha: torch.Tensor,
    y: torch.Tensor,
    config: Config,
) -> tuple[torch.Tensor, dict[str, float]]:
    dist = Dirichlet(alpha)
    nll = -dist.log_prob(y).mean()
    prior_penalty = config.prior_lambda * _l2_penalty(model)
    loss = nll + prior_penalty
    return loss, {"nll": float(nll.item()), "prior_penalty": float(prior_penalty.item())}


def _evaluate(model: BayesianDirichletMAP, split: SplitData, config: Config) -> dict[str, float]:
    device = next(model.parameters()).device
    model.eval()
    loader = _loader(split, shuffle=False)
    preds = []
    total_loss = 0.0
    n_batches = 0
    with torch.no_grad():
        for x, y, prior, _prior_short in loader:
            x = x.to(device)
            y = y.to(device)
            prior = prior.to(device)
            mean, alpha, concentration = model(x, prior)
            loss, _parts = _loss(model, mean, alpha, y, config)
            preds.append(mean.cpu().numpy())
            total_loss += float(loss.item())
            n_batches += 1
    pred = np.vstack(preds)
    return {
        "loss": total_loss / max(n_batches, 1),
        "mae": base._mean_abs_error(split.y, pred),
        "jsd": base._jsd_mean(split.y, pred),
        "kl_proxy": float(np.mean(np.sum(split.y * (np.log(split.y + 1e-12) - np.log(pred + 1e-12)), axis=1))),
    }


def _train_single(
    train: SplitData,
    val: SplitData,
    config: Config,
) -> tuple[BayesianDirichletMAP, list[dict[str, float]], dict[str, float]]:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = BayesianDirichletMAP(input_dim=train.x.shape[1], out_dim=train.y.shape[1], init_concentration=config.init_concentration).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=config.lr)
    train_loader = _loader(train, shuffle=True)

    best_state = None
    best_val = float("inf")
    patience_left = PATIENCE
    history: list[dict[str, float]] = []

    for epoch in range(1, MAX_EPOCHS + 1):
        model.train()
        running = 0.0
        batches = 0
        for x, y, prior, _prior_short in train_loader:
            x = x.to(device)
            y = y.to(device)
            prior = prior.to(device)
            optimizer.zero_grad()
            mean, alpha, concentration = model(x, prior)
            loss, parts = _loss(model, mean, alpha, y, config)
            loss.backward()
            optimizer.step()
            running += float(loss.item())
            batches += 1

        val_metrics = _evaluate(model, val, config)
        history.append(
            {
                "epoch": epoch,
                "train_loss": running / max(batches, 1),
                "val_loss": val_metrics["loss"],
                "val_mae": val_metrics["mae"],
                "val_jsd": val_metrics["jsd"],
            }
        )
        if val_metrics["mae"] < best_val - 1e-6:
            best_val = val_metrics["mae"]
            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)
    final_val = _evaluate(model, val, config)
    return model, history, final_val


def _run_grid(
    train: SplitData,
    val: SplitData,
    configs: list[Config],
) -> tuple[BayesianDirichletMAP, list[dict[str, float]], dict[str, float], Config]:
    best = None
    best_model = None
    best_history = None
    best_val = None
    for config in configs:
        model, history, val_metrics = _train_single(train, val, config)
        score = val_metrics["mae"]
        if best is None or score < best:
            best = score
            best_model = model
            best_history = history
            best_val = val_metrics
            best_config = config
    assert best_model is not None and best_history is not None and best_val is not None
    return best_model, best_history, best_val, best_config


def _predict(model: BayesianDirichletMAP, split: SplitData) -> np.ndarray:
    device = next(model.parameters()).device
    model.eval()
    preds = []
    loader = _loader(split, shuffle=False)
    with torch.no_grad():
        for x, _y, prior, _prior_short in loader:
            mean, alpha, concentration = model(x.to(device), prior.to(device))
            preds.append(mean.cpu().numpy())
    return np.vstack(preds)


def _metrics(y_true: np.ndarray, pred_model: np.ndarray, pred_long: np.ndarray, pred_short: np.ndarray) -> dict[str, float]:
    return {
        "model_mae": base._mean_abs_error(y_true, pred_model),
        "season_baseline_mae": base._mean_abs_error(y_true, pred_long),
        "short8_baseline_mae": base._mean_abs_error(y_true, pred_short),
        "model_jsd": base._jsd_mean(y_true, pred_model),
        "season_baseline_jsd": base._jsd_mean(y_true, pred_long),
        "short8_baseline_jsd": base._jsd_mean(y_true, pred_short),
        "model_kl_proxy": float(np.mean(np.sum(y_true * (np.log(y_true + 1e-12) - np.log(pred_model + 1e-12)), axis=1))),
        "season_baseline_kl_proxy": float(np.mean(np.sum(y_true * (np.log(y_true + 1e-12) - np.log(pred_long + 1e-12)), axis=1))),
        "short8_baseline_kl_proxy": float(np.mean(np.sum(y_true * (np.log(y_true + 1e-12) - np.log(pred_short + 1e-12)), axis=1))),
    }


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]
    short_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[short_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:
        row["fecha"] = pd.Timestamp(row["fecha"]).strftime("%Y-%m-%d")
    return rows


def _build_attack_task() -> tuple[SplitData, SplitData, SplitData, pd.DataFrame, dict]:
    df, features, attack_targets = _load_attack_data()
    train_df = df[df["split"] == "train"].copy()
    train_idx, val_idx, val_start_date = base._build_temporal_validation(train_df)
    test_idx = df.index[df["split"] == "test"]
    x_scaled, scaler_bundle = base._standardize_features(features, train_idx)
    y = _smooth_target(df[attack_targets].to_numpy(dtype=np.float32))
    prior_long = train_model_mod._normalize_rows(df[[f"long_mean__actual_attack_share__{zone}" for zone in attack_mod.ZONE_ORDER]].to_numpy(dtype=float)).astype(np.float32)
    prior_short = train_model_mod._normalize_rows(df[[f"short_mean__actual_attack_share__{zone}" for zone in attack_mod.ZONE_ORDER]].to_numpy(dtype=float)).astype(np.float32)
    bundle = {
        "attack_targets": attack_targets,
        "val_start_date": val_start_date,
        "scaler_bundle": scaler_bundle,
    }
    return (
        _make_split(df, train_idx, x_scaled, y, prior_long, prior_short),
        _make_split(df, val_idx, x_scaled, y, prior_long, prior_short),
        _make_split(df, test_idx, x_scaled, y, prior_long, prior_short),
        df,
        bundle,
    )


def _build_pv_task() -> tuple[SplitData, SplitData, SplitData, pd.DataFrame, dict]:
    df, features, y_dist = _load_pv_data()
    train_df = df[df["split"] == "train"].copy()
    train_idx, val_idx, val_start_date = base._build_temporal_validation(train_df)
    test_idx = df.index[df["split"] == "test"]
    x_scaled, scaler_bundle = base._standardize_features(features, train_idx)
    y = _smooth_target(y_dist.astype(np.float32))
    raw_long = df[[f"long_mean__zone_pvAdded__{zone}" for zone in pv_base.ZONE_ORDER]].to_numpy(dtype=float)
    raw_short = df[[f"short_mean__zone_pvAdded__{zone}" for zone in pv_base.ZONE_ORDER]].to_numpy(dtype=float)
    prior_long = train_model_mod._normalize_rows(np.clip(raw_long, 0.0, None)).astype(np.float32)
    prior_short = train_model_mod._normalize_rows(np.clip(raw_short, 0.0, None)).astype(np.float32)
    bundle = {
        "val_start_date": val_start_date,
        "scaler_bundle": scaler_bundle,
    }
    return (
        _make_split(df, train_idx, x_scaled, y, prior_long, prior_short),
        _make_split(df, val_idx, x_scaled, y, prior_long, prior_short),
        _make_split(df, test_idx, x_scaled, y, prior_long, prior_short),
        df,
        bundle,
    )


def _report_html(summary: dict) -> str:
    def rows_html(rows: list[dict[str, float | str]]) -> str:
        return "".join(
            f"<tr><td>{html.escape(str(r['fecha']))}</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>Bayesian Dirichlet MAP</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>Dirichlet MAP con prior de partido</h1>
    <p class="lead">Modelo bayesiano práctico: likelihood Dirichlet sobre la distribución observada, media a priori centrada en el baseline de temporada del partido, y priors gaussianos sobre los pesos. Los features ajustan esa distribución base vía correcciones lineales en logit-space.</p>
    <div class="grid2">
      <section class="card">
        <h2>Ataque - test</h2>
        <table>
          <tr><th>Metrica</th><th>Modelo</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>
        <p>Config elegida: {html.escape(summary['attack']['best_config']['name'])}</p>
      </section>
      <section class="card">
        <h2>PV - test</h2>
        <table>
          <tr><th>Metrica</th><th>Modelo</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>
        <p>Config elegida: {html.escape(summary['pv']['best_config']['name'])}</p>
      </section>
    </div>
    <div class="grid2">
      <section class="card">
        <h2>Ataque - ultimos 3 de Racing</h2>
        <table>
          <tr><th>Fecha</th><th>Rival</th><th>Modelo</th><th>Temporada</th><th>Ultimos 8</th><th>GNN previo</th></tr>
          {rows_html(summary['attack']['last3_racing'])}
        </table>
      </section>
      <section class="card">
        <h2>PV - ultimos 3 de Racing</h2>
        <table>
          <tr><th>Fecha</th><th>Rival</th><th>Modelo</th><th>Temporada</th><th>Ultimos 8</th><th>GNN previo</th></tr>
          {rows_html(summary['pv']['last3_racing'])}
        </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_train, attack_val, attack_test, attack_df, attack_bundle = _build_attack_task()
    attack_configs = [
        Config("a_soft_prior", lr=5e-3, prior_lambda=1e-5, init_concentration=20.0),
        Config("a_medium_prior", lr=3e-3, prior_lambda=3e-5, init_concentration=30.0),
        Config("a_strong_prior", lr=2e-3, prior_lambda=1e-4, init_concentration=40.0),
    ]
    attack_model, attack_history, attack_val_metrics, best_attack_config = _run_grid(attack_train, attack_val, attack_configs)
    attack_test_pred = _predict(attack_model, attack_test)
    attack_test_metrics = _metrics(attack_test.y, attack_test_pred, attack_test.prior, attack_test.prior_short)

    attack_pred_df = attack_test.metadata.copy().reset_index(drop=True)
    attack_targets = attack_bundle["attack_targets"]
    for i, zone in enumerate(attack_mod.ZONE_ORDER):
        target_col = attack_targets[attack_targets.index(f"target_attack_share__{zone}")]
        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.prior[:, i]
        attack_pred_df[f"pred_short8__target_attack_share__{zone}"] = attack_test.prior_short[:, i]
    attack_pred_df.to_parquet(ATTACK_PRED_PATH, index=False)

    pv_train, pv_val, pv_test, pv_df, pv_bundle = _build_pv_task()
    pv_configs = [
        Config("pv_soft_prior", lr=5e-3, prior_lambda=1e-5, init_concentration=18.0),
        Config("pv_medium_prior", lr=3e-3, prior_lambda=3e-5, init_concentration=28.0),
        Config("pv_strong_prior", lr=2e-3, prior_lambda=1e-4, init_concentration=36.0),
    ]
    pv_model, pv_history, pv_val_metrics, best_pv_config = _run_grid(pv_train, pv_val, pv_configs)
    pv_test_pred = _predict(pv_model, pv_test)
    pv_test_metrics = _metrics(pv_test.y, pv_test_pred, pv_test.prior, pv_test.prior_short)

    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.prior[:, i]
        pv_pred_df[f"pred_short8__target_pv_dist__{zone}"] = pv_test.prior_short[:, i]
    pv_pred_df.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 = {
        "model": "bayesian_dirichlet_map",
        "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": attack_bundle["val_start_date"],
            "test_metrics": attack_test_metrics,
            "val_metrics": attack_val_metrics,
            "previous_test_metrics": attack_prev_metrics,
            "best_config": vars(best_attack_config),
            "last3_racing": _last3_rows(attack_last3),
        },
        "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_bundle["val_start_date"],
            "test_metrics": pv_test_metrics,
            "val_metrics": pv_val_metrics,
            "previous_test_metrics": pv_prev_metrics,
            "best_config": vars(best_pv_config),
            "last3_racing": _last3_rows(pv_last3),
        },
        "artifacts": {
            "attack_predictions": str(ATTACK_PRED_PATH),
            "pv_predictions": str(PV_PRED_PATH),
        },
    }

    torch.save(
        {
            "attack_state_dict": attack_model.state_dict(),
            "pv_state_dict": pv_model.state_dict(),
            "attack_best_config": vars(best_attack_config),
            "pv_best_config": vars(best_pv_config),
            "summary": summary,
        },
        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()