utils/column_rules.py

"""리더보드 컬럼 규칙.

벤치마크 컬럼은 `<데이터셋>_<카테고리>` 명명 규칙(예: Coupang_Fire, Soil_v2_Smoke,
AVG_Falldown)을 따른다. 이 규칙으로 df.columns 만 보고 카테고리/데이터셋 축을
자동 분류하고 표시/숨김 컬럼을 계산한다 — 신규 벤치마크가 등록되어도
config 목록을 손볼 필요가 없다 (기존 ON_LOAD_COLUMNS/HIDE_COLUMNS 대체).
"""
from typing import List, Optional, Sequence, Tuple

import pandas as pd

import enviroments.config as config

# AVG_* 컬럼의 prefix. 데이터셋 축에서는 가상 데이터셋으로 취급한다.
AVG_DATASET = "AVG"


def split_benchmark_column(column: str) -> Optional[Tuple[str, str]]:
    """`<데이터셋>_<카테고리>` 컬럼을 (데이터셋, 카테고리)로 분해.

    규칙 밖 컬럼(TASK, Model, "Model name" 등)은 None.
    """
    if "_" not in column or " " in column:
        return None
    dataset, _, category = column.rpartition("_")
    if not dataset or not category:
        return None
    return dataset, category


def build_hide_columns(columns: Sequence[str]) -> List[str]:
    """테이블에서 아예 제외할 컬럼 (숨김 카테고리 패턴 + 개별 지정)."""
    hidden = []
    for col in columns:
        if col in config.HIDE_COLUMN_EXACT:
            hidden.append(col)
            continue
        parsed = split_benchmark_column(col)
        if parsed and parsed[1] in config.HIDE_CATEGORIES:
            hidden.append(col)
    return hidden


def parse_axes(columns: Sequence[str]) -> Tuple[List[str], List[str]]:
    """df.columns → (카테고리 목록, 데이터셋 목록). 숨김 컬럼/AVG 데이터셋 제외."""
    hidden = set(build_hide_columns(columns))
    categories, datasets = set(), set()
    for col in columns:
        if col in hidden:
            continue
        parsed = split_benchmark_column(col)
        if parsed is None:
            continue
        dataset, category = parsed
        categories.add(category)
        if dataset != AVG_DATASET:
            datasets.add(dataset)
    return sorted(categories), sorted(datasets)


def extra_columns(columns: Sequence[str]) -> List[str]:
    """벤치마크 명명 규칙 밖이면서 기본/필수 컬럼도 아닌 부가 컬럼.

    예: "Datasets Used", "Infer Speed". 셀렉터의 '추가 정보' 축이 된다.
    """
    always_on = set(config.LEADERBOARD_BASE_COLUMNS) | set(config.OFF_LOAD_COLUMNS)
    hidden = set(build_hide_columns(columns))
    return [
        col
        for col in columns
        if col not in always_on
        and col not in hidden
        and split_benchmark_column(col) is None
    ]


def avg_categories(columns: Sequence[str]) -> List[str]:
    """AVG_<카테고리> 컬럼이 존재하는 카테고리 목록 (숨김 제외, df 순서)."""
    hidden = set(build_hide_columns(columns))
    cats = []
    for col in columns:
        if col in hidden:
            continue
        parsed = split_benchmark_column(col)
        if parsed and parsed[0] == AVG_DATASET:
            cats.append(parsed[1])
    return cats


def datasets_for_category(columns: Sequence[str], category: str) -> List[str]:
    """해당 카테고리 벤치마크(<데이터셋>_<카테고리>)를 가진 전체 데이터셋."""
    out = set()
    for col in columns:
        parsed = split_benchmark_column(col)
        if parsed and parsed[1] == category and parsed[0] != AVG_DATASET:
            out.add(parsed[0])
    return sorted(out)


def default_avg_map(columns: Sequence[str]) -> dict:
    """카테고리별 AVG 계산 기준의 기본값: 정식 세트 ∩ 해당 카테고리 보유 데이터셋."""
    return {
        cat: [
            d
            for d in config.DEFAULT_AVG_DATASETS
            if d in datasets_for_category(columns, cat)
        ]
        for cat in avg_categories(columns)
    }


def apply_avg_override(df: pd.DataFrame, avg_map: Optional[dict] = None) -> pd.DataFrame:
    """AVG_<카테고리> 컬럼을 카테고리별 기준 데이터셋의 행별 평균으로 대체.

    구글 시트의 AVG 원본은 건드리지 않고 화면 표시값만 교체한다.
    - avg_map: {카테고리: [데이터셋, ...]}. 카테고리가 없거나 빈 리스트면
      config.DEFAULT_AVG_DATASETS (정식 벤치마크 세트) 기준으로 계산.
    - 그룹 컬럼 값이 하나라도 비어 있는 모델은 빈칸 (완주한 모델만 AVG 표시).
    - 기준 데이터셋에 해당 카테고리 컬럼이 하나도 없으면 전체 빈칸.
    """
    avg_map = avg_map or {}
    df = df.copy()
    for col in df.columns:
        parsed = split_benchmark_column(col)
        if parsed is None or parsed[0] != AVG_DATASET:
            continue
        category = parsed[1]
        selected = avg_map.get(category) or config.DEFAULT_AVG_DATASETS
        members = [
            f"{d}_{category}" for d in selected if f"{d}_{category}" in df.columns
        ]
        if not members:
            df[col] = ""
            continue
        vals = df[members].apply(pd.to_numeric, errors="coerce")
        avg = vals.mean(axis=1)
        avg[vals.isna().any(axis=1)] = float("nan")
        df[col] = avg.map(lambda v: "" if pd.isna(v) else f"{v:.4f}")
    return df


def compute_display_columns(
    columns: Sequence[str],
    categories: Sequence[str],
    datasets: Sequence[str],
    extras: Sequence[str] = (),
) -> List[str]:
    """선택된 카테고리 × 데이터셋 교집합 + 부가 컬럼의 표시 목록 (df 컬럼 순서 유지).

    datasets 가 비어 있으면 전체 데이터셋(AVG 포함)으로 취급.
    기본/필수 컬럼(TASK, Model, Model name 등)은 항상 표시.
    """
    selected_cats = set(categories or [])
    selected_ds = set(datasets or [])
    selected_extras = set(extras or [])
    always_on = set(config.LEADERBOARD_BASE_COLUMNS) | set(config.OFF_LOAD_COLUMNS)
    hidden = set(build_hide_columns(columns))

    display = []
    for col in columns:
        if col in hidden:
            continue
        parsed = split_benchmark_column(col)
        if parsed is None:
            if col in always_on or col in selected_extras:
                display.append(col)
            continue
        dataset, category = parsed
        if category not in selected_cats:
            continue
        if selected_ds and dataset not in selected_ds:
            continue
        display.append(col)
    return display