src/data.py

import os

import pandas as pd
from dotenv import load_dotenv
from huggingface_hub import HfApi, hf_hub_download

from src.constants import TASK_SEQ_LENGTHS, TASK_TO_CATEGORY

load_dotenv()

HF_TOKEN = os.environ.get("HF_TOKEN")
LEADERBOARD_DATASET = os.environ.get("LEADERBOARD_DATASET", "lokahq/genomic-benchmark-metrics")


def load_data() -> pd.DataFrame:
    """Load all model parquet files from the HF Hub dataset and return a concatenated DataFrame.

    Each file under data/*.parquet corresponds to one model. Derived columns task_category and
    task_seq_length are added from the constants maps.
    """
    api = HfApi(token=HF_TOKEN)
    files = [
        f
        for f in api.list_repo_files(LEADERBOARD_DATASET, repo_type="dataset")
        if f.startswith("results/") and f.endswith(".parquet")
    ]
    frames = [
        pd.read_parquet(hf_hub_download(LEADERBOARD_DATASET, f, repo_type="dataset", token=HF_TOKEN))
        for f in files
    ]
    df = pd.concat(frames, ignore_index=True)

    # Normalize task names to lowercase so TASK_TO_CATEGORY mapping works regardless
    # of how the benchmark runner capitalizes them (e.g. "H3K27ac" → "h3k27ac")
    df["task_name"] = df["task_name"].str.lower()

    df["task_category"] = df["task_name"].map(TASK_TO_CATEGORY).fillna("other")
    df["task_seq_length"] = df["task_name"].map(TASK_SEQ_LENGTHS)

    if "run_at" in df.columns:
        df["run_at"] = pd.to_datetime(df["run_at"], utc=True, errors="coerce")

    return df


def _subsample_priority(val) -> float:
    """Map subsample_train value to sort priority. NULL/0/1 = production = highest priority."""
    if val is None or (isinstance(val, float) and pd.isna(val)) or val == 0 or val == 1:
        return float("inf")
    return float(val)


def deduplicate(df: pd.DataFrame) -> pd.DataFrame:
    """Keep best run per (embedding_config_hash, head_type).

    embedding_config_hash is task-specific: it hashes (model, pooling, layer_selection, dataset, task).
    Two rows for the same config but different tasks will have different hashes, so dedup correctly
    keeps one row per (config, task, head) — not one row per config globally.
    'Best' = production run preferred (subsample_train NULL/0/1 → inf priority),
    then highest subsample fraction, then most recent run_at.
    """
    if "run_at" not in df.columns:
        return df
    dedup_cols = []
    if "embedding_config_hash" in df.columns:
        dedup_cols.append("embedding_config_hash")
    if "head_type" in df.columns:
        dedup_cols.append("head_type")
    if not dedup_cols:
        return df
    df = df.copy()
    df["_sub_priority"] = df.get("subsample_train", pd.Series(dtype=float)).map(_subsample_priority)
    df = df.sort_values(["_sub_priority", "run_at"], ascending=[False, False])
    return df.drop_duplicates(subset=dedup_cols).drop(columns=["_sub_priority"]).reset_index(drop=True)



def apply_filters(df: pd.DataFrame, filters: dict) -> pd.DataFrame:
    """Apply user-selected filters and return the filtered DataFrame."""
    result = df.copy()

    subsample_mode = filters.get("subsample_train")
    if subsample_mode == "full" and "subsample_train" in result.columns:
        result = result[result["subsample_train"].isna() | (result["subsample_train"] == 0)]
    elif subsample_mode == "subsampled" and "subsample_train" in result.columns:
        result = result[result["subsample_train"].notna() & (result["subsample_train"] > 0)]

    for col, min_key, max_key in [
        ("model_params", "model_params_min", "model_params_max"),
        ("embedding_dim", "embedding_dim_min", "embedding_dim_max"),
        ("max_context_size", "max_context_size_min", "max_context_size_max"),
    ]:
        if col not in result.columns:
            continue
        lo = filters.get(min_key)
        hi = filters.get(max_key)
        if lo is not None:
            result = result[result[col] >= lo]
        if hi is not None:
            result = result[result[col] <= hi]

    excluded_models = filters.get("exclude_models") or []
    if excluded_models:
        result = result[~result["model_alias"].isin(excluded_models)]

    pooling_strategies = filters.get("pooling_strategies") or []
    if pooling_strategies and "pooling_strategy" in result.columns:
        result = result[result["pooling_strategy"].isin(pooling_strategies)]

    excluded_runs = filters.get("exclude_run_ids") or []
    if excluded_runs and "run_id" in result.columns:
        result = result[~result["run_id"].isin(excluded_runs)]

    return result.reset_index(drop=True)