data.py

"""Load published Praxa behavioral experiment metrics from Hugging Face Hub.

Source of truth: EdyVision/praxa-behavioral-experiments (private dataset).
Does not depend on the training harness package.

Load strategy: walk ``domains/.../results/<model_key>/<condition>/metrics.json``
on the Hub so newly published conditions (e.g. RFDT) appear after Refresh
without editing this file. Known labels stay curated; unknown dirs get a
readable fallback label.

Caching: in-process frames plus a disk snapshot under ``PRAXA_CACHE_DIR``
(default ``/tmp/praxa_hub_cache``) shared across gunicorn workers. Ordinary
page loads reuse memory/disk; the Refresh button force-clears both and
re-lists the Hub.
"""

from __future__ import annotations

import fcntl
import json
import os
import pickle
import threading
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import pandas as pd

DEFAULT_EXPERIMENTS_REPO = os.environ.get(
    "PRAXA_EXPERIMENTS_REPO",
    "EdyVision/praxa-behavioral-experiments",
)

DOMAIN_PREFIX: dict[str, str] = {
    "medicareplus": "domains/healthcare/medicareplus_v1",
    "citygov": "domains/government/citygov_v1",
}

DOMAIN_LABELS: dict[str, str] = {
    "medicareplus": "MediCarePlus (healthcare)",
    "citygov": "CityGov (government)",
}

MODEL_KEYS: tuple[str, ...] = (
    "qwen3-8b",
    "granite41-8b",
    "ministral3-8b-instruct-bf16",
    "modernbert-base",
)

MATCHED_8B_KEYS: tuple[str, ...] = (
    "qwen3-8b",
    "granite41-8b",
    "ministral3-8b-instruct-bf16",
)

ENCODER_MODEL_KEYS: tuple[str, ...] = ("modernbert-base",)

MODEL_LABELS: dict[str, str] = {
    "qwen3-8b": "Qwen3-8B",
    "granite41-8b": "Granite4.1-8B",
    "ministral3-8b-instruct-bf16": "Ministral3-8B",
    "modernbert-base": "ModernBERT-base",
}

FROZEN_CONDITIONS: tuple[str, ...] = (
    "base",
    "explicit_full_json",
    "explicit_retrieved",
    "1shot",
    "2shot",
    "4shot",
)

LEARNED_CONDITIONS: tuple[str, ...] = ("lora", "cartridge")

RFDT_CONDITIONS: tuple[str, ...] = ("rfdt_base", "rfdt_lora")

ENCODER_CONDITIONS: tuple[str, ...] = ("encoder_ft",)

ALL_CONDITIONS: tuple[str, ...] = (
    FROZEN_CONDITIONS + LEARNED_CONDITIONS + RFDT_CONDITIONS + ENCODER_CONDITIONS
)

CONDITION_LABELS: dict[str, str] = {
    "base": "Base",
    "explicit_full_json": "Explicit policy",
    "explicit_retrieved": "Explicit retrieved",
    "1shot": "1-shot",
    "2shot": "2-shot",
    "4shot": "4-shot",
    "lora": "LoRA",
    "cartridge": "Behavioral KV",
    "rfdt_base": "RFDT Base",
    "rfdt_lora": "RFDT-LoRA",
    "encoder_ft": "Encoder-FT",
}

METRIC_LABELS: dict[str, str] = {
    "macro_f1": "Macro-F1",
    "accuracy": "Accuracy",
    "answer_f1": "Answer F1",
    "refuse_f1": "Refuse F1",
    "allowed_edge": "Allowed-edge Acc",
    "denied_edge": "Denied-edge Acc",
    "denied_edge_refuse_recall": "Denied-edge Refuse Recall",
}

QUERY_TYPE_SLICES: tuple[str, ...] = (
    "allowed_base",
    "allowed_edge",
    "denied_base",
    "denied_edge",
)

QUERY_TYPE_LABELS: dict[str, str] = {
    "allowed_base": "Allowed base",
    "allowed_edge": "Allowed edge",
    "denied_base": "Denied base",
    "denied_edge": "Denied edge",
}

SLICE_HEATMAP_CONDITIONS: tuple[str, ...] = (
    "base",
    "4shot",
    "lora",
    "cartridge",
    "rfdt_lora",
    "encoder_ft",
)

# Seed catalog (also used when Hub tree listing fails). Discovery can find more.
_DOMAIN_RUNS: dict[str, list[tuple[str, str, str]]] = {
    "medicareplus": [
        ("medicareplus_qwen3_8b_base", "qwen3-8b", "base"),
        ("medicareplus_qwen3_8b_explicit_policy", "qwen3-8b", "explicit_full_json"),
        ("medicareplus_qwen3_8b_explicit_retrieved", "qwen3-8b", "explicit_retrieved"),
        ("medicareplus_qwen3_8b_few_shot_1", "qwen3-8b", "1shot"),
        ("medicareplus_qwen3_8b_few_shot_2", "qwen3-8b", "2shot"),
        ("medicareplus_qwen3_8b_few_shot_4", "qwen3-8b", "4shot"),
        ("medicareplus_qwen3_8b_lora", "qwen3-8b", "lora"),
        ("medicareplus_qwen3_8b_cartridge", "qwen3-8b", "cartridge"),
        ("medicareplus_qwen3_8b_rfdt_base", "qwen3-8b", "rfdt_base"),
        ("medicareplus_qwen3_8b_rfdt_lora", "qwen3-8b", "rfdt_lora"),
        ("medicareplus_granite41_8b_base", "granite41-8b", "base"),
        (
            "medicareplus_granite41_8b_explicit_policy",
            "granite41-8b",
            "explicit_full_json",
        ),
        (
            "medicareplus_granite41_8b_explicit_retrieved",
            "granite41-8b",
            "explicit_retrieved",
        ),
        ("medicareplus_granite41_8b_few_shot_1", "granite41-8b", "1shot"),
        ("medicareplus_granite41_8b_few_shot_2", "granite41-8b", "2shot"),
        ("medicareplus_granite41_8b_few_shot_4", "granite41-8b", "4shot"),
        ("medicareplus_granite41_8b_lora", "granite41-8b", "lora"),
        ("medicareplus_granite41_8b_cartridge", "granite41-8b", "cartridge"),
        ("medicareplus_granite41_8b_rfdt_base", "granite41-8b", "rfdt_base"),
        ("medicareplus_granite41_8b_rfdt_lora", "granite41-8b", "rfdt_lora"),
        ("medicareplus_ministral3_8b_base", "ministral3-8b-instruct-bf16", "base"),
        (
            "medicareplus_ministral3_8b_explicit_policy",
            "ministral3-8b-instruct-bf16",
            "explicit_full_json",
        ),
        (
            "medicareplus_ministral3_8b_explicit_retrieved",
            "ministral3-8b-instruct-bf16",
            "explicit_retrieved",
        ),
        ("medicareplus_ministral3_8b_few_shot_1", "ministral3-8b-instruct-bf16", "1shot"),
        ("medicareplus_ministral3_8b_few_shot_2", "ministral3-8b-instruct-bf16", "2shot"),
        ("medicareplus_ministral3_8b_few_shot_4", "ministral3-8b-instruct-bf16", "4shot"),
        ("medicareplus_ministral3_8b_lora", "ministral3-8b-instruct-bf16", "lora"),
        (
            "medicareplus_ministral3_8b_cartridge",
            "ministral3-8b-instruct-bf16",
            "cartridge",
        ),
        (
            "medicareplus_ministral3_8b_rfdt_base",
            "ministral3-8b-instruct-bf16",
            "rfdt_base",
        ),
        (
            "medicareplus_ministral3_8b_rfdt_lora",
            "ministral3-8b-instruct-bf16",
            "rfdt_lora",
        ),
        (
            "medicareplus_modernbert_base_encoder_ft",
            "modernbert-base",
            "encoder_ft",
        ),
    ],
    "citygov": [
        ("citygov_qwen3_8b_base", "qwen3-8b", "base"),
        ("citygov_qwen3_8b_explicit_policy", "qwen3-8b", "explicit_full_json"),
        ("citygov_qwen3_8b_explicit_retrieved", "qwen3-8b", "explicit_retrieved"),
        ("citygov_qwen3_8b_few_shot_1", "qwen3-8b", "1shot"),
        ("citygov_qwen3_8b_few_shot_2", "qwen3-8b", "2shot"),
        ("citygov_qwen3_8b_few_shot_4", "qwen3-8b", "4shot"),
        ("citygov_qwen3_8b_lora", "qwen3-8b", "lora"),
        ("citygov_qwen3_8b_cartridge", "qwen3-8b", "cartridge"),
        ("citygov_qwen3_8b_rfdt_base", "qwen3-8b", "rfdt_base"),
        ("citygov_qwen3_8b_rfdt_lora", "qwen3-8b", "rfdt_lora"),
        ("citygov_granite41_8b_base", "granite41-8b", "base"),
        ("citygov_granite41_8b_explicit_policy", "granite41-8b", "explicit_full_json"),
        (
            "citygov_granite41_8b_explicit_retrieved",
            "granite41-8b",
            "explicit_retrieved",
        ),
        ("citygov_granite41_8b_few_shot_1", "granite41-8b", "1shot"),
        ("citygov_granite41_8b_few_shot_2", "granite41-8b", "2shot"),
        ("citygov_granite41_8b_few_shot_4", "granite41-8b", "4shot"),
        ("citygov_granite41_8b_lora", "granite41-8b", "lora"),
        ("citygov_granite41_8b_cartridge", "granite41-8b", "cartridge"),
        ("citygov_granite41_8b_rfdt_base", "granite41-8b", "rfdt_base"),
        ("citygov_granite41_8b_rfdt_lora", "granite41-8b", "rfdt_lora"),
        ("citygov_ministral3_8b_base", "ministral3-8b-instruct-bf16", "base"),
        (
            "citygov_ministral3_8b_explicit_policy",
            "ministral3-8b-instruct-bf16",
            "explicit_full_json",
        ),
        (
            "citygov_ministral3_8b_explicit_retrieved",
            "ministral3-8b-instruct-bf16",
            "explicit_retrieved",
        ),
        ("citygov_ministral3_8b_few_shot_1", "ministral3-8b-instruct-bf16", "1shot"),
        ("citygov_ministral3_8b_few_shot_2", "ministral3-8b-instruct-bf16", "2shot"),
        ("citygov_ministral3_8b_few_shot_4", "ministral3-8b-instruct-bf16", "4shot"),
        ("citygov_ministral3_8b_lora", "ministral3-8b-instruct-bf16", "lora"),
        ("citygov_ministral3_8b_cartridge", "ministral3-8b-instruct-bf16", "cartridge"),
        (
            "citygov_ministral3_8b_rfdt_base",
            "ministral3-8b-instruct-bf16",
            "rfdt_base",
        ),
        (
            "citygov_ministral3_8b_rfdt_lora",
            "ministral3-8b-instruct-bf16",
            "rfdt_lora",
        ),
        (
            "citygov_modernbert_base_encoder_ft",
            "modernbert-base",
            "encoder_ft",
        ),
    ],
}


def hf_token() -> str | None:
    return (
        os.environ.get("HF_TOKEN")
        or os.environ.get("HUGGING_FACE_HUB_TOKEN")
        or None
    )


def hub_results_path(domain: str, model_key: str, condition: str) -> str:
    prefix = DOMAIN_PREFIX[domain]
    return f"{prefix}/results/{model_key}/{condition}"


def condition_label(condition: str) -> str:
    if condition in CONDITION_LABELS:
        return CONDITION_LABELS[condition]
    return condition.replace("_", " ").strip().title()


def condition_family(condition: str) -> str:
    if condition in FROZEN_CONDITIONS:
        return "frozen"
    if condition in RFDT_CONDITIONS:
        return "rfdt"
    if condition in ENCODER_CONDITIONS:
        return "encoder"
    if condition in LEARNED_CONDITIONS:
        return "learned"
    if condition.endswith("_lora") or condition in {"lora", "cartridge"}:
        return "learned"
    return "other"


def condition_order_index(condition: str) -> int:
    order = {name: i for i, name in enumerate(ALL_CONDITIONS)}
    if condition in order:
        return order[condition]
    return len(ALL_CONDITIONS) + abs(hash(condition)) % 1000


def _run_id_for(domain: str, model_key: str, condition: str) -> str:
    for run_id, mk, cond in _DOMAIN_RUNS.get(domain, []):
        if mk == model_key and cond == condition:
            return run_id
    return f"{domain}_{model_key}_{condition}".replace("-", "_")


def discover_domain_runs(
    domain: str,
    *,
    repo_id: str,
    token: str | None,
) -> list[tuple[str, str, str]]:
    """List (run_id, model_key, condition) from Hub results tree."""

    from huggingface_hub import HfApi

    prefix = DOMAIN_PREFIX[domain]
    results_root = f"{prefix}/results"
    api = HfApi(token=token)
    found: list[tuple[str, str, str]] = []
    seen: set[tuple[str, str]] = set()

    try:
        entries = api.list_repo_tree(
            repo_id=repo_id,
            repo_type="dataset",
            path_in_repo=results_root,
            recursive=True,
        )
    except Exception:
        return []

    for entry in entries:
        path = getattr(entry, "path", None) or str(entry)
        if not path.endswith("/metrics.json"):
            continue
        marker = "/results/"
        if marker not in path:
            continue
        rel = path.split(marker, 1)[1]
        parts = rel.split("/")
        if len(parts) != 3 or parts[2] != "metrics.json":
            continue
        model_key, condition = parts[0], parts[1]
        # Only auto-ingest under known model keys. Legacy Hub dirs like
        # ``qwen3-8b-base/lora`` are orphans from older publishes and must not
        # appear as fake models on the ladder.
        if model_key not in MODEL_KEYS:
            continue
        key = (model_key, condition)
        if key in seen:
            continue
        seen.add(key)
        found.append((_run_id_for(domain, model_key, condition), model_key, condition))

    found.sort(key=lambda row: (row[1], condition_order_index(row[2]), row[2]))
    return found


def _catalog_runs(domain: str) -> list[tuple[str, str, str]]:
    return list(_DOMAIN_RUNS.get(domain, []))


def _summarize_metrics(metrics: dict[str, Any]) -> dict[str, Any]:
    pc = metrics.get("per_class") or {}
    bq = metrics.get("by_query_type") or {}
    denied = (bq.get("denied_edge") or {}).get("per_class") or {}
    out: dict[str, Any] = {
        "macro_f1": metrics.get("macro_f1"),
        "accuracy": metrics.get("accuracy"),
        "answer_f1": (pc.get("answer") or {}).get("f1"),
        "refuse_f1": (pc.get("refuse") or {}).get("f1"),
        "allowed_edge": (bq.get("allowed_edge") or {}).get("accuracy"),
        "denied_edge": (bq.get("denied_edge") or {}).get("accuracy"),
        "denied_edge_refuse_recall": (denied.get("refuse") or {}).get("recall"),
        "selection_status": metrics.get("selection_status"),
    }
    for slice_name in QUERY_TYPE_SLICES:
        block = bq.get(slice_name) or {}
        refuse = (block.get("per_class") or {}).get("refuse") or {}
        out[f"slice_acc_{slice_name}"] = block.get("accuracy")
        out[f"slice_refuse_recall_{slice_name}"] = refuse.get("recall")
        out[f"slice_n_{slice_name}"] = block.get("n")
    return out


def _download_metrics_json(
    *,
    repo_id: str,
    path_in_repo: str,
    token: str | None,
) -> dict[str, Any] | None:
    from huggingface_hub import hf_hub_download
    from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError

    try:
        local = hf_hub_download(
            repo_id=repo_id,
            repo_type="dataset",
            filename=path_in_repo,
            token=token,
        )
    except (EntryNotFoundError, RepositoryNotFoundError, OSError, ValueError):
        return None
    except Exception:
        return None

    return json.loads(open(local, encoding="utf-8").read())


def _download_jsonl(
    *,
    repo_id: str,
    path_in_repo: str,
    token: str | None,
) -> list[dict[str, Any]]:
    from huggingface_hub import hf_hub_download
    from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError

    try:
        local = hf_hub_download(
            repo_id=repo_id,
            repo_type="dataset",
            filename=path_in_repo,
            token=token,
        )
    except (EntryNotFoundError, RepositoryNotFoundError, OSError, ValueError):
        return []
    except Exception:
        return []

    rows: list[dict[str, Any]] = []
    with open(local, encoding="utf-8") as handle:
        for line in handle:
            line = line.strip()
            if not line:
                continue
            rows.append(json.loads(line))
    return rows


def load_run_predictions(
    domain: str,
    model_key: str,
    condition: str,
    *,
    repo_id: str | None = None,
) -> pd.DataFrame:
    """Fetch predictions.jsonl for one published run."""

    if domain not in DOMAIN_PREFIX:
        raise ValueError(f"unknown domain {domain!r}")
    repo = repo_id or DEFAULT_EXPERIMENTS_REPO
    hub_path = hub_results_path(domain, model_key, condition)
    rows = _download_jsonl(
        repo_id=repo,
        path_in_repo=f"{hub_path}/predictions.jsonl",
        token=hf_token(),
    )
    if not rows:
        return pd.DataFrame()
    frame = pd.DataFrame(rows)
    frame["domain"] = domain
    frame["model_key"] = model_key
    frame["model"] = MODEL_LABELS.get(model_key, model_key)
    frame["condition"] = condition
    frame["condition_label"] = condition_label(condition)
    return frame


def load_domain_metrics(
    domain: str,
    *,
    repo_id: str | None = None,
    skip_missing: bool = True,
) -> pd.DataFrame:
    """Fetch metrics.json for every published run under a domain."""
    if domain not in DOMAIN_PREFIX:
        raise ValueError(f"unknown domain {domain!r}")

    repo = repo_id or DEFAULT_EXPERIMENTS_REPO
    token = hf_token()
    rows: list[dict[str, Any]] = []

    discovered = discover_domain_runs(domain, repo_id=repo, token=token)
    catalog = _catalog_runs(domain)
    # Prefer discovery when it finds anything; otherwise fall back to catalog
    # so empty Hub / offline still has a deterministic attempt list.
    runs = discovered if discovered else catalog

    # Union: also try catalog slots discovery missed (listing glitches).
    if discovered:
        seen = {(mk, cond) for _rid, mk, cond in discovered}
        for run in catalog:
            if (run[1], run[2]) not in seen:
                runs.append(run)

    for run_id, model_key, condition in runs:
        hub_path = hub_results_path(domain, model_key, condition)
        metrics = _download_metrics_json(
            repo_id=repo,
            path_in_repo=f"{hub_path}/metrics.json",
            token=token,
        )
        if metrics is None:
            if skip_missing:
                continue
            raise FileNotFoundError(f"missing {repo}/{hub_path}/metrics.json")

        summary = _summarize_metrics(metrics)
        family = condition_family(condition)
        rows.append(
            {
                "domain": domain,
                "domain_label": DOMAIN_LABELS[domain],
                "run_id": run_id,
                "model_key": model_key,
                "model": MODEL_LABELS.get(model_key, model_key),
                "condition": condition,
                "condition_label": condition_label(condition),
                "hub_path": hub_path,
                "n": metrics.get("n"),
                "matched_8b": model_key in MATCHED_8B_KEYS,
                "is_encoder": model_key in ENCODER_MODEL_KEYS
                or family == "encoder",
                "family": family,
                **summary,
            }
        )

    return pd.DataFrame(rows)


@dataclass
class MetricsCache:
    frames: dict[str, pd.DataFrame]
    predictions: dict[tuple[str, str, str], pd.DataFrame]
    loaded_at: float | None
    error: str | None
    repo_id: str


_CACHE_LOCK = threading.Lock()
_CACHE = MetricsCache(
    frames={},
    predictions={},
    loaded_at=None,
    error=None,
    repo_id=DEFAULT_EXPERIMENTS_REPO,
)

# Disk snapshot shared across gunicorn workers. Memory is still L1; Refresh
# (force=True) deletes the snapshot and re-pulls from Hub.
_DISK_CACHE_DIR = Path(
    os.environ.get("PRAXA_CACHE_DIR", "/tmp/praxa_hub_cache")
)


def _disk_cache_dir() -> Path:
    _DISK_CACHE_DIR.mkdir(parents=True, exist_ok=True)
    return _DISK_CACHE_DIR


def _metrics_snapshot_path() -> Path:
    return _disk_cache_dir() / "metrics_frames.pkl"


def _analysis_snapshot_path() -> Path:
    return _disk_cache_dir() / "analysis_cache.pkl"


def _refresh_lock_path() -> Path:
    return _disk_cache_dir() / ".refresh.lock"


def _clear_disk_snapshots() -> None:
    for path in (_metrics_snapshot_path(), _analysis_snapshot_path()):
        try:
            path.unlink()
        except FileNotFoundError:
            pass
        except OSError:
            pass


def _read_metrics_snapshot(
    wanted: tuple[str, ...],
) -> tuple[dict[str, pd.DataFrame], float, str | None] | None:
    path = _metrics_snapshot_path()
    if not path.is_file():
        return None
    try:
        blob = pickle.loads(path.read_bytes())
    except Exception:
        return None
    if not isinstance(blob, dict):
        return None
    if blob.get("repo_id") != DEFAULT_EXPERIMENTS_REPO:
        return None
    frames = blob.get("frames") or {}
    if not all(domain in frames for domain in wanted):
        return None
    if not all(isinstance(frames[domain], pd.DataFrame) for domain in wanted):
        return None
    loaded_at = float(blob.get("loaded_at") or time.time())
    error = blob.get("error")
    if error is not None and not isinstance(error, str):
        error = str(error)
    return frames, loaded_at, error


def _write_metrics_snapshot(
    frames: dict[str, pd.DataFrame],
    *,
    loaded_at: float,
    error: str | None,
) -> None:
    payload = {
        "repo_id": DEFAULT_EXPERIMENTS_REPO,
        "loaded_at": loaded_at,
        "error": error,
        "frames": frames,
    }
    path = _metrics_snapshot_path()
    tmp = path.with_suffix(".pkl.tmp")
    tmp.write_bytes(pickle.dumps(payload, protocol=pickle.HIGHEST_PROTOCOL))
    tmp.replace(path)


def _read_analysis_snapshot() -> dict[str, Any] | None:
    path = _analysis_snapshot_path()
    if not path.is_file():
        return None
    try:
        blob = pickle.loads(path.read_bytes())
    except Exception:
        return None
    if not isinstance(blob, dict):
        return None
    if blob.get("repo_id") != DEFAULT_EXPERIMENTS_REPO:
        return None
    if blob.get("loaded_at") is None:
        return None
    return blob


def _write_analysis_snapshot(state: dict[str, Any]) -> None:
    payload = {
        "repo_id": DEFAULT_EXPERIMENTS_REPO,
        "loaded_at": state.get("loaded_at"),
        "error": state.get("error"),
        "catalog": state.get("catalog"),
        "by_id": state.get("by_id") or {},
    }
    path = _analysis_snapshot_path()
    tmp = path.with_suffix(".pkl.tmp")
    tmp.write_bytes(pickle.dumps(payload, protocol=pickle.HIGHEST_PROTOCOL))
    tmp.replace(path)


def cache_status() -> dict[str, Any]:
    with _CACHE_LOCK:
        return {
            "loaded_at": _CACHE.loaded_at,
            "error": _CACHE.error,
            "repo_id": _CACHE.repo_id,
            "domains": {
                d: (0 if df is None else len(df)) for d, df in _CACHE.frames.items()
            },
            "has_token": bool(hf_token()),
            "disk_snapshot": _metrics_snapshot_path().is_file(),
            "cache_dir": str(_DISK_CACHE_DIR),
        }


def refresh_cache(*, force: bool = False, domains: tuple[str, ...] | None = None) -> MetricsCache:
    wanted = domains or tuple(DOMAIN_PREFIX.keys())
    with _CACHE_LOCK:
        if (
            not force
            and _CACHE.loaded_at is not None
            and all(d in _CACHE.frames for d in wanted)
        ):
            warm = True
        else:
            warm = False
    if warm:
        refresh_analysis_cache(force=False)
        return _CACHE

    # Cross-process lock so gunicorn workers share one Hub pull / disk write.
    with open(_refresh_lock_path(), "a+", encoding="utf-8") as lock_file:
        fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX)
        with _CACHE_LOCK:
            if (
                not force
                and _CACHE.loaded_at is not None
                and all(d in _CACHE.frames for d in wanted)
            ):
                pass  # fall through to analysis refresh below
            else:
                if force:
                    _clear_disk_snapshots()
                    _CACHE.predictions = {}
                    with _ANALYSIS_LOCK:
                        _ANALYSIS_CACHE["loaded_at"] = None
                        _ANALYSIS_CACHE["error"] = None
                        _ANALYSIS_CACHE["catalog"] = None
                        _ANALYSIS_CACHE["by_id"] = {}

                loaded_from_disk = False
                if not force:
                    snap = _read_metrics_snapshot(wanted)
                    if snap is not None:
                        frames, loaded_at, error = snap
                        _CACHE.frames = dict(frames)
                        _CACHE.loaded_at = loaded_at
                        _CACHE.error = error
                        _CACHE.repo_id = DEFAULT_EXPERIMENTS_REPO
                        loaded_from_disk = True

                if not loaded_from_disk:
                    frames = dict(_CACHE.frames)
                    error = None
                    try:
                        if not hf_token():
                            error = (
                                "HF_TOKEN is not set. Add it as a Space secret "
                                "to read the private experiments dataset."
                            )
                        for domain in wanted:
                            frames[domain] = load_domain_metrics(domain)
                            if frames[domain].empty and error is None:
                                # Empty can also mean not published yet.
                                pass
                    except Exception as exc:  # noqa: BLE001 — surface to UI
                        error = f"{type(exc).__name__}: {exc}"

                    loaded_at = time.time()
                    _CACHE.frames = frames
                    _CACHE.loaded_at = loaded_at
                    _CACHE.error = error
                    _CACHE.repo_id = DEFAULT_EXPERIMENTS_REPO
                    try:
                        _write_metrics_snapshot(
                            frames, loaded_at=loaded_at, error=error
                        )
                    except OSError:
                        pass

    refresh_analysis_cache(force=force)
    return _CACHE


def get_metrics(
    domain: str,
    *,
    matched_8b_only: bool = False,
    families: tuple[str, ...] | None = None,
) -> pd.DataFrame:
    cache = refresh_cache()
    df = cache.frames.get(domain)
    if df is None or df.empty:
        return pd.DataFrame()

    out = df.copy()
    if matched_8b_only:
        out = out[out["matched_8b"]]
    if families:
        out = out[out["family"].isin(families)]
    return out.reset_index(drop=True)


def get_predictions(
    domain: str,
    model_key: str,
    condition: str,
    *,
    repo_id: str | None = None,
) -> pd.DataFrame:
    """Cached predictions.jsonl for one run (lazy; Hub download on first hit)."""

    key = (domain, model_key, condition)
    with _CACHE_LOCK:
        hit = _CACHE.predictions.get(key)
        if hit is not None:
            return hit.copy()

    frame = load_run_predictions(
        domain, model_key, condition, repo_id=repo_id
    )
    with _CACHE_LOCK:
        _CACHE.predictions[key] = frame
        return frame.copy()


def _pred_label(row: pd.Series) -> str | None:
    for col in ("parsed_prediction", "pred_outcome"):
        val = row.get(col)
        if val is not None and not (isinstance(val, float) and pd.isna(val)):
            return str(val).strip().lower() or None
    return None


def _gold_label(row: pd.Series) -> str | None:
    for col in ("target", "gold_outcome"):
        val = row.get(col)
        if val is not None and not (isinstance(val, float) and pd.isna(val)):
            return str(val).strip().lower() or None
    return None


def slice_accuracy_frame(
    df: pd.DataFrame,
    *,
    conditions: tuple[str, ...] | None = None,
) -> pd.DataFrame:
    """Long frame: method × query_type accuracy from metrics summary columns."""

    wanted = conditions or SLICE_HEATMAP_CONDITIONS
    if df.empty:
        return pd.DataFrame()

    focus = df[df["condition"].isin(wanted)].copy()
    if focus.empty:
        return pd.DataFrame()

    rows: list[dict[str, Any]] = []
    for row in focus.itertuples(index=False):
        method = f"{row.model} · {row.condition_label}"
        for slice_name in QUERY_TYPE_SLICES:
            col = f"slice_acc_{slice_name}"
            value = getattr(row, col, None)
            if value is None or (isinstance(value, float) and pd.isna(value)):
                continue
            n_col = f"slice_n_{slice_name}"
            rows.append(
                {
                    "method": method,
                    "model": row.model,
                    "model_key": row.model_key,
                    "condition": row.condition,
                    "condition_label": row.condition_label,
                    "query_type": slice_name,
                    "query_type_label": QUERY_TYPE_LABELS.get(
                        slice_name, slice_name
                    ),
                    "accuracy": float(value),
                    "n": getattr(row, n_col, None),
                    "_cond_ord": condition_order_index(row.condition),
                }
            )
    if not rows:
        return pd.DataFrame()
    out = pd.DataFrame(rows)
    return out.sort_values(["_cond_ord", "model", "query_type"]).reset_index(
        drop=True
    )


def rfdt_encoder_disagreement(
    domain: str,
    *,
    matched_8b_only: bool = True,
    encoder_model_key: str = "modernbert-base",
) -> pd.DataFrame:
    """Join RFDT-LoRA vs Encoder-FT predictions on sample_id."""

    encoder = get_predictions(domain, encoder_model_key, "encoder_ft")
    if encoder.empty:
        return pd.DataFrame()

    enc = encoder.copy()
    enc["enc_pred"] = enc.apply(_pred_label, axis=1)
    enc["gold"] = enc.apply(_gold_label, axis=1)
    enc["enc_correct"] = enc["enc_pred"] == enc["gold"]
    enc = enc[
        ["sample_id", "query_type", "category", "gold", "enc_pred", "enc_correct"]
    ].drop_duplicates("sample_id")

    model_keys = MATCHED_8B_KEYS if matched_8b_only else MODEL_KEYS
    rows: list[dict[str, Any]] = []
    for model_key in model_keys:
        if model_key in ENCODER_MODEL_KEYS:
            continue
        rfdt = get_predictions(domain, model_key, "rfdt_lora")
        if rfdt.empty:
            continue
        left = rfdt.copy()
        left["rfdt_pred"] = left.apply(_pred_label, axis=1)
        left["rfdt_gold"] = left.apply(_gold_label, axis=1)
        left["rfdt_correct"] = left["rfdt_pred"] == left["rfdt_gold"]
        merged = left.merge(enc, on="sample_id", how="inner", suffixes=("", "_enc"))
        if merged.empty:
            continue
        # Prefer RFDT query_type; fall back to encoder column if missing.
        if "query_type" not in merged.columns or merged["query_type"].isna().all():
            merged["query_type"] = merged.get("query_type_enc")
        for row in merged.itertuples(index=False):
            agree_pred = row.rfdt_pred == row.enc_pred
            if row.rfdt_correct and row.enc_correct:
                outcome = "both_correct"
            elif (not row.rfdt_correct) and (not row.enc_correct):
                outcome = "both_wrong"
            elif row.rfdt_correct and (not row.enc_correct):
                outcome = "rfdt_only"
            else:
                outcome = "encoder_only"
            qt = getattr(row, "query_type", None) or "unknown"
            rows.append(
                {
                    "domain": domain,
                    "sample_id": row.sample_id,
                    "model_key": model_key,
                    "model": MODEL_LABELS.get(model_key, model_key),
                    "query_type": qt,
                    "query_type_label": QUERY_TYPE_LABELS.get(qt, qt),
                    "gold": row.gold if hasattr(row, "gold") else row.rfdt_gold,
                    "rfdt_pred": row.rfdt_pred,
                    "enc_pred": row.enc_pred,
                    "rfdt_correct": bool(row.rfdt_correct),
                    "enc_correct": bool(row.enc_correct),
                    "agree_pred": bool(agree_pred),
                    "outcome": outcome,
                }
            )
    if not rows:
        return pd.DataFrame()
    return pd.DataFrame(rows)


def disagreement_summary(disagreement: pd.DataFrame) -> pd.DataFrame:
    """Per-model outcome rates for the stacked disagreement chart."""

    if disagreement.empty:
        return pd.DataFrame()
    counts = (
        disagreement.groupby(["model", "model_key", "outcome"], as_index=False)
        .size()
        .rename(columns={"size": "n"})
    )
    totals = (
        disagreement.groupby(["model", "model_key"], as_index=False)
        .size()
        .rename(columns={"size": "total"})
    )
    out = counts.merge(totals, on=["model", "model_key"])
    out["rate"] = out["n"] / out["total"]
    return out


def disagreement_by_slice(disagreement: pd.DataFrame) -> pd.DataFrame:
    """Per-model × query_type: share where only one side is correct."""

    if disagreement.empty:
        return pd.DataFrame()
    rows: list[dict[str, Any]] = []
    for (model, model_key, qt), group in disagreement.groupby(
        ["model", "model_key", "query_type"]
    ):
        n = len(group)
        if n == 0:
            continue
        rfdt_only = int((group["outcome"] == "rfdt_only").sum())
        enc_only = int((group["outcome"] == "encoder_only").sum())
        rows.append(
            {
                "model": model,
                "model_key": model_key,
                "query_type": qt,
                "query_type_label": QUERY_TYPE_LABELS.get(qt, qt),
                "n": n,
                "rfdt_only_rate": rfdt_only / n,
                "encoder_only_rate": enc_only / n,
                "disagree_rate": (rfdt_only + enc_only) / n,
                "net_rfdt_edge": (rfdt_only - enc_only) / n,
            }
        )
    if not rows:
        return pd.DataFrame()
    return pd.DataFrame(rows)


def lift_vs_base(df: pd.DataFrame, metric: str = "macro_f1") -> pd.DataFrame:
    if df.empty or metric not in df.columns:
        return pd.DataFrame()

    base = (
        df[df["condition"] == "base"][["model_key", "model", metric]]
        .rename(columns={metric: "base_value"})
        .drop_duplicates("model_key")
    )
    merged = df.merge(base, on=["model_key", "model"], how="left")
    # Lift chart: generation ladder vs Base. Keep RFDT / encoder out
    # (different objectives / architectures).
    merged = merged[
        (merged["condition"] != "base")
        & (~merged["condition"].isin(RFDT_CONDITIONS))
        & (~merged["condition"].isin(ENCODER_CONDITIONS))
    ].copy()
    merged["delta"] = merged[metric] - merged["base_value"]
    return merged


def generative_ladder_frame(df: pd.DataFrame) -> pd.DataFrame:
    """Drop encoder-comparator rows from the 8B generation ladder charts."""

    if df.empty or "family" not in df.columns:
        return df
    return df[df["family"] != "encoder"].reset_index(drop=True)


def encoder_comparator_frame(df: pd.DataFrame) -> pd.DataFrame:
    """RFDT-LoRA (matched 8B) + ModernBERT Encoder-FT for the twin chart."""

    if df.empty:
        return df
    keep = df[
        (df["condition"] == "rfdt_lora")
        | ((df["condition"] == "encoder_ft") & (df["model_key"].isin(ENCODER_MODEL_KEYS)))
    ].copy()
    return keep.reset_index(drop=True)


def best_by_family(df: pd.DataFrame, metric: str = "macro_f1") -> dict[str, Any]:
    """Headline stats for datacards."""
    empty = {
        "best_frozen": None,
        "best_lora": None,
        "best_kv": None,
        "best_rfdt": None,
        "best_encoder": None,
        "n_runs": 0,
    }
    if df.empty or metric not in df.columns:
        return empty

    usable = df.dropna(subset=[metric])
    if usable.empty:
        return empty

    frozen = usable[usable["family"] == "frozen"]
    lora = usable[usable["condition"] == "lora"]
    kv = usable[usable["condition"] == "cartridge"]
    rfdt = usable[usable["condition"] == "rfdt_lora"]
    encoder = usable[usable["condition"] == "encoder_ft"]

    def _best(frame: pd.DataFrame) -> dict[str, Any] | None:
        if frame.empty:
            return None
        row = frame.loc[frame[metric].idxmax()]
        return {
            "model": row["model"],
            "condition_label": row["condition_label"],
            "value": float(row[metric]),
            "run_id": row["run_id"],
        }

    return {
        "best_frozen": _best(frozen),
        "best_lora": _best(lora),
        "best_kv": _best(kv),
        "best_rfdt": _best(rfdt),
        "best_encoder": _best(encoder),
        "n_runs": int(len(usable)),
    }


# ---------------------------------------------------------------------------
# Published analysis artifacts (analysis/catalog.json on the experiments dataset)
# ---------------------------------------------------------------------------

ANALYSIS_CATALOG_PATH = "analysis/catalog.json"

_ANALYSIS_LOCK = threading.Lock()
_ANALYSIS_CACHE: dict[str, Any] = {
    "loaded_at": None,
    "error": None,
    "catalog": None,
    "by_id": {},
}


def refresh_analysis_cache(*, force: bool = False) -> dict[str, Any]:
    """Load analysis/catalog.json and each results.json listed there."""
    with _ANALYSIS_LOCK:
        if not force and _ANALYSIS_CACHE["loaded_at"] is not None:
            return _ANALYSIS_CACHE

        if not force:
            snap = _read_analysis_snapshot()
            if snap is not None:
                _ANALYSIS_CACHE["catalog"] = snap.get("catalog")
                _ANALYSIS_CACHE["by_id"] = dict(snap.get("by_id") or {})
                _ANALYSIS_CACHE["loaded_at"] = snap.get("loaded_at")
                _ANALYSIS_CACHE["error"] = snap.get("error")
                return _ANALYSIS_CACHE

        token = hf_token()
        catalog = _download_metrics_json(
            repo_id=DEFAULT_EXPERIMENTS_REPO,
            path_in_repo=ANALYSIS_CATALOG_PATH,
            token=token,
        )
        by_id: dict[str, dict[str, Any]] = {}
        error: str | None = None
        if catalog is None:
            error = (
                "No published analysis catalog yet "
                f"({DEFAULT_EXPERIMENTS_REPO}/{ANALYSIS_CATALOG_PATH})."
            )
        else:
            for entry in catalog.get("analyses") or []:
                path = entry.get("path")
                analysis_id = entry.get("analysis_id")
                if not path or not analysis_id:
                    continue
                payload = _download_metrics_json(
                    repo_id=DEFAULT_EXPERIMENTS_REPO,
                    path_in_repo=str(path),
                    token=token,
                )
                if payload is not None:
                    by_id[str(analysis_id)] = payload

        _ANALYSIS_CACHE["catalog"] = catalog
        _ANALYSIS_CACHE["by_id"] = by_id
        _ANALYSIS_CACHE["loaded_at"] = time.time()
        _ANALYSIS_CACHE["error"] = error
        try:
            _write_analysis_snapshot(_ANALYSIS_CACHE)
        except OSError:
            pass
        return _ANALYSIS_CACHE


def get_analysis_catalog() -> dict[str, Any] | None:
    cache = refresh_analysis_cache()
    return cache.get("catalog")


def get_analysis(analysis_id: str) -> dict[str, Any] | None:
    cache = refresh_analysis_cache()
    return cache.get("by_id", {}).get(analysis_id)


def bootstrap_contrasts_frame(
    domain: str | None = None,
    *,
    analysis_id: str = "leakage_bootstrap",
) -> pd.DataFrame:
    """Flatten published leakage-bootstrap contrasts; optional domain filter."""
    blob = get_analysis(analysis_id)
    if not blob:
        return pd.DataFrame()
    payload = blob.get("payload") or {}
    rows = payload.get("contrasts") or []
    if not rows:
        return pd.DataFrame()
    frame = pd.DataFrame(rows)
    if domain and "domain" in frame.columns:
        frame = frame[frame["domain"] == domain].reset_index(drop=True)
    return frame


def mcnemar_contrasts_frame(
    domain: str | None = None,
    *,
    analysis_id: str = "mcnemar",
) -> pd.DataFrame:
    """Flatten published McNemar contrasts; optional domain filter."""
    blob = get_analysis(analysis_id)
    if not blob:
        return pd.DataFrame()
    payload = blob.get("payload") or {}
    rows = payload.get("contrasts") or []
    if not rows:
        return pd.DataFrame()
    frame = pd.DataFrame(rows)
    if domain and "domain" in frame.columns:
        frame = frame[frame["domain"] == domain].reset_index(drop=True)
    return frame


def analysis_cache_status() -> dict[str, Any]:
    with _ANALYSIS_LOCK:
        catalog = _ANALYSIS_CACHE.get("catalog") or {}
        return {
            "loaded_at": _ANALYSIS_CACHE.get("loaded_at"),
            "error": _ANALYSIS_CACHE.get("error"),
            "n_analyses": len(_ANALYSIS_CACHE.get("by_id") or {}),
            "analysis_ids": sorted((_ANALYSIS_CACHE.get("by_id") or {}).keys()),
            "catalog_updated_at_utc": catalog.get("updated_at_utc"),
        }