utils.py

"""Helpers to load from and push rows to a HuggingFace dataset."""

from datetime import datetime, timezone

import pandas as pd
from datasets import Dataset, load_dataset

from config import HF_DATASET_PATH, REGISTRATION_COLUMNS, SUBMISSION_COLUMNS, TOKEN


# ── Internal helpers ──────────────────────────────────────────────────────────

def _load_subset(subset: str) -> pd.DataFrame:
    """Load one config/subset of the HF dataset as a DataFrame.
    Returns an empty DataFrame (with the right columns) if nothing exists yet.
    """
    columns = REGISTRATION_COLUMNS if subset == "registrations" else SUBMISSION_COLUMNS
    try:
        ds = load_dataset(HF_DATASET_PATH, subset, split="train", token=TOKEN)
        return ds.to_pandas()
    except Exception:
        return pd.DataFrame(columns=columns)


def _push_subset(df: pd.DataFrame, subset: str) -> None:
    """Push a DataFrame back to HF as a dataset config."""
    ds = Dataset.from_pandas(df, preserve_index=False)
    ds.push_to_hub(HF_DATASET_PATH, config_name=subset, split="train", token=TOKEN)


# ── Public API ────────────────────────────────────────────────────────────────

def load_registrations() -> pd.DataFrame:
    return _load_subset("registrations")


def add_registration(name: str, email: str, affiliation: str, team_name: str) -> None:
    df = load_registrations()
    new_row = {
        "timestamp":   datetime.now(timezone.utc).isoformat(),
        "name":        name,
        "email":       email,
        "affiliation": affiliation,
        "team_name":   team_name,
    }
    df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
    _push_subset(df, "registrations")


def load_submissions() -> pd.DataFrame:
    return _load_subset("submissions")


def add_submission(team_name: str, method: str, score: float, file_name: str) -> None:
    df = load_submissions()
    new_row = {
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "team_name": team_name,
        "method":    method,
        "score":     score,
        "file_name": file_name,
    }
    df = pd.concat([df, pd.DataFrame([new_row])], ignore_index=True)
    _push_subset(df, "submissions")