src/tsfm_bench/eval/gift_eval_aggregation.py

from __future__ import annotations

import json
import re
from datetime import datetime, timezone
from pathlib import Path
from typing import Final

import numpy as np
import pandas as pd
from scipy import stats


MSE_COLUMN: Final = "eval_metrics/MSE[mean]"
CRPS_COLUMN: Final = "eval_metrics/mean_weighted_sum_quantile_loss"
DEFAULT_BASELINE_MODEL: Final = "Seasonal-Naive"
GEOMEAN_EPSILON: Final = 1e-12

CONFIG_COLUMNS: Final = ["dataset", "term_length", "frequency", "prediction_length"]
DOMAIN_DISPLAY: Final = {
    "Finance": "Finance",
    "Econ/Fin": "Finance",
    "Sales": "Finance",
    "Climate": "Climate",
    "Weather": "Climate",
    "Air Quality": "Climate",
    "Nature": "Climate",
    "Energy": "Climate",
    "Ocean": "Ocean",
    "Hydrology": "Hydrology",
    "Transport": "Transportation",
    "Transportation": "Transportation",
    "Web": "Web",
    "Web/CloudOps": "Web",
    "Hazards": "Hazards",
    "Healthcare": "Hazards",
}
AGGREGATE_COLUMNS: Final = [
    "model",
    "MSE",
    "CRPS",
    "Rank",
    "n_configs",
    "n_group_configs",
    "coverage",
]

RESULT_CONFIG_COLUMNS: Final = [
    "dataset",
    "model",
    "domain",
    MSE_COLUMN,
    CRPS_COLUMN,
    "prediction_length",
]

GROUP_SPECS: Final = {
    "prediction_length": ("prediction_length", "results_by_prediction_length.csv"),
    "domain": ("domain", "results_by_domain.csv"),
    "frequency": ("frequency", "results_by_frequency.csv"),
}


def _model_key(value: object) -> str:
    return re.sub(r"[^a-z0-9]+", "", str(value).lower())


def _load_model_allowlist(root: Path) -> set[str]:
    path = root / "model_allowlist.json"
    if not path.exists():
        return set()
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except (OSError, json.JSONDecodeError):
        return set()
    models = payload.get("models", payload) if isinstance(payload, dict) else payload
    if not isinstance(models, list):
        return set()
    return {_model_key(model) for model in models if str(model).strip()}


def _geometric_mean(series: pd.Series) -> float:
    values = pd.to_numeric(series, errors="coerce").dropna().to_numpy(dtype=float)
    values = values[np.isfinite(values) & (values >= 0.0)]
    if values.size == 0:
        return np.nan
    values = np.clip(values, GEOMEAN_EPSILON, None)
    return float(stats.gmean(values))


def _split_dataset_config(value: object) -> tuple[str, str, str]:
    parts = str(value).rsplit("/", 2)
    if len(parts) == 3:
        return parts[0], parts[1], parts[2]
    if len(parts) == 2:
        return parts[0], parts[1], "short"
    return parts[0], "unknown", "short"


def canonical_domain(value: object) -> str:
    domain = str(value).strip()
    if domain in DOMAIN_DISPLAY:
        return DOMAIN_DISPLAY[domain]
    lowered = domain.lower()
    if any(token in lowered for token in ("finance", "econ", "sales")):
        return "Finance"
    if any(token in lowered for token in ("climate", "weather", "air", "nature", "energy")):
        return "Climate"
    if "ocean" in lowered:
        return "Ocean"
    if "hydro" in lowered:
        return "Hydrology"
    if any(token in lowered for token in ("transport", "taxi")):
        return "Transportation"
    if any(token in lowered for token in ("web", "cloud", "wiki")):
        return "Web"
    return "Hazards"


def _load_prediction_lengths(model_dir: Path) -> dict[str, int]:
    """Load current numeric horizons from per-dataset forecast snapshots."""

    lengths: dict[str, int] = {}
    forecasts_dir = model_dir / "forecasts"
    if not forecasts_dir.exists():
        return lengths
    for snapshot_path in sorted(forecasts_dir.glob("*.json")):
        try:
            snapshot = json.loads(snapshot_path.read_text())
            dataset_config = str(snapshot["dataset"])
            prediction_length = int(snapshot["prediction_length"])
        except (OSError, ValueError, TypeError, KeyError, json.JSONDecodeError):
            continue
        if prediction_length > 0:
            lengths[dataset_config] = prediction_length
    return lengths


def _read_result_config_csv(path: Path) -> pd.DataFrame:
    return pd.read_csv(
        path,
        usecols=lambda column: column in RESULT_CONFIG_COLUMNS,
        index_col=False,
    )


def _dataset_key(value: object) -> str:
    return str(value).split("/")[0].strip().lower()


def _load_live_dataset_keys(root: Path) -> set[str]:
    path = root / "dataset_properties.csv"
    if not path.exists():
        return set()
    try:
        props = pd.read_csv(path)
    except (OSError, pd.errors.ParserError, pd.errors.EmptyDataError):
        return set()
    if "dataset" not in props.columns:
        return set()
    return {_dataset_key(value) for value in props["dataset"].dropna()}


def load_result_configs(root_dir: str | Path) -> pd.DataFrame:
    """Load one result row per model and GIFT-Eval-style dataset config."""

    root = Path(root_dir)
    frames: list[pd.DataFrame] = []
    live_dataset_keys = _load_live_dataset_keys(root)
    for csv_path in sorted(root.glob("*/all_results.csv")):
        try:
            frame = _read_result_config_csv(csv_path)
        except (OSError, pd.errors.ParserError, pd.errors.EmptyDataError):
            continue
        required = {"dataset", "model", "domain", MSE_COLUMN, CRPS_COLUMN}
        if frame.empty or not required.issubset(frame.columns):
            continue
        snapshot_lengths = _load_prediction_lengths(csv_path.parent)
        existing_lengths = (
            pd.to_numeric(frame["prediction_length"], errors="coerce")
            if "prediction_length" in frame.columns
            else pd.Series(np.nan, index=frame.index, dtype=float)
        )
        snapshot_series = frame["dataset"].astype(str).map(snapshot_lengths)
        frame["prediction_length"] = existing_lengths.fillna(snapshot_series)
        frames.append(frame)

    if not frames:
        return pd.DataFrame(
            columns=[
                *CONFIG_COLUMNS,
                "prediction_length",
                "config_id",
                "model",
                "domain",
                MSE_COLUMN,
                CRPS_COLUMN,
            ]
        )

    df = pd.concat(frames, ignore_index=True)
    configs = df["dataset"].map(_split_dataset_config)
    df[["dataset", "frequency", "term_length"]] = pd.DataFrame(
        configs.tolist(),
        index=df.index,
    )
    model_allowlist = _load_model_allowlist(root)
    if model_allowlist:
        df = df[df["model"].map(_model_key).isin(model_allowlist)].copy()
    if live_dataset_keys:
        df = df[df["dataset"].map(_dataset_key).isin(live_dataset_keys)].copy()
    df["frequency"] = df["frequency"].astype(str).str.strip()
    df["term_length"] = df["term_length"].astype(str).str.strip().str.lower()
    df["domain"] = df["domain"].map(canonical_domain)
    df["model"] = df["model"].astype(str).str.strip()
    df["prediction_length"] = pd.to_numeric(df["prediction_length"], errors="coerce").astype(
        "Int64"
    )
    df["config_id"] = df[CONFIG_COLUMNS].astype(str).agg("/".join, axis=1)
    for metric in (MSE_COLUMN, CRPS_COLUMN):
        df[metric] = pd.to_numeric(df[metric], errors="coerce")
        df.loc[~np.isfinite(df[metric]), metric] = np.nan

    # A model should contribute only once to a dataset configuration. Keeping
    # the last row makes the loader robust to accidentally concatenated files.
    return (
        df.dropna(subset=["dataset", "model"])
        .drop_duplicates(subset=[*CONFIG_COLUMNS, "model"], keep="last")
        .reset_index(drop=True)
    )


def normalize_against_seasonal_naive(
    df: pd.DataFrame,
    baseline_model: str = DEFAULT_BASELINE_MODEL,
) -> pd.DataFrame:
    """Apply the GIFT-Eval per-configuration baseline normalization.

    MSE and CRPS are divided by the Seasonal-Naive value from the same
    ``(dataset, term_length, frequency)`` configuration. Rank follows
    GIFT-Eval and is computed from CRPS before group aggregation.
    """

    if df.empty:
        result = df.copy()
        result["MSE"] = pd.Series(dtype=float)
        result["CRPS"] = pd.Series(dtype=float)
        result["Rank"] = pd.Series(dtype=float)
        return result

    baseline_key = _model_key(baseline_model)
    baseline_rows = df[df["model"].map(_model_key) == baseline_key]
    if baseline_rows.empty:
        available = ", ".join(sorted(df["model"].dropna().unique()))
        raise ValueError(
            f"Normalization baseline {baseline_model!r} was not found. Available models: {available}"
        )

    baseline = (
        baseline_rows.groupby(CONFIG_COLUMNS, as_index=False)[[MSE_COLUMN, CRPS_COLUMN]]
        .mean()
        .rename(columns={MSE_COLUMN: "__baseline_mse", CRPS_COLUMN: "__baseline_crps"})
    )
    normalized = df.merge(baseline, on=CONFIG_COLUMNS, how="left", validate="many_to_one")

    valid_mse = (
        np.isfinite(normalized["__baseline_mse"])
        & normalized["__baseline_mse"].ge(0)
        & np.isfinite(normalized[MSE_COLUMN])
        & normalized[MSE_COLUMN].ge(0)
    )
    valid_crps = (
        np.isfinite(normalized["__baseline_crps"])
        & normalized["__baseline_crps"].ge(0)
        & np.isfinite(normalized[CRPS_COLUMN])
        & normalized[CRPS_COLUMN].ge(0)
    )
    mse_denominator = normalized["__baseline_mse"].clip(lower=GEOMEAN_EPSILON)
    crps_denominator = normalized["__baseline_crps"].clip(lower=GEOMEAN_EPSILON)
    mse_numerator = normalized[MSE_COLUMN].clip(lower=GEOMEAN_EPSILON)
    crps_numerator = normalized[CRPS_COLUMN].clip(lower=GEOMEAN_EPSILON)
    normalized["MSE"] = np.where(
        valid_mse,
        mse_numerator / mse_denominator,
        np.nan,
    )
    normalized["CRPS"] = np.where(
        valid_crps,
        crps_numerator / crps_denominator,
        np.nan,
    )
    rank_metric = normalized[CRPS_COLUMN].where(
        normalized[CRPS_COLUMN].notna(),
        normalized[MSE_COLUMN],
    )
    normalized["Rank"] = rank_metric.groupby(
        [normalized[column] for column in CONFIG_COLUMNS],
        dropna=False,
    ).rank(
        method="average",
        ascending=True,
        na_option="keep",
    )
    return normalized.drop(columns=["__baseline_mse", "__baseline_crps"])


def aggregate_by(normalized: pd.DataFrame, group_column: str) -> pd.DataFrame:
    """Aggregate normalized configuration scores using the GIFT-Eval rule.

    Per-configuration ranks are averaged for ordering, then written as
    positive integer leaderboard ranks within each group.
    """

    output_columns = [group_column, *AGGREGATE_COLUMNS]
    if normalized.empty:
        return pd.DataFrame(columns=output_columns)
    if group_column not in normalized.columns:
        raise ValueError(f"Unknown aggregation column: {group_column}")

    source = normalized
    if group_column == "prediction_length":
        source = normalized.dropna(subset=[group_column])
    if source.empty:
        return pd.DataFrame(columns=output_columns)

    group_config_counts = source.groupby(group_column, dropna=False)["config_id"].nunique()
    grouped = (
        source.groupby([group_column, "model"], dropna=False)
        .agg(
            MSE=("MSE", _geometric_mean),
            CRPS=("CRPS", _geometric_mean),
            Rank=("Rank", "mean"),
            n_configs=("config_id", "nunique"),
        )
        .reset_index()
    )
    grouped["n_group_configs"] = grouped[group_column].map(group_config_counts).astype(int)
    grouped["coverage"] = grouped["n_configs"] / grouped["n_group_configs"]

    rank_sort_columns = ["Rank", "MSE", "model"]
    if group_column == "term_length":
        term_order = {"short": 0, "medium": 1, "long": 2}
        grouped["__order"] = grouped[group_column].map(term_order).fillna(99)
        grouped = grouped.sort_values(["__order", group_column, *rank_sort_columns]).drop(
            columns="__order"
        )
    elif group_column == "prediction_length":
        grouped = grouped.sort_values([group_column, *rank_sort_columns], na_position="last")
    else:
        grouped = grouped.sort_values([group_column, *rank_sort_columns], na_position="last")
    grouped["Rank"] = grouped.groupby(group_column, dropna=False).cumcount().add(1).astype(int)
    return grouped.loc[:, output_columns].reset_index(drop=True)


def build_aggregated_results(
    df: pd.DataFrame,
    baseline_model: str = DEFAULT_BASELINE_MODEL,
) -> dict[str, pd.DataFrame]:
    df = df.copy()
    if "domain" in df.columns:
        df["domain"] = df["domain"].map(canonical_domain)
    normalized = normalize_against_seasonal_naive(df, baseline_model)
    tables = {
        name: aggregate_by(normalized, column).rename(columns={column: name})
        for name, (column, _) in GROUP_SPECS.items()
    }
    overall_source = normalized.assign(overall="Overall")
    tables["overall"] = aggregate_by(overall_source, "overall")
    return tables


def _write_csv_atomic(df: pd.DataFrame, path: Path) -> None:
    tmp_path = path.with_suffix(path.suffix + ".tmp")
    df.to_csv(tmp_path, index=False)
    tmp_path.replace(path)


def write_aggregated_results(
    root_dir: str | Path,
    baseline_model: str = DEFAULT_BASELINE_MODEL,
) -> dict[str, Path]:
    """Build and persist current GIFT-style aggregate tables."""

    root = Path(root_dir)
    aggregate_dir = root / "aggregates"
    aggregate_dir.mkdir(parents=True, exist_ok=True)

    configs = load_result_configs(root)
    tables = build_aggregated_results(configs, baseline_model)
    paths: dict[str, Path] = {}
    for name, (_, filename) in GROUP_SPECS.items():
        path = aggregate_dir / filename
        _write_csv_atomic(tables[name], path)
        paths[name] = path

    overall_path = aggregate_dir / "results_overall.csv"
    _write_csv_atomic(tables["overall"], overall_path)
    paths["overall"] = overall_path

    metadata = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "baseline_model": baseline_model,
        "normalization": "Per-config metric / Seasonal-Naive metric",
        "metric_aggregation": "Geometric mean over dataset configurations",
        "rank_aggregation": (
            "Positive integer order by mean per-config CRPS rank; "
            "falls back to MSE when CRPS is unavailable"
        ),
        "config_key": CONFIG_COLUMNS,
        "num_configs": int(configs["config_id"].nunique()) if not configs.empty else 0,
        "num_models": int(configs["model"].nunique()) if not configs.empty else 0,
        "num_rows_missing_prediction_length": (
            int(configs["prediction_length"].isna().sum()) if not configs.empty else 0
        ),
    }
    metadata_path = aggregate_dir / "metadata.json"
    tmp_metadata_path = metadata_path.with_suffix(".json.tmp")
    tmp_metadata_path.write_text(json.dumps(metadata, indent=2) + "\n", encoding="utf-8")
    tmp_metadata_path.replace(metadata_path)
    paths["metadata"] = metadata_path
    return paths