"""Build a standard labeled test set from Imagenette (permissive, no auth).

Samples N images across the 10 Imagenette classes, saves JPEGs to
examples/testset/, and writes examples/testset.csv (image_path,truth_label).

Imagenette classes map to plain single-word truths a VLM should nail on a clean
image, which is exactly what we want to try to poison.

Usage:
  python scripts/build_testset.py --n 40 --split validation
"""

from __future__ import annotations

import argparse
import csv
from pathlib import Path

# Imagenette label index -> a clean human truth word.
IMAGENETTE_TRUTH = {
    0: "tench",           # a fish
    1: "english springer",  # a dog
    2: "cassette player",
    3: "chainsaw",
    4: "church",
    5: "french horn",
    6: "garbage truck",
    7: "gas pump",
    8: "golf ball",
    9: "parachute",
}

# Rough difficulty tier per class (for a balanced sample):
#   iconic/unambiguous (hard to poison), ordinary, ambiguous-ish
TIER = {
    "iconic": [8, 4, 6],       # golf ball, church, garbage truck
    "ordinary": [1, 3, 5, 2],  # dog, chainsaw, french horn, cassette player
    "ambiguous": [0, 7, 9],    # tench, gas pump, parachute
}


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=40)
    ap.add_argument("--split", default="validation")
    ap.add_argument("--out", default="examples/testset")
    args = ap.parse_args()

    from datasets import load_dataset

    ds = load_dataset("frgfm/imagenette", "320px", split=args.split)
    label_names = ds.features["label"].names if hasattr(ds.features["label"], "names") else None

    out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)
    rows: list[tuple[str, str]] = []

    # Round-robin across tiers so the set spans difficulty evenly.
    per_class = max(1, args.n // 10)
    counts: dict[int, int] = {}
    for ex in ds:
        lbl = int(ex["label"])
        if counts.get(lbl, 0) >= per_class:
            continue
        img = ex["image"].convert("RGB")
        truth = IMAGENETTE_TRUTH.get(lbl, (label_names[lbl] if label_names else str(lbl)))
        fname = out_dir / f"{lbl:02d}_{counts.get(lbl,0):02d}.jpg"
        img.save(fname, quality=95)
        rows.append((str(fname), truth))
        counts[lbl] = counts.get(lbl, 0) + 1
        if len(rows) >= args.n:
            break

    csv_path = Path("examples/testset.csv")
    with csv_path.open("w", newline="") as f:
        w = csv.writer(f)
        for path, truth in rows:
            w.writerow([path, truth])
    print(f"wrote {len(rows)} images to {out_dir} and manifest {csv_path}")


if __name__ == "__main__":
    main()