leaderboard_compute.py

"""Compute the imputation leaderboard table live from the per-user substrate.

Downloads each method's per-user error substrate from the OpenMHC leaderboard
dataset, then reuses the canonical point-flow reducers from ``openmhc`` to
produce paired skill scores, cross-method ranks, and disparity-ratio (MAPD)
fair skill scores. No bootstrap (point only).

LOCF is the Track-2 baseline, so when it is the only method present the skill
table is empty by construction (baseline-vs-self) — the baseline row is emitted
with skill 0.0 and rank 1.
"""

from __future__ import annotations

import glob
import json
import os

import pandas as pd

DATASET_REPO = "MyHeartCounts/OpenMHC-leaderboard-data"
BASELINE = "locf"

HEADLINE_SCOPE = "overall"
# skill-score scope -> leaderboard subgroup column
SUBGROUP_SCOPES = {
    "cat:activity": "activity",
    "cat:physiology": "physiology",
    "cat:sleep": "sleep",
    "cat:workouts": "workout",
    "semantic": "semantic",
}
# Columns identifying one (method, task, subgroup) cell. The fairness reducer
# consumes per-cell MEAN errors, not raw per-user rows.
PER_CELL_COLS = [
    "method",
    "scenario",
    "channel",
    "channel_type",
    "subgroup_attr",
    "subgroup_value",
]

# Hardcoded HF model links, keyed by internal method id. Only methods with a
# published checkpoint appear here — statistical baselines (locf / mean / mode /
# linear / temporal_* / personalized_*) have none, and the forecasting (-fc)
# models belong to a different track.
MODEL_URLS = {
    "brits": "https://huggingface.co/MyHeartCounts/openmhc-brits-imp",
    "dlinear": "https://huggingface.co/MyHeartCounts/openmhc-dlinear-imp",
    "dlinear_weekly": "https://huggingface.co/MyHeartCounts/openmhc-dlinear-7day-imp",
    "fedformer": "https://huggingface.co/MyHeartCounts/openmhc-fedformer-imp",
    "lsm2": "https://huggingface.co/MyHeartCounts/openmhc-lsm2-daily",
    "lsm2_weekly": "https://huggingface.co/MyHeartCounts/openmhc-lsm2-weekly",
    "lsm2_weekly_sparse": "https://huggingface.co/MyHeartCounts/openmhc-lsm2-weekly-sparse",
    "timesnet": "https://huggingface.co/MyHeartCounts/openmhc-timesnet-imp",
}
def load_substrate(track: str = "imputation") -> tuple[pd.DataFrame, str]:
    """Download + concat every method's per-user substrate parquet for a track.

    Returns ``(df, local_root)``; the per-method display sidecars
    (``<method>.meta.json``) are downloaded alongside into ``local_root``.
    """
    from huggingface_hub import snapshot_download

    root = snapshot_download(
        DATASET_REPO,
        repo_type="dataset",
        allow_patterns=[f"{track}/*.parquet", f"{track}/*.meta.json"],
    )
    files = sorted(glob.glob(os.path.join(root, track, "*.parquet")))
    if not files:
        raise RuntimeError(f"No substrate under {track}/ in {DATASET_REPO}")
    df = pd.concat((pd.read_parquet(f) for f in files), ignore_index=True)
    return df, root


def read_method_meta(root: str, track: str, method: str) -> tuple[str, str, str, str, float | None]:
    """Return ``(display_name, type, submitter, subtrack, fallback_rate)`` from the sidecar.

    ``fallback_rate`` is the fraction of predictions substituted by the fallback
    baseline (``None`` when the sidecar omits it). Falls back to
    ``(raw_id, "—", "—", "other", None)`` when no sidecar is present.
    """
    path = os.path.join(root, track, f"{method}.meta.json")
    if os.path.exists(path):
        with open(path) as f:
            m = json.load(f)
        fr = m.get("fallback_rate")
        return (
            m.get("display_name", method),
            m.get("type", "—"),
            m.get("submitter", "—"),
            m.get("subtrack", "other"),
            float(fr) if isinstance(fr, (int, float)) else None,
        )
    return method, "—", "—", "other", None


def compute_imputation_rows() -> list[dict]:
    """Return leaderboard rows (sorted by skill desc) for the imputation track."""
    from imputation_evaluation.evaluation.bootstrap_skill_rank import (
        compute_per_task_paired_R,
    )
    from imputation_evaluation.evaluation.paper_metrics_core import (
        compute_average_rankings,
        compute_fair_skill_scores,
        compute_skill_scores,
    )

    df_full, root = load_substrate("imputation")
    df = df_full.rename(columns={"E_per_user": "E"})
    # skill + rank use the global ("all") cell; demographic cells are for fairness.
    df_all = df[df["subgroup_attr"] == "all"].copy()
    methods = sorted(df_all["method"].unique())

    per_task_r = compute_per_task_paired_R(df_all, baseline_method=BASELINE)
    skill = compute_skill_scores(per_task_r, mode="paired")  # [method, scope, skill_score, ...]
    ranks = compute_average_rankings(df_all)  # [method, scope, avg_rank, ...]
    # Disparity-ratio fair skill (MAPD across age_group + sex). The reducer
    # expects per-cell mean errors; feeding raw per-user rows triggers a
    # cartesian-merge blow-up, so collapse users to the cell mean first.
    per_cell = df.groupby(PER_CELL_COLS, observed=True)["E"].mean().reset_index()
    fair = compute_fair_skill_scores(per_cell, baseline_method=BASELINE)  # [method, scope, fair_skill_score, ...]

    def skill_at(method: str, scope: str) -> float | None:
        hit = skill[(skill["method"] == method) & (skill["scope"] == scope)]
        if not hit.empty:
            return float(hit["skill_score"].iloc[0])
        return 0.0 if method == BASELINE else None  # baseline-vs-self -> 0

    def rank_at(method: str, scope: str) -> float | None:
        hit = ranks[(ranks["method"] == method) & (ranks["scope"] == scope)]
        return float(hit["avg_rank"].iloc[0]) if not hit.empty else None

    def fair_at(method: str, scope: str) -> float | None:
        hit = fair[(fair["method"] == method) & (fair["scope"] == scope)]
        if not hit.empty:
            return float(hit["fair_skill_score"].iloc[0])
        return 0.0 if method == BASELINE else None  # baseline-vs-self -> 0

    rows = []
    for m in methods:
        name, mtype, submitter, subtrack, fallback = read_method_meta(root, "imputation", m)
        rows.append(
            {
                "method": name,
                "mtype": mtype,
                "submitter": submitter,
                "subtrack": subtrack,
                "model_url": MODEL_URLS.get(m),
                "skill": skill_at(m, HEADLINE_SCOPE),
                "fair_skill": fair_at(m, HEADLINE_SCOPE),
                "rank": rank_at(m, HEADLINE_SCOPE),
                "fallback": fallback,
                **{col: skill_at(m, scope) for scope, col in SUBGROUP_SCOPES.items()},
            }
        )
    rows.sort(key=lambda r: (r["skill"] is None, -(r["skill"] or 0.0)))
    return rows


# ---------------------------------------------------------------------------
# Track 3 — forecasting
# ---------------------------------------------------------------------------

FC_BASELINE = "seasonal_naive"
# Published HF checkpoints (MyHeartCounts org), keyed by internal method id. Only
# methods with a trained MyHeartCounts checkpoint appear: the deep-learning models
# and the *fine-tuned* foundation variants. Zero-shot foundation runs use the
# off-the-shelf base model (no MyHeartCounts checkpoint), and the statistical
# baselines (seasonal_naive / autoARIMA / autoETS) have none.
FC_MODEL_URLS: dict[str, str] = {
    "chronos2_finetuned": "https://huggingface.co/MyHeartCounts/openmhc-chronos2-fc",
    "toto_finetuned_ctx4096": "https://huggingface.co/MyHeartCounts/openmhc-toto-fc",
    "dlinear": "https://huggingface.co/MyHeartCounts/openmhc-dlinear-fc",
    "mixlinear": "https://huggingface.co/MyHeartCounts/openmhc-mixlinear-fc",
    "segrnn": "https://huggingface.co/MyHeartCounts/openmhc-segrnn-fc",
}


def _load_demographics() -> dict[str, dict[str, str]]:
    """Build ``{user_id: {age_group, sex}}`` from the imputation substrate.

    The forecasting per-user substrate carries no demographic columns, but every
    forecasting user is also in the imputation substrate (which bakes per-user
    ``age_group`` / ``sex`` subgroup rows under the canonical 18/30/40/50/60 age
    bins). Reusing that partition lets forecasting fairness use the *identical*
    user->subgroup mapping as imputation fairness, with no extra artifact.
    """
    from huggingface_hub import hf_hub_download

    path = hf_hub_download(DATASET_REPO, "imputation/locf.parquet", repo_type="dataset")
    imp = pd.read_parquet(path, columns=["subgroup_attr", "subgroup_value", "user_id"])
    demo: dict[str, dict[str, str]] = {}
    for attr in ("age_group", "sex"):
        sub = imp[imp["subgroup_attr"].astype(str) == attr][
            ["user_id", "subgroup_value"]
        ].drop_duplicates()
        for uid, val in sub.itertuples(index=False):
            demo.setdefault(str(uid), {})[attr] = str(val)
    return demo


def compute_forecasting_rows() -> list[dict]:
    """Return leaderboard rows (sorted by skill desc) for the forecasting track."""
    from forecasting_evaluation.metrics import metric_spec as spec
    from forecasting_evaluation.metrics.fair_skill_score import (
        compute_fair_skill_scores_from_errors,
    )
    from forecasting_evaluation.metrics.grouped_metric_rank_summary import (
        build_grouped_metric_rank_tables,
    )
    from forecasting_evaluation.metrics.per_user_errors import to_error_df
    from forecasting_evaluation.metrics.skill_score_summary import compute_skill_score_tables

    df, root = load_substrate("forecasting")
    methods = sorted(df["model"].astype(str).unique())
    models = {m: {"path": "", "display_name": m} for m in methods}

    # Skill: overall (category-balanced) + per-sensor-category scopes, vs the
    # Seasonal Naive baseline. summary_df is one row per model with *_score cols.
    _, summary, _ = compute_skill_score_tables(
        models=models,
        baseline_model=FC_BASELINE,
        continuous_metrics=list(spec.PAPER_CONTINUOUS_METRICS),
        binary_metrics=list(spec.PAPER_BINARY_METRICS),
        continuous_channel_indices=spec.CONTINUOUS_CHANNELS,
        binary_channel_indices=spec.BINARY_CHANNELS,
        clip_lower=spec.PAPER_CLIP_LOWER,
        clip_upper=spec.PAPER_CLIP_UPPER,
        min_pairs=spec.PAPER_MIN_PAIRS,
        aggregation_unit="user",
        within_user_aggregation="micro",
        per_user_metrics=df,
    )
    summary["model"] = summary["model"].astype(str)
    summary = summary.set_index("model")

    # Cross-method average rank; the headline is the category-balanced "overall".
    _, rank_long, _ = build_grouped_metric_rank_tables(
        models=models,
        continuous_metrics=list(spec.PAPER_CONTINUOUS_METRICS),
        binary_metrics=list(spec.PAPER_BINARY_METRICS),
        continuous_channel_indices=spec.CONTINUOUS_CHANNELS,
        binary_groups=[(name, tuple(idx)) for name, idx in spec.BINARY_GROUPS],
        within_user_aggregation="micro",
        per_user_metrics=df,
    )
    rank_long["model"] = rank_long["model"].astype(str)
    overall_rank = rank_long[rank_long["scope"] == "overall"].set_index("model")["rank"]

    # Disparity-ratio (MAPD) fair skill across age_group + sex. Demographics come
    # from the imputation substrate (same user->subgroup partition); the reducer
    # is the canonical forecasting one (point estimate == leaderboard value).
    err = to_error_df(df, user_col="user_id")
    fair = compute_fair_skill_scores_from_errors(
        err, _load_demographics(), baseline_method=FC_BASELINE
    )
    fair_overall = fair[fair["scope"] == "overall"].set_index("model")["fair_skill_score"]

    def num(value) -> float | None:
        if value is None or (isinstance(value, float) and pd.isna(value)):
            return None
        return float(value)

    rows = []
    for m in methods:
        name, mtype, submitter, subtrack, fallback = read_method_meta(root, "forecasting", m)
        srow = summary.loc[m] if m in summary.index else None
        rows.append(
            {
                "method": name,
                "mtype": mtype,
                "submitter": submitter,
                "subtrack": subtrack,
                "model_url": FC_MODEL_URLS.get(m),
                "skill": num(srow["overall_score"]) if srow is not None else None,
                "fair_skill": num(fair_overall.get(m)),
                "rank": num(overall_rank.get(m)),
                "activity": num(srow["activity_score"]) if srow is not None else None,
                "physiology": num(srow["physiology_score"]) if srow is not None else None,
                "sleep": num(srow["sleep_score"]) if srow is not None else None,
                "workout": num(srow["workout_score"]) if srow is not None else None,
                "fallback": num(fallback),
            }
        )
    rows.sort(key=lambda r: (r["skill"] is None, -(r["skill"] or 0.0)))
    return rows


# ---------------------------------------------------------------------------
# Track 1 — predictive tasks (downstream)
# ---------------------------------------------------------------------------

DOWNSTREAM_BASELINE = "linear"
# skill-score scope -> leaderboard subgroup column
DOWNSTREAM_SUBGROUP_SCOPES = {
    "Demographics": "demographics",
    "Medical conditions": "conditions",
    "Body metrics and biomarkers": "vitals",
    "Mental well-being": "mental",
    "Sleep and lifestyle": "lifestyle",
}
# Display grouping + model link per OpenMHC-team method. The substrate sidecar's `type`
# is the raw model family; the board groups methods into these display categories. A
# new submitter's numbers come straight from the reduced substrate; their label/link
# default to their own sidecar `type` + no url.
DOWNSTREAM_META = {
    "lsm2": ("Wearable Foundation Models", "https://huggingface.co/MyHeartCounts/openmhc-lsm2-daily"),
    "wbm": ("Wearable Foundation Models", "https://huggingface.co/MyHeartCounts/openmhc-wbm-dp"),
    "gru_d": ("Supervised Neural Models", None),
    "xgboost": ("Statistical Models", None),
    "multirocket": ("Statistical Models", None),
    "linear": ("Statistical Models", None),
    "chronos2": ("Time-Series Foundation Models", "https://huggingface.co/MyHeartCounts/openmhc-chronos2-fc"),
    "toto": ("Time-Series Foundation Models", "https://huggingface.co/MyHeartCounts/openmhc-toto-fc"),
}


def compute_downstream_rows() -> list[dict]:
    """Return leaderboard rows (sorted by skill desc) for the Track-1 (predictive tasks) track.

    Reduces the per-user-pairs substrate live — point skill / rank / disparity-ratio
    fairness vs the Linear baseline — mirroring the imputation/forecasting tracks, so a
    newly-uploaded substrate appears without a maintainer regenerating a precomputed
    rows file. Point only (the bootstrap CIs live under downstream/bootstrap/); reuses
    the paper's reducers, so the board equals the paper's point estimate.
    """
    from downstream_evaluation.evaluation.bootstrap_skill_rank import reduce_substrate_to_point

    df, root = load_substrate("downstream")
    point = reduce_substrate_to_point(df, baseline=DOWNSTREAM_BASELINE)

    rows = []
    for m in sorted(point):
        name, meta_type, submitter, subtrack, fallback = read_method_meta(root, "downstream", m)
        mtype, model_url = DOWNSTREAM_META.get(m, (meta_type, None))
        skill = point[m]["skill"]
        rows.append({
            "method": name,
            "mtype": mtype,
            "submitter": submitter,
            "subtrack": subtrack,
            "model_url": model_url,
            "skill": skill.get("Overall"),
            "fair_skill": point[m]["fair_skill"],
            "rank": point[m]["rank"].get("Overall"),
            "fallback": fallback,
            **{col: skill.get(scope) for scope, col in DOWNSTREAM_SUBGROUP_SCOPES.items()},
        })
    rows.sort(key=lambda r: (r["skill"] is None, -(r["skill"] or 0.0)))
    return rows


if __name__ == "__main__":
    import json

    print(json.dumps({
        "downstream": compute_downstream_rows(),
        "imputation": compute_imputation_rows(),
        "forecasting": compute_forecasting_rows(),
    }, indent=2))