vendor/scripts/build_attack_prediction_dataset.py

from __future__ import annotations

from pathlib import Path
import json

import numpy as np
import pandas as pd

import build_head_to_head_report as base


PROJECT_ROOT = Path(__file__).resolve().parents[1]
ZONE_FEATURES_PATH = PROJECT_ROOT / "data" / "team_match_zone_features.parquet"
MATCH_FEATURES_PATH = PROJECT_ROOT / "data" / "match_features.parquet"
EVENTS_ROOT = PROJECT_ROOT / "data" / "processed" / "events_parquet"
OUTPUT_DIR = PROJECT_ROOT / "data" / "modeling"
OUTPUT_PATH = OUTPUT_DIR / "attack_prediction_dataset.parquet"
SUMMARY_PATH = OUTPUT_DIR / "attack_prediction_dataset_summary.json"

SHORT_WINDOW = 8
RACING_TEAM_ID = "bzkwzatvwahmbzok1ymm5vqa1"

TARGET_LEAGUES = [
    "Austrian Bundesliga",
    "Belgian Challenger Pro League",
    "Danish Superligaen",
    "French Ligue 2",
    "Italian Serie B",
    "Liga Profesional Argentina",
    "Polish Ekstraklasa",
    "Serbian Super Liga",
    "Spanish La Liga",
    "Spanish Segunda Division",
    "Swedish Allsvenskan",
]

OPTIONAL_MISSING_LEAGUES = {
    "Polish Ekstraklasa",
    "Swedish Allsvenskan",
}


def _parse_date(series: pd.Series) -> pd.Series:
    return pd.to_datetime(series.astype(str).str.replace("Z", "", regex=False), errors="coerce")


def _safe_share(df: pd.DataFrame, cols: list[str], prefix: str) -> pd.DataFrame:
    mat = df[cols].fillna(0.0).to_numpy(dtype=float)
    total = mat.sum(axis=1, keepdims=True)
    with np.errstate(divide="ignore", invalid="ignore"):
        share = np.where(total > 0, mat / total, 0.0)
    out = pd.DataFrame(index=df.index)
    for i, col in enumerate(cols):
        suffix = col.split("__", 1)[1]
        if suffix.endswith("__opp"):
            suffix = suffix[:-5]
        out[f"{prefix}{suffix}"] = share[:, i]
    return out


def _build_match_features_long(df_matches: pd.DataFrame) -> pd.DataFrame:
    records: list[dict] = []
    meta_cols = [
        "matchId",
        "fecha",
        "competencia",
        "temporada",
        "home_team_id",
        "away_team_id",
        "equipo_local",
        "equipo_visitante",
    ]
    for _, row in df_matches.iterrows():
        base_meta = {col: row.get(col) for col in meta_cols}
        for prefix, team_id_col, opp_id_col, team_name_col, opp_name_col, is_home in [
            ("home_", "home_team_id", "away_team_id", "equipo_local", "equipo_visitante", True),
            ("away_", "away_team_id", "home_team_id", "equipo_visitante", "equipo_local", False),
        ]:
            rec = {
                "matchId": base_meta["matchId"],
                "fecha": base_meta["fecha"],
                "league": base_meta["competencia"],
                "season": str(base_meta["temporada"]),
                "teamId": row.get(team_id_col),
                "opponent_team_id": row.get(opp_id_col),
                "team_name": row.get(team_name_col),
                "opponent_name": row.get(opp_name_col),
                "is_home": is_home,
            }
            for col, value in row.items():
                if not isinstance(col, str) or not col.startswith(prefix):
                    continue
                raw_name = col[len(prefix) :]
                if raw_name in {"team_id", "matchId", "fecha", "competencia", "temporada"}:
                    continue
                if pd.api.types.is_number(value) and not isinstance(value, bool):
                    rec[f"mf__{raw_name}"] = float(value)
            records.append(rec)

    long_df = pd.DataFrame.from_records(records)
    long_df["fecha"] = _parse_date(long_df["fecha"])
    long_df["teamId"] = long_df["teamId"].astype(str)
    long_df["opponent_team_id"] = long_df["opponent_team_id"].astype(str)
    return long_df


def _season_events_dir(league: str, season: str) -> Path:
    return EVENTS_ROOT / f"league={league.replace(' ', '%20')}" / f"season={season}"


def _build_name_to_id_maps(df_matches: pd.DataFrame) -> dict[tuple[str, str], dict[str, str]]:
    maps: dict[tuple[str, str], dict[str, str]] = {}
    if df_matches.empty:
        return maps

    temp = df_matches.copy()
    temp["temporada"] = temp["temporada"].astype(str)
    for (league, season), sub in temp.groupby(["competencia", "temporada"], dropna=False):
        name_to_id: dict[str, str] = {}
        for _, row in sub[["equipo_local", "home_team_id"]].drop_duplicates().iterrows():
            if pd.notna(row["equipo_local"]) and pd.notna(row["home_team_id"]):
                name_to_id[str(row["equipo_local"])] = str(row["home_team_id"])
        for _, row in sub[["equipo_visitante", "away_team_id"]].drop_duplicates().iterrows():
            if pd.notna(row["equipo_visitante"]) and pd.notna(row["away_team_id"]):
                name_to_id[str(row["equipo_visitante"])] = str(row["away_team_id"])
        maps[(str(league), str(season))] = name_to_id
    return maps


def _load_event_role_lookup(
    league_season_pairs: list[tuple[str, str]],
    name_to_id_maps: dict[tuple[str, str], dict[str, str]],
) -> pd.DataFrame:
    frames: list[pd.DataFrame] = []
    for league, season in league_season_pairs:
        events_dir = _season_events_dir(league, season)
        if not events_dir.exists():
            continue
        raw = pd.read_parquet(events_dir, columns=["matchId", "fecha", "source_path"])
        lookup = raw[["matchId", "fecha", "source_path"]].drop_duplicates("matchId").copy()
        if lookup.empty:
            continue
        lookup["league"] = league
        lookup["season"] = str(season)
        lookup["file_name"] = lookup["source_path"].astype(str).str.split("/").str[-1]
        lookup[["home_name", "away_name"]] = lookup["file_name"].str.extract(
            r"\d{4}-\d{2}-\d{2} - (.*) vs (.*)\.xlsx$"
        )
        name_to_id = name_to_id_maps.get((str(league), str(season)), {})
        lookup["evt_home_team_id"] = lookup["home_name"].map(name_to_id)
        lookup["evt_away_team_id"] = lookup["away_name"].map(name_to_id)
        frames.append(
            lookup[
                [
                    "matchId",
                    "fecha",
                    "league",
                    "season",
                    "home_name",
                    "away_name",
                    "evt_home_team_id",
                    "evt_away_team_id",
                ]
            ]
        )

    if not frames:
        return pd.DataFrame(
            columns=[
                "matchId",
                "fecha",
                "league",
                "season",
                "home_name",
                "away_name",
                "evt_home_team_id",
                "evt_away_team_id",
            ]
        )
    return pd.concat(frames, ignore_index=True)


def _load_zone_dataset() -> pd.DataFrame:
    df = pd.read_parquet(ZONE_FEATURES_PATH)
    df = base._add_match_results(df)
    df["season"] = df["season"].astype(str)
    df["fecha"] = _parse_date(df["fecha"])

    available_leagues = set(df["league"].dropna().astype(str).unique())
    requested_available = [league for league in TARGET_LEAGUES if league in available_leagues]
    missing = [league for league in TARGET_LEAGUES if league not in available_leagues]
    unexpected_missing = [league for league in missing if league not in OPTIONAL_MISSING_LEAGUES]
    if unexpected_missing:
        raise ValueError(f"Faltan ligas necesarias en zone_features: {unexpected_missing}")

    df = df[df["league"].isin(requested_available)].copy()
    return df


def _attach_opponent_row(df_zone: pd.DataFrame) -> pd.DataFrame:
    opp_cols = [
        "matchId",
        "teamId",
        "goals_for",
        "shots_total",
        "pvAdded_total",
        "zone_actions__Creativity_Zone",
        "zone_actions__Cross__Der_",
        "zone_actions__Cross__Izq_",
        "zone_actions__Cut_Back__Der_",
        "zone_actions__Cut_Back__Izq_",
        "zone_actions__Deep_Cross__Der_",
        "zone_actions__Deep_Cross__Izq_",
        "zone_actions__Half_Space__Der_",
        "zone_actions__Half_Space__Izq_",
        "zone_actions__Scoring_Zone",
        "zone_pvAdded__Creativity_Zone",
        "zone_pvAdded__Cross__Der_",
        "zone_pvAdded__Cross__Izq_",
        "zone_pvAdded__Cut_Back__Der_",
        "zone_pvAdded__Cut_Back__Izq_",
        "zone_pvAdded__Deep_Cross__Der_",
        "zone_pvAdded__Deep_Cross__Izq_",
        "zone_pvAdded__Half_Space__Der_",
        "zone_pvAdded__Half_Space__Izq_",
        "zone_pvAdded__Scoring_Zone",
    ]
    opp = df_zone[opp_cols].rename(columns={"teamId": "opponent_team_id"})
    merged = df_zone.merge(opp, on="matchId", how="left", suffixes=("", "__opp"))
    merged = merged[merged["teamId"] != merged["opponent_team_id"]].copy()
    return merged


def _build_base_dataset() -> tuple[pd.DataFrame, dict]:
    df_zone = _load_zone_dataset()
    df_zone = _attach_opponent_row(df_zone)

    zone_action_cols = sorted(c for c in df_zone.columns if c.startswith("zone_actions__") and not c.endswith("__opp"))
    zone_pv_cols = sorted(c for c in df_zone.columns if c.startswith("zone_pvAdded__") and not c.endswith("__opp") and "shots" not in c)
    press_cols = sorted(c for c in df_zone.columns if c.startswith("press_cnt__"))
    rec_cols = sorted(c for c in df_zone.columns if c.startswith("interception_cnt__"))

    attack_share = _safe_share(df_zone, zone_action_cols, "actual_attack_share__")
    pv_share = _safe_share(df_zone, zone_pv_cols, "actual_pv_share__")
    press_share = _safe_share(df_zone, press_cols, "actual_press_share__")
    rec_share = _safe_share(df_zone, rec_cols, "actual_recovery_share__")

    opp_zone_action_cols = [f"{c}__opp" for c in zone_action_cols]
    opp_zone_pv_cols = [f"{c}__opp" for c in zone_pv_cols]
    conceded_attack_share = _safe_share(df_zone, opp_zone_action_cols, "actual_conceded_attack_share__")
    conceded_pv_share = _safe_share(df_zone, opp_zone_pv_cols, "actual_conceded_pv_share__")

    df_zone = pd.concat(
        [df_zone, attack_share, pv_share, press_share, rec_share, conceded_attack_share, conceded_pv_share],
        axis=1,
    )

    summary = {
        "available_zone_leagues": sorted(df_zone["league"].dropna().astype(str).unique().tolist()),
        "rows_zone": int(len(df_zone)),
        "matches_zone": int(df_zone["matchId"].nunique()),
    }
    return df_zone, summary


def _merge_match_features(df_zone: pd.DataFrame) -> pd.DataFrame:
    if not MATCH_FEATURES_PATH.exists():
        return df_zone

    df_matches = pd.read_parquet(MATCH_FEATURES_PATH)
    df_matches["temporada"] = df_matches["temporada"].astype(str)
    df_matches = df_matches[df_matches["competencia"].isin(TARGET_LEAGUES)].copy()
    df_long = _build_match_features_long(df_matches)

    out = df_zone.merge(
        df_long,
        on=["matchId", "teamId"],
        how="left",
        suffixes=("", "__mf"),
    )
    out["team_name"] = out.get("team_name")
    out["opponent_name"] = out.get("opponent_name")
    out["is_home"] = out.get("is_home")
    out["opponent_team_id"] = out["opponent_team_id"].fillna(out.get("opponent_team_id__mf"))
    return out


def _fill_roles_from_events(df: pd.DataFrame) -> pd.DataFrame:
    if not MATCH_FEATURES_PATH.exists():
        return df

    df_matches = pd.read_parquet(MATCH_FEATURES_PATH)
    name_to_id_maps = _build_name_to_id_maps(df_matches)
    league_season_pairs = [
        (str(row["league"]), str(row["season"]))
        for _, row in df[["league", "season"]].drop_duplicates().iterrows()
    ]
    evt_lookup = _load_event_role_lookup(league_season_pairs, name_to_id_maps)
    if evt_lookup.empty:
        return df

    out = df.merge(
        evt_lookup[
            ["matchId", "home_name", "away_name", "evt_home_team_id", "evt_away_team_id"]
        ],
        on="matchId",
        how="left",
    )

    evt_is_home = pd.Series(pd.NA, index=out.index, dtype="boolean")
    mask_evt = out["evt_home_team_id"].notna()
    evt_is_home.loc[mask_evt] = (
        out.loc[mask_evt, "teamId"].astype(str) == out.loc[mask_evt, "evt_home_team_id"].astype(str)
    ).to_numpy()
    if "is_home" not in out.columns:
        out["is_home"] = pd.Series(pd.NA, index=out.index, dtype="boolean")
    else:
        out["is_home"] = out["is_home"].astype("boolean")
    out["is_home"] = out["is_home"].where(out["is_home"].notna(), evt_is_home)

    fill_team_name = pd.Series(pd.NA, index=out.index, dtype="object")
    fill_opponent_name = pd.Series(pd.NA, index=out.index, dtype="object")
    known_home = out["is_home"].notna()
    fill_team_name.loc[known_home] = np.where(
        out.loc[known_home, "is_home"],
        out.loc[known_home, "home_name"],
        out.loc[known_home, "away_name"],
    )
    fill_opponent_name.loc[known_home] = np.where(
        out.loc[known_home, "is_home"],
        out.loc[known_home, "away_name"],
        out.loc[known_home, "home_name"],
    )

    if "team_name" not in out.columns:
        out["team_name"] = np.nan
    if "opponent_name" not in out.columns:
        out["opponent_name"] = np.nan

    out["team_name"] = out["team_name"].fillna(fill_team_name)
    out["opponent_name"] = out["opponent_name"].fillna(fill_opponent_name)
    return out


def _rolling_features(df: pd.DataFrame, source_cols: list[str], std_cols: list[str]) -> pd.DataFrame:
    df = df.sort_values(["league", "season", "teamId", "fecha", "matchId"]).copy()
    group_keys = ["league", "season", "teamId"]
    grouped = df.groupby(group_keys, sort=False)
    df["n_prior_matches"] = grouped.cumcount()
    rolling_frames: dict[str, pd.Series] = {}

    for col in source_cols:
        shifted = grouped[col].shift(1)
        rolling_frames[f"short_mean__{col}"] = (
            shifted.groupby([df["league"], df["season"], df["teamId"]], sort=False).transform(
                lambda s: s.rolling(SHORT_WINDOW, min_periods=1).mean()
            )
        )
        rolling_frames[f"long_mean__{col}"] = (
            shifted.groupby([df["league"], df["season"], df["teamId"]], sort=False).transform(
                lambda s: s.expanding(min_periods=1).mean()
            )
        )

    for col in std_cols:
        shifted = grouped[col].shift(1)
        rolling_frames[f"long_std__{col}"] = (
            shifted.groupby([df["league"], df["season"], df["teamId"]], sort=False).transform(
                lambda s: s.expanding(min_periods=2).std(ddof=0)
            )
        )

    return pd.concat([df, pd.DataFrame(rolling_frames, index=df.index)], axis=1)


def _add_opponent_histories(df: pd.DataFrame, hist_cols: list[str]) -> pd.DataFrame:
    opp = df[["matchId", "teamId", "n_prior_matches"] + hist_cols].copy()
    rename_map = {"teamId": "opponent_team_id", "n_prior_matches": "opp_n_prior_matches"}
    rename_map.update({col: f"opp__{col}" for col in hist_cols})
    opp = opp.rename(columns=rename_map)
    return df.merge(opp, on=["matchId", "opponent_team_id"], how="left")


def _assign_temporal_split(df: pd.DataFrame) -> tuple[pd.DataFrame, dict]:
    match_dates = (
        df[["matchId", "fecha"]]
        .drop_duplicates()
        .sort_values(["fecha", "matchId"])
        .reset_index(drop=True)
    )
    if match_dates.empty:
        raise ValueError("No hay partidos con fecha valida para asignar el split temporal.")
    split_idx = min(max(int(np.floor(len(match_dates) * 0.9)), 1), len(match_dates) - 1)
    default_cutoff = match_dates.loc[split_idx, "fecha"]

    racing_dates = (
        df[df["teamId"] == RACING_TEAM_ID][["matchId", "fecha"]]
        .drop_duplicates()
        .sort_values(["fecha", "matchId"])
        .reset_index(drop=True)
    )
    if len(racing_dates) >= 3:
        racing_cutoff = racing_dates.tail(3)["fecha"].min()
        cutoff = min(default_cutoff, racing_cutoff)
    else:
        cutoff = default_cutoff

    out = df.copy()
    out["split"] = np.where(out["fecha"] >= cutoff, "test", "train")

    summary = {
        "test_start_date": cutoff.strftime("%Y-%m-%d") if pd.notna(cutoff) else None,
        "train_rows": int((out["split"] == "train").sum()),
        "test_rows": int((out["split"] == "test").sum()),
        "train_matches": int(out.loc[out["split"] == "train", "matchId"].nunique()),
        "test_matches": int(out.loc[out["split"] == "test", "matchId"].nunique()),
    }
    return out, summary


def build_dataset() -> tuple[pd.DataFrame, dict]:
    df_zone, summary = _build_base_dataset()
    df = _merge_match_features(df_zone)
    df = _fill_roles_from_events(df)

    base_numeric_cols = []
    for col in df.columns:
        if col in {
            "matchId",
            "teamId",
            "league",
            "season",
            "fecha",
            "team_name",
            "opponent_name",
            "opponent_team_id",
            "is_home",
        }:
            continue
        if pd.api.types.is_numeric_dtype(df[col]):
            base_numeric_cols.append(col)

    target_attack_cols = sorted(c for c in df.columns if c.startswith("actual_attack_share__"))
    target_pv_cols = sorted(c for c in df.columns if c.startswith("actual_pv_share__"))
    pressure_hist_cols = sorted(
        c for c in df.columns if c.startswith("actual_press_share__") or c.startswith("actual_recovery_share__")
    )

    df = _rolling_features(df, source_cols=base_numeric_cols, std_cols=pressure_hist_cols)
    hist_cols = sorted(
        c for c in df.columns if c.startswith("short_mean__") or c.startswith("long_mean__") or c.startswith("long_std__")
    )
    df = _add_opponent_histories(df, hist_cols)

    rename_targets = {col: col.replace("actual_", "target_") for col in target_attack_cols + target_pv_cols}
    df = df.rename(columns=rename_targets)

    df, split_summary = _assign_temporal_split(df)

    df["usable_for_model"] = (
        (df["n_prior_matches"] >= 1)
        & (df["opp_n_prior_matches"].fillna(0) >= 1)
        & df["fecha"].notna()
    )

    summary.update(split_summary)
    summary["rows_total"] = int(len(df))
    summary["matches_total"] = int(df["matchId"].nunique())
    summary["leagues_total"] = sorted(df["league"].dropna().astype(str).unique().tolist())
    summary["rows_usable_for_model"] = int(df["usable_for_model"].sum())
    summary["feature_columns"] = int(
        sum(col.startswith("short_mean__") or col.startswith("long_mean__") or col.startswith("long_std__") or col.startswith("opp__") for col in df.columns)
    )
    summary["target_columns"] = int(sum(col.startswith("target_attack_share__") or col.startswith("target_pv_share__") for col in df.columns))

    racing_test = (
        df[(df["teamId"] == RACING_TEAM_ID) & (df["split"] == "test")][["matchId", "fecha"]]
        .drop_duplicates()
        .sort_values(["fecha", "matchId"])
        .tail(3)
    )
    summary["racing_last_3_test_matches"] = [
        {"matchId": row["matchId"], "fecha": row["fecha"].strftime("%Y-%m-%d")}
        for _, row in racing_test.iterrows()
    ]

    return df, summary


def main() -> None:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    df, summary = build_dataset()
    df.to_parquet(OUTPUT_PATH, index=False)
    SUMMARY_PATH.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")

    print(f"Dataset guardado en: {OUTPUT_PATH}")
    print(f"Resumen guardado en: {SUMMARY_PATH}")
    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()