"""
Folder-level eval coverage for lm-evaluation-harness tasks.

SKIP_LIST: folders intentionally excluded from evaluation, with reasons.
Run as __main__ to regenerate coverage_summary.json from the HF dataset.
"""

from __future__ import annotations

# ── All task folders from lm_eval/tasks/ ─────────────────────────────────────
# (lm-evaluation-harness @ current HEAD, non-folder entries excluded)
ALL_FOLDERS: list[str] = [
    "aclue", "acpbench", "aexams", "afrimgsm", "afrimmlu", "afrixnli",
    "afrobench", "agieval", "aime", "alghafa", "anli", "arab_culture",
    "arab_culture_completion", "arabic_leaderboard_complete",
    "arabic_leaderboard_light", "arabicmmlu", "aradice", "arc", "arc_mt",
    "arithmetic", "asdiv", "babi", "babilong", "bangla", "basque_bench",
    "basqueglue", "bbh", "bbq", "bear", "belebele", "benchmarks", "bertaqa",
    "bhs", "bigbench", "blimp", "blimp_nl", "c4", "cabbq", "careqa",
    "catalan_bench", "ceval", "chartqa", "click", "cmmlu", "cnn_dailymail",
    "code_x_glue", "commonsense_qa", "common_voice", "copal_id", "coqa",
    "crows_pairs", "csatqa", "darija_bench", "darijahellaswag", "darijammlu",
    "discrim_eval", "drop", "e2lmc", "egyhellaswag", "egymmlu", "eq_bench",
    "esbbq", "eus_exams", "eus_proficiency", "eus_reading", "eus_trivia",
    "evalita_llm", "fda", "fld", "french_bench", "galician_bench",
    "glianorex", "global_mmlu", "global_piqa", "glue", "gpqa", "graphwalks",
    "groundcocoa", "gsm8k", "gsm8k_platinum", "gsm_plus", "haerae",
    "headqa", "hellaswag", "hendrycks_ethics", "hendrycks_math",
    "histoires_morales", "hrm8k", "humaneval", "humaneval_infilling",
    "icelandic_winogrande", "ifeval", "inverse_scaling",
    "japanese_leaderboard", "jfinqa", "jsonschema_bench", "kbl", "kmmlu",
    "kobest", "kormedmcqa", "lambada", "lambada_cloze",
    "lambada_multilingual", "lambada_multilingual_stablelm", "leaderboard",
    "libra", "lingoly", "llama3", "lm_syneval", "logiqa", "logiqa2",
    "longbench", "longbench2", "mastermind", "mathqa", "mbpp", "mc_taco",
    "med_concepts_qa", "meddialog", "mediqa_qa2019", "medmcqa",
    "med_prescriptions", "medqa", "medtext", "med_text_classification",
    "mela", "meqsum", "metabench", "mgsm", "mimic_repsum", "minerva_math",
    "mlqa", "mmlu", "mmlu_pro", "mmlu-pro-plus", "mmlu_prox", "mmlu-redux",
    "mmlu-redux-spanish", "mmlusr", "mmmu", "model_written_evals",
    "moral_stories", "mts_dialog", "multiblimp", "mutual", "noreval",
    "noticia", "nq_open", "okapi", "olaph", "openai-mmmlu", "openbookqa",
    "paloma", "paws-x", "pile", "pile_10k", "piqa", "pisa", "polemo2",
    "portuguese_bench", "prost", "pubmedqa", "qa4mre", "qasper", "race",
    "realtoxicityprompts", "ruler", "sciq", "score", "scrolls",
    "simple_cooccurrence_bias", "siqa", "slr_bench", "spanish_bench",
    "squad_completion", "squadv2", "storycloze", "super_glue", "swag",
    "swde", "tinyBenchmarks", "tmlu", "tmmluplus", "toxigen", "translation",
    "triviaqa", "truthfulqa", "truthfulqa-multi", "turblimp", "turkishmmlu",
    "ulqa", "unitxt", "unscramble", "webqs", "wikitext", "winogender",
    "winogrande", "wmdp", "wmt2016", "wsc273", "xcopa", "xnli", "xnli_eu",
    "xquad", "xstorycloze", "xwinograd", "zhoblimp",
]

# ── Skip list ─────────────────────────────────────────────────────────────────
# Value = human-readable reason (fill in as needed).
SKIP_LIST: dict[str, str] = {
    # ── Non-English / regional ────────────────────────────────────────────────
    "aclue": "",
    "afrimgsm": "",
    "afrimmlu": "",
    "afrixnli": "",
    "afrobench": "",
    "alghafa": "",
    "arab_culture": "",
    "arab_culture_completion": "",
    "arabic_leaderboard_complete": "",
    "arabic_leaderboard_light": "",
    "arabicmmlu": "",
    "aradice": "",
    "arc_mt": "",
    "bangla": "",
    "basque_bench": "",
    "basqueglue": "",
    "belebele": "",
    "bertaqa": "",
    "bhs": "",
    "cabbq": "",
    "catalan_bench": "",
    "ceval": "",
    "cmmlu": "",
    "copal_id": "",
    "darija_bench": "",
    "darijahellaswag": "",
    "darijammlu": "",
    "egyhellaswag": "",
    "egymmlu": "",
    "esbbq": "",
    "eus_exams": "",
    "eus_proficiency": "",
    "eus_reading": "",
    "eus_trivia": "",
    "evalita_llm": "",
    "french_bench": "",
    "galician_bench": "",
    "haerae": "",
    "histoires_morales": "",
    "hrm8k": "",
    "icelandic_winogrande": "",
    "japanese_leaderboard": "",
    "jfinqa": "",
    "kbl": "",
    "kmmlu": "",
    "kobest": "",
    "kormedmcqa": "",
    "lambada_multilingual": "",
    "lambada_multilingual_stablelm": "",
    "libra": "",
    "lingoly": "",
    "mela": "",
    "mgsm": "",
    "mlqa": "",
    "noreval": "",
    "noticia": "",
    "okapi": "",
    "openai-mmmlu": "",
    "paws-x": "",
    "polemo2": "",
    "portuguese_bench": "",
    "spanish_bench": "",
    "tmlu": "",
    "tmmluplus": "",
    "turkishmmlu": "",
    "turblimp": "",
    "xcopa": "",
    "xnli_eu": "",
    "xwinograd": "",
    "zhoblimp": "",
    # ── Code execution required ───────────────────────────────────────────────
    "code_x_glue": "",
    "humaneval": "",
    "humaneval_infilling": "",
    "mbpp": "",
    # ── Aggregate / meta benchmark collections ────────────────────────────────
    "benchmarks": "",
    "leaderboard": "",
    "llama3": "",
    "metabench": "",
    "tinyBenchmarks": "",
    # ── Large corpus / perplexity-only ────────────────────────────────────────
    "c4": "",
    "paloma": "",
    "pile": "",
    "pile_10k": "",
    "wikitext": "",
}

# ── Core logic ────────────────────────────────────────────────────────────────

def _folder_for_task(task: str, folder_set: set[str]) -> str | None:
    """Return the longest folder name that is a prefix of task (with separator)."""
    best: str | None = None
    for folder in folder_set:
        if task == folder or task.startswith(folder + "_") or task.startswith(folder + "-"):
            if best is None or len(folder) > len(best):
                best = folder
    return best


def compute_coverage(evaluated_tasks: set[str]) -> dict[str, dict]:
    """
    Returns a dict mapping each folder name to its coverage info:
      status:  "evaluated" | "partial" | "skipped" | "not_started"
      skip_reason: str | None
      matched_tasks: list[str]   (evaluated task names mapped to this folder)
    """
    folder_set = set(ALL_FOLDERS)

    # Map every evaluated task to its folder
    folder_hits: dict[str, list[str]] = {f: [] for f in ALL_FOLDERS}
    for task in evaluated_tasks:
        folder = _folder_for_task(task, folder_set)
        if folder:
            folder_hits[folder].append(task)

    result: dict[str, dict] = {}
    for folder in ALL_FOLDERS:
        hits = folder_hits[folder]
        if folder in SKIP_LIST:
            status = "skipped"
        elif not hits:
            status = "not_started"
        elif folder in evaluated_tasks:
            # The group task itself was evaluated → full run
            status = "evaluated"
        else:
            status = "partial"

        result[folder] = {
            "status": status,
            "skip_reason": SKIP_LIST.get(folder),
            "matched_tasks": sorted(hits),
        }

    return result


def summary_stats(coverage: dict[str, dict]) -> dict[str, int]:
    counts: dict[str, int] = {"evaluated": 0, "partial": 0, "skipped": 0, "not_started": 0}
    for info in coverage.values():
        counts[info["status"]] += 1
    return counts


# ── CLI: regenerate coverage_summary.json ────────────────────────────────────

if __name__ == "__main__":
    import json, os
    from datetime import datetime, timezone
    from pathlib import Path
    from pathlib import PurePosixPath
    from huggingface_hub import HfApi, hf_hub_download

    REPO = "MIMIR-AI-ROUTER/harness_evals"
    TOKEN = os.getenv("HF_TOKEN")

    print("Fetching file list...")
    api = HfApi(token=TOKEN)
    all_files = sorted(api.list_repo_files(REPO, repo_type="dataset"))
    result_files = [
        f for f in all_files
        if PurePosixPath(f).name.startswith("results_") and f.endswith(".json")
    ]
    print(f"  {len(result_files)} result files found")

    evaluated: set[str] = set()
    for filename in result_files:
        local = hf_hub_download(REPO, filename, repo_type="dataset", token=TOKEN)
        data = json.load(open(local))
        evaluated.update(data.get("results", {}).keys())
    print(f"  {len(evaluated)} unique evaluated tasks")

    coverage = compute_coverage(evaluated)
    stats = summary_stats(coverage)

    out = {
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "stats": stats,
        "folders": coverage,
    }

    dest = Path(__file__).parent / "coverage_summary.json"
    dest.write_text(json.dumps(out, indent=2))
    print(f"\nWrote {dest}")
    print(f"  evaluated:   {stats['evaluated']}")
    print(f"  partial:     {stats['partial']}")
    print(f"  skipped:     {stats['skipped']}")
    print(f"  not_started: {stats['not_started']}")