"""Benchmark ontology — families, variants, metric bounds, trust levels and the comparability of two result configurations.

Two scores are only *directly comparable* when they measure the same benchmark variant with the same metric under configurations that
do not change the task: same dataset revision/variant, same evaluator or harness class, same scaffold when the benchmark is agentic,
same shot/pass regime. Reasoning effort, sampling temperature and judge differences make them *partially comparable*. Different
variants (SWE-bench Verified vs Lite), different metrics, or different pass regimes are *not directly comparable*.
"""
from __future__ import annotations

import hashlib
import json
import re
from typing import Any

COMPARABLE, PARTIAL, NOT_COMPARABLE = "comparable", "partially-comparable", "not-comparable"

# Benchmark family → known variants (registry keys). Used to build `variant_of` relations and the family attribute; the registry YAML
# carries the authoritative per-entry `family`/`variant` fields, this table is the fallback for benchmarks created from other sources.
FAMILIES: dict[str, dict[str, Any]] = {
    "swe-bench": {"label": "SWE-bench", "variants": {"swe-bench-full": "full", "swe-bench-verified": "Verified", "swe-bench-lite": "Lite",
                                                       "swe-bench-multimodal": "Multimodal", "swe-bench-multilingual": "Multilingual", "swe-bench-pro": "Pro"}},
    "aime": {"label": "AIME", "variants": {"aime-2024": "2024", "aime-2025": "2025", "aime-2026": "2026"}},
    "mmlu": {"label": "MMLU", "variants": {"mmlu": "original", "mmlu-pro": "Pro", "mmlu-redux": "Redux", "mmmlu": "multilingual"}},
    "mmmu": {"label": "MMMU", "variants": {"mmmu": "original", "mmmu-pro": "Pro"}},
    "gpqa": {"label": "GPQA", "variants": {"gpqa": "main", "gpqa-diamond": "Diamond"}},
    "tau-bench": {"label": "τ-bench", "variants": {"tau-bench": "v1", "tau2-bench": "τ²"}},
    "livebench": {"label": "LiveBench", "variants": {"livebench": "global", "livebench-2": "global"}},
    "arc-agi": {"label": "ARC-AGI", "variants": {"arc-agi": "1", "arc-agi-2": "2", "arc-agi-3": "3"}},
    "humaneval": {"label": "HumanEval", "variants": {"humaneval": "original", "humaneval-plus": "Plus", "mbpp": None}},
    "livecodebench": {"label": "LiveCodeBench", "variants": {"livecodebench": "rolling"}},
    "math": {"label": "MATH", "variants": {"math": "full", "math-500": "500"}},
    "terminal-bench": {"label": "Terminal-Bench", "variants": {"terminal-bench": "1.0", "terminal-bench-2": "2.0"}},
    "ifeval": {"label": "IFEval", "variants": {"ifeval": "original", "ifbench": "IFBench"}},
    "artificial-analysis-intelligence-index": {"label": "Artificial Analysis Intelligence Index", "variants": {"artificial-analysis-intelligence-index": "index"}},
    "humanitys-last-exam": {"label": "Humanity's Last Exam", "variants": {"humanitys-last-exam": "full"}},
    "scicode": {"label": "SciCode", "variants": {"scicode": "main"}},
    "aider-polyglot": {"label": "Aider polyglot", "variants": {"aider-polyglot": "polyglot"}},
    "lmarena": {"label": "LMArena", "variants": {"lmarena-text": "text", "lmarena-vision": "vision", "lmarena-webdev": "webdev"}},
    "mteb": {"label": "MTEB", "variants": {"mteb": "v1", "mteb-v2": "v2", "mmteb": "MMTEB"}},
}
_KEY_TO_FAMILY: dict[str, tuple[str, str | None]] = {k: (fam, v) for fam, spec in FAMILIES.items() for k, v in spec["variants"].items()}


def family_of(benchmark_key: str) -> tuple[str | None, str | None]:
    """(family key, variant label) for a registry key / slug."""
    if benchmark_key in _KEY_TO_FAMILY:
        return _KEY_TO_FAMILY[benchmark_key]
    for fam in FAMILIES:
        if benchmark_key.startswith(fam):
            return fam, benchmark_key[len(fam):].strip("-") or None
    return None, None


# ---------------------------------------------------------------------------------------------- metrics
METRICS: dict[str, dict[str, Any]] = {
    # canonical metric → bounds and direction. `None` bound = unbounded.
    "accuracy": {"min": 0, "max": 100, "higher_is_better": True, "unit": "%"},
    "pass@1": {"min": 0, "max": 100, "higher_is_better": True, "unit": "%"},
    "pass^1": {"min": 0, "max": 100, "higher_is_better": True, "unit": "%"},
    "pass^k": {"min": 0, "max": 100, "higher_is_better": True, "unit": "%"},
    "resolved": {"min": 0, "max": 100, "higher_is_better": True, "unit": "%"},
    "pass_rate_2": {"min": 0, "max": 100, "higher_is_better": True, "unit": "%"},
    "percent_cases_well_formed": {"min": 0, "max": 100, "higher_is_better": True, "unit": "%"},
    "global_average": {"min": 0, "max": 100, "higher_is_better": True, "unit": "%"},
    "average score": {"min": 0, "max": 100, "higher_is_better": True, "unit": "%"},
    "mean score": {"min": 0, "max": 100, "higher_is_better": True, "unit": "%"},
    "index": {"min": 0, "max": 100, "higher_is_better": True, "unit": ""},
    "elo": {"min": 0, "max": None, "higher_is_better": True, "unit": ""},
    "score": {"min": None, "max": None, "higher_is_better": True, "unit": ""},
    "win_rate": {"min": 0, "max": 100, "higher_is_better": True, "unit": "%"},
    "f1": {"min": 0, "max": 100, "higher_is_better": True, "unit": "%"},
    "ndcg": {"min": 0, "max": 100, "higher_is_better": True, "unit": "%"},
    "perplexity": {"min": 0, "max": None, "higher_is_better": False, "unit": ""},
    "latency_ms": {"min": 0, "max": None, "higher_is_better": False, "unit": "ms"},
    "cost_usd": {"min": 0, "max": None, "higher_is_better": False, "unit": "USD"},
}
_METRIC_ALIASES = {
    "acc": "accuracy", "exact match": "accuracy", "em": "accuracy", "pass@1": "pass@1", "pass at 1": "pass@1", "pass^1": "pass^1", "pass1": "pass^1",
    "pass rate (2 attempts)": "pass_rate_2", "pass rate": "pass_rate_2", "percent resolved": "resolved", "% resolved": "resolved", "resolved rate": "resolved",
    "elo / bradley–terry score": "elo", "elo / bradley-terry score": "elo", "arena score": "elo", "bradley-terry": "elo", "mean score": "mean score",
    "average score": "average score", "global average": "global_average", "prompt-level strict accuracy": "accuracy", "strict accuracy": "accuracy",
    "intelligence index": "index", "composite": "index", "ppl": "perplexity",
}


def normalize_metric(raw: str | None) -> str | None:
    if not raw:
        return None
    s = raw.strip().lower()
    if s.startswith("category:"):
        return raw.strip()  # LiveBench per-category averages keep their label; they are a separate metric each
    if s in METRICS:
        return s
    return _METRIC_ALIASES.get(s) or (s if re.fullmatch(r"[a-z0-9_@^%.\- ]+", s) else None)


def metric_bounds(metric: str | None, unit: str | None = None) -> tuple[float | None, float | None]:
    m = normalize_metric(metric)
    if m and m.startswith("category:"):
        return 0, 100
    spec = METRICS.get(m or "")
    if spec:
        return spec["min"], spec["max"]
    if unit == "%":
        return 0, 100
    return None, None


# ---------------------------------------------------------------------------------------------- trust levels
TRUST_LEVELS = ("official-model-card", "official-benchmark", "peer-reviewed", "independent-evaluator", "community", "unverified")
TRUST_LABELS = {
    "official-model-card": "Official model card / technical report (self-reported)",
    "official-benchmark": "Official benchmark leaderboard (submissions checked by the benchmark owner)",
    "peer-reviewed": "Peer-reviewed paper",
    "independent-evaluator": "Independent third-party evaluator",
    "community": "Community-run leaderboard or submission",
    "unverified": "Unverified / unknown provenance",
}
# source key (domain) → trust level of results it publishes
SOURCE_TRUST: dict[str, str] = {
    "swebench.com": "official-benchmark", "aider.chat": "official-benchmark", "livebench.ai": "official-benchmark", "artificialanalysis.ai": "independent-evaluator",
    "lmarena.ai": "independent-evaluator", "scale.com": "independent-evaluator", "epoch.ai": "independent-evaluator", "vals.ai": "independent-evaluator",
    "huggingface.co": "community", "github.com": "community", "arxiv.org": "peer-reviewed", "openreview.net": "peer-reviewed",
}
OFFICIAL_LAB_SOURCES = {"docs.claude.com", "platform.openai.com", "ai.google.dev", "docs.mistral.ai", "api-docs.deepseek.com", "docs.cohere.com",
                        "docs.x.ai", "ai.meta.com", "llama.com", "qwenlm.github.io", "developer.nvidia.com", "machinelearning.apple.com"}


def trust_level(source_key: str | None, config: dict[str, Any] | None = None, *, extractor: str = "deterministic") -> str:
    cfg = config or {}
    if source_key in SOURCE_TRUST:
        level = SOURCE_TRUST[source_key]
        if level == "official-benchmark" and cfg.get("checked_by_swebench") is False:
            return "community"
        if level == "official-benchmark" and str(cfg.get("submission", "")).lower() in ("self-reported", "self reported", "unverified"):
            return "community"
        return level
    if source_key in OFFICIAL_LAB_SOURCES or cfg.get("self_reported") or cfg.get("source_kind") == "model_card":
        return "official-model-card"
    if extractor == "llm":
        return "unverified"
    return "unverified"


# ---------------------------------------------------------------------------------------------- comparability
# Config keys that change the *task* (must match for full comparability).
TASK_KEYS = ("variant", "board", "harness", "evaluator", "subset", "split", "shots", "pass_count", "attempts", "language", "scaffold", "agent", "system")
# Config keys that identify the *run* (a LiveBench release, an AA index version, a dataset revision): a newer run supersedes the older one
# for the same task (`run_group`, one current row per task), and two runs of the same task are only partially comparable.
RUN_KEYS = ("release", "index_version", "version", "dataset_revision")
# Config keys that change the *conditions* (mismatch → partially comparable).
CONDITION_KEYS = ("reasoning_effort", "reasoning", "thinking_budget", "temperature", "judge", "tools", "tool_use", "max_tokens", "context_length",
                  "sampling", "aggregation", "edit_format", "model_tag", *RUN_KEYS)
# Keys that are pure bookkeeping (never affect comparability).
IGNORED_KEYS = {"aa_slug", "livebench_model_id", "api_model_id", "date", "submission", "checked_by_swebench", "open_source_system", "system_org",
                "total_cost_usd", "cost_per_instance_usd", "dirname", "command", "versions", "test_cases", "seconds_per_case", "estimated",
                "livebench_hf_link", "source_kind", "self_reported", "subtasks", "notes", "url"}


def _clean(v: Any) -> Any:
    if isinstance(v, str):
        return v.strip().lower()
    return v


def config_key(config: dict[str, Any] | None, metric: str | None = None) -> str:
    """Stable hash of the comparability-relevant part of a result configuration (task keys + metric). Run keys (release, index
    version) are deliberately excluded: they define `run_group`, so the newest run of a task replaces the older ones on leaderboards."""
    cfg = config or {}
    core = {k: _clean(cfg[k]) for k in TASK_KEYS if k in cfg and cfg[k] not in (None, "", [], {})}
    core["metric"] = normalize_metric(metric) or ""
    return hashlib.sha1(json.dumps(core, sort_keys=True, default=str).encode()).hexdigest()[:12]


def comparability(a_cfg: dict[str, Any] | None, b_cfg: dict[str, Any] | None, a_metric: str | None = None, b_metric: str | None = None,
                  *, same_benchmark: bool = True) -> tuple[str, list[str]]:
    """Return (level, reasons). `same_benchmark` False → not comparable outright."""
    reasons: list[str] = []
    if not same_benchmark:
        return NOT_COMPARABLE, ["different benchmark variants"]
    ma, mb = normalize_metric(a_metric), normalize_metric(b_metric)
    if ma != mb and (ma or mb):
        return NOT_COMPARABLE, [f"different metrics ({a_metric} vs {b_metric})"]
    a, b = a_cfg or {}, b_cfg or {}
    for k in TASK_KEYS:
        if k in a or k in b:
            va, vb = _clean(a.get(k)), _clean(b.get(k))
            if va != vb and va not in (None, "") and vb not in (None, ""):
                reasons.append(f"{k}: {a.get(k)} vs {b.get(k)}")
    if reasons:
        return NOT_COMPARABLE, reasons
    for k in CONDITION_KEYS:
        if k in a or k in b:
            va, vb = _clean(a.get(k)), _clean(b.get(k))
            if va != vb:
                reasons.append(f"{k}: {a.get(k) if k in a else 'unspecified'} vs {b.get(k) if k in b else 'unspecified'}")
    if reasons:
        return PARTIAL, reasons
    return COMPARABLE, ["same variant, metric and evaluation conditions"]


def variant_from_config(config: dict[str, Any] | None) -> str | None:
    cfg = config or {}
    for k in ("variant", "board", "subset", "split"):
        v = cfg.get(k)
        if isinstance(v, str) and v.strip():
            return v.strip()
    return None


def run_group_from_config(config: dict[str, Any] | None) -> str | None:
    """The 'run' a result belongs to (a LiveBench release, an aider run date, an AA index version…). One current row per run group."""
    cfg = config or {}
    for k in ("release", "index_version", "version", "date", "dataset_revision"):
        v = cfg.get(k)
        if v not in (None, ""):
            return str(v)
    return None


__all__ = ["COMPARABLE", "CONDITION_KEYS", "FAMILIES", "IGNORED_KEYS", "METRICS", "NOT_COMPARABLE", "PARTIAL", "RUN_KEYS", "TASK_KEYS", "TRUST_LABELS",
           "TRUST_LEVELS", "comparability", "config_key", "family_of", "metric_bounds", "normalize_metric", "run_group_from_config", "trust_level",
           "variant_from_config"]