"""Build the v0.2 evaluation data from Imagenette (permissive, no auth):

  1. An EXEMPLAR pool  -> examples/exemplars/<class_word>/*.jpg   (for M4 centroids)
  2. An ATTACK set     -> examples/testset60/ + examples/testset60.csv

The two splits are disjoint per class (exemplars taken first, attack images taken
after), so decoy/truth feature centroids are never computed from an attacked image.

Usage (run on a box with `datasets`):
  python scripts/build_eval_data.py --exemplars 24 --attack 6 --split validation
"""

from __future__ import annotations

import argparse
import csv
from pathlib import Path

IMAGENETTE_TRUTH = {
    0: "tench", 1: "english springer", 2: "cassette player", 3: "chainsaw",
    4: "church", 5: "french horn", 6: "garbage truck", 7: "gas pump",
    8: "golf ball", 9: "parachute",
}


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--exemplars", type=int, default=24, help="imgs/class for centroid pool")
    ap.add_argument("--attack", type=int, default=6, help="imgs/class for the attack set")
    ap.add_argument("--split", default="validation")
    ap.add_argument("--exemplar-dir", default="examples/exemplars")
    ap.add_argument("--attack-dir", default="examples/testset60")
    ap.add_argument("--manifest", default="examples/testset60.csv")
    args = ap.parse_args()

    from datasets import load_dataset

    ds = load_dataset("frgfm/imagenette", "320px", split=args.split)

    ex_root = Path(args.exemplar_dir)
    at_root = Path(args.attack_dir)
    at_root.mkdir(parents=True, exist_ok=True)

    seen: dict[int, int] = {}
    rows: list[tuple[str, str]] = []
    need = args.exemplars + args.attack
    for ex in ds:
        lbl = int(ex["label"])
        k = seen.get(lbl, 0)
        if k >= need:
            continue
        truth = IMAGENETTE_TRUTH[lbl]
        img = ex["image"].convert("RGB")
        if k < args.exemplars:
            d = ex_root / truth.replace(" ", "_")
            d.mkdir(parents=True, exist_ok=True)
            img.save(d / f"{k:02d}.jpg", quality=95)
        else:
            j = k - args.exemplars
            fp = at_root / f"{lbl:02d}_{j:02d}.jpg"
            img.save(fp, quality=95)
            rows.append((str(fp), truth))
        seen[lbl] = k + 1
        if all(seen.get(i, 0) >= need for i in IMAGENETTE_TRUTH):
            break

    with Path(args.manifest).open("w", newline="") as f:
        w = csv.writer(f)
        for path, truth in rows:
            w.writerow([path, truth])
    ex_total = sum(min(v, args.exemplars) for v in seen.values())
    print(f"exemplars: {ex_total} imgs across {len(IMAGENETTE_TRUTH)} classes -> {ex_root}")
    print(f"attack set: {len(rows)} imgs -> {at_root} + {args.manifest}")


if __name__ == "__main__":
    main()