src/display/columns.py

from __future__ import annotations

import pandas as pd

from src.display.i18n import t
from src.display.utils import AutoEvalColumn, COLS, fields

INTERNAL_MODEL_COL = AutoEvalColumn.model.name
INTERNAL_TEAM_COL = AutoEvalColumn.team.name
INTERNAL_DATASET_MODEL_COL = "model"
INTERNAL_DATASET_TEAM_COL = "team"
MODEL_COL = t("col_model")
TEAM_COL = t("col_team")

_NUMBER_COLS = {
    AutoEvalColumn.average.name,
    AutoEvalColumn.params.name,
    AutoEvalColumn.likes.name,
}
_BOOL_COLS = {
    AutoEvalColumn.still_on_hub.name,
}


def _rename_model_team_columns(
    df: pd.DataFrame,
    model_key: str,
    team_key: str,
) -> pd.DataFrame:
    rename: dict[str, str] = {}
    for source, target in (
        (model_key, MODEL_COL),
        (team_key, TEAM_COL),
        (INTERNAL_MODEL_COL, MODEL_COL),
        (INTERNAL_TEAM_COL, TEAM_COL),
        ("Model", MODEL_COL),
        ("Team", TEAM_COL),
    ):
        if source in df.columns and source != target:
            rename[source] = target
    if not rename:
        return df
    return df.rename(columns=rename)


def localize_main_leaderboard_df(df: pd.DataFrame) -> pd.DataFrame:
    return _rename_model_team_columns(df, INTERNAL_MODEL_COL, INTERNAL_TEAM_COL)


def localize_dataset_leaderboard_df(df: pd.DataFrame) -> pd.DataFrame:
    return _rename_model_team_columns(
        df, INTERNAL_DATASET_MODEL_COL, INTERNAL_DATASET_TEAM_COL
    )


def main_board_datatype_for_df(df: pd.DataFrame) -> list[str]:
    """Build Leaderboard datatypes aligned to the actual dataframe columns."""
    from src.about import Tasks

    number_cols = set(_NUMBER_COLS) | {task.value.col_name for task in Tasks}
    result: list[str] = []
    for col in df.columns:
        if col in number_cols:
            result.append("number")
        elif col in _BOOL_COLS:
            result.append("bool")
        else:
            result.append("str")
    return result


def main_board_datatype() -> list[str]:
    col_types = {c.name: c.type for c in fields(AutoEvalColumn)}
    return [col_types[name] for name in COLS]


def main_board_default_cols() -> list[str]:
    from src.about import Tasks

    return [
        MODEL_COL,
        TEAM_COL,
        AutoEvalColumn.average.name,
        *[task.value.col_name for task in Tasks],
    ]


def main_board_hidden_columns() -> list[str]:
    """Metadata columns hidden from the table and the column picker."""
    from src.about import Tasks

    keep_visible = {
        AutoEvalColumn.model_type_symbol.name,
        MODEL_COL,
        TEAM_COL,
        AutoEvalColumn.average.name,
        *[task.value.col_name for task in Tasks],
    }
    localized_cols = []
    for col in COLS:
        if col == INTERNAL_MODEL_COL:
            localized_cols.append(MODEL_COL)
        elif col == INTERNAL_TEAM_COL:
            localized_cols.append(TEAM_COL)
        else:
            localized_cols.append(col)
    return [col for col in localized_cols if col not in keep_visible]