"""Load and flatten the SaylorTwift/llm-benchmark-usage dataset for exploration."""

from datetime import datetime

import pandas as pd
from datasets import load_dataset

DATASET_ID = "SaylorTwift/llm-benchmark-usage"


def _bucket(category: str) -> str:
    """Collapse the free-text per-source categories into a fixed set of broad buckets
    (same scheme used to build the original study's category-evolution chart)."""
    c = category.lower()
    if any(k in c for k in ["safety", "preparedness", "sandbagging", "bias", "hallucination", "health", "alignment"]):
        return "safety"
    if any(k in c for k in ["vision", "video", "audio", "multimodal"]):
        return "vision"
    if any(k in c for k in ["multilingual", "korean", "russian"]):
        return "multilingual"
    if "long_context" in c:
        return "long_context"
    if "math" in c:
        return "math"
    if any(k in c for k in ["agentic", "agent", "tool", "computer_use", "search"]) or c in {
        "north_mini_code", "grm2.6plus", "terminus_qwen3_8b", "m2.5_search", "m2.5_office",
        "m2.7_office", "coder_implied",
    }:
        return "agentic"
    if "cod" in c:
        return "coding"
    if any(
        k in c
        for k in [
            "knowledge", "reasoning", "general", "commonsense", "factuality", "scien",
            "domain_specific", "human_eval", "macaron_v1", "intern_s2", "aggregate",
            "reading_comprehension", "qa_reading", "simple_evals", "internal_framework",
            "internal_proprietary", "open_llm_leaderboard", "out_of_distribution",
            "llm_judge", "pretraining_ablation", "abstraction_reasoning", "human_preference",
            "emotional_intelligence", "creative_writing",
        ]
    ) or c in {"base", "pretrained", "shared_base", "engineering"}:
        return "knowledge"
    return "other"


CATEGORY_ORDER = ["knowledge", "coding", "agentic", "math", "safety", "vision", "multilingual", "long_context", "other"]


def _half(d: datetime) -> str:
    return f"{d.year}-{'H1' if d.month <= 6 else 'H2'}"


def load_data():
    """Returns (models_df, usage_df).

    models_df: one row per model (model_id, lab, release_date, source, is_open, period)
    usage_df: one row per (model_id, benchmark, category, bucket) — the fully resolved
              benchmark usage table used by every tab.
    """
    print("load_data: fetching models config...", flush=True)
    models_ds = load_dataset(DATASET_ID, "models", token=False)["models"]
    print("load_data: done", flush=True)

    models_raw = list(models_ds)

    models_df = pd.DataFrame(
        [{k: v for k, v in m.items() if k != "benchmarks"} for m in models_raw]
    )
    models_df["release_date"] = pd.to_datetime(models_df["release_date"])
    models_df["is_open"] = models_df["model_id"].str.contains("/")
    models_df["period"] = models_df["release_date"].apply(_half)
    models_df = models_df.sort_values("release_date").reset_index(drop=True)

    model_meta = models_df.set_index("model_id").to_dict("index")

    rows = []
    for m in models_raw:
        meta = model_meta[m["model_id"]]
        for b in m["benchmarks"]:
            rows.append(
                {
                    "model_id": m["model_id"],
                    "lab": meta["lab"],
                    "release_date": meta["release_date"],
                    "period": meta["period"],
                    "is_open": meta["is_open"],
                    "source": m["source"],
                    "benchmark": b["name"],
                    "category": b["category"],
                    "bucket": _bucket(b["category"]),
                }
            )
    usage_df = pd.DataFrame(rows)
    return models_df, usage_df


if __name__ == "__main__":
    models_df, usage_df = load_data()
    print("models:", len(models_df), "| usage rows:", len(usage_df))
    print(models_df.head())
    print(usage_df.head())