"""
datasets.py — Ground-truth providers.

`GTSample` is the uniform shape every loader yields: an image (PIL, or None for
the torch-free smoke set), the per-category ground truth, and the image size used
for coordinate normalization. The packaged smoke set runs on CPU with no network
so tests and `--dataset smoke --runner stub` work offline. Real loaders stream
public datasets via HF `datasets` (imported lazily) for Phase 1+.

GT shapes (per category):
  image_classification : {"labels": [acceptable label strings]}
  bbox_grounding       : {"boxes": [{"label": str, "bbox": [x,y,w,h], "fmt": "xywh"}]}
  ocr_text             : {"text": "the reference transcription / answer"}
  (stub categories)    : None  (no GT wired yet)
"""

from __future__ import annotations

import itertools
from dataclasses import dataclass, field
from typing import Any, Callable, Optional


@dataclass
class GTSample:
    image: Any                  # PIL.Image or None (smoke / stub)
    prompt: str
    gt: Any
    category: str
    image_id: str
    size: tuple[int, int]       # (W, H) for coordinate normalization
    meta: dict = field(default_factory=dict)


# ──────────────────────────────────────────────────────────────────────────────
# Packaged CPU smoke set (no network, no torch). Small but exercises every shape.
# ──────────────────────────────────────────────────────────────────────────────

_SMOKE: dict[str, list[GTSample]] = {
    "image_classification": [
        GTSample(None, "Classify this image.", {"labels": ["golden retriever", "dog"]},
                 "image_classification", "smk_cls_0", (640, 480)),
        GTSample(None, "Classify this image.", {"labels": ["espresso", "coffee"]},
                 "image_classification", "smk_cls_1", (512, 512)),
        GTSample(None, "Classify this image.", {"labels": ["school bus", "bus"]},
                 "image_classification", "smk_cls_2", (800, 600)),
    ],
    "bbox_grounding": [
        GTSample(None, "Detect all objects.",
                 {"boxes": [{"label": "dog", "bbox": [64, 48, 128, 96], "fmt": "xywh"}]},
                 "bbox_grounding", "smk_box_0", (640, 480)),
        GTSample(None, "Detect all objects.",
                 {"boxes": [{"label": "cat", "bbox": [10, 10, 40, 40], "fmt": "xywh"},
                            {"label": "ball", "bbox": [200, 150, 50, 50], "fmt": "xywh"}]},
                 "bbox_grounding", "smk_box_1", (640, 480)),
    ],
    "ocr_text": [
        GTSample(None, "Read all text.", {"text": "STOP"}, "ocr_text", "smk_ocr_0", (200, 200)),
        GTSample(None, "Read all text.", {"text": "no entry"}, "ocr_text", "smk_ocr_1", (300, 200)),
    ],
}


def smoke_samples(category: str, n: Optional[int] = None) -> list[GTSample]:
    """Smoke samples for a category. Stub categories (no GT) get synthetic blanks."""
    if category in _SMOKE:
        out = _SMOKE[category]
    else:
        out = [GTSample(None, "Analyze this image.", None, category, f"smk_{category}_{i}", (64, 64))
               for i in range(2)]
    return out[:n] if n else list(out)


# ──────────────────────────────────────────────────────────────────────────────
# Real loaders (Phase 1+). HF `datasets` is imported lazily so the smoke path and
# the CPU tests never require it.
# ──────────────────────────────────────────────────────────────────────────────

def _hf_stream(repo: str, split: str, n: int, **kw):
    from datasets import load_dataset  # lazy
    ds = load_dataset(repo, split=split, streaming=True, **kw)
    return list(itertools.islice(ds, n))


def load_imagenet_val(n: int = 200, split: str = "validation") -> list[GTSample]:
    """Classification GT. ImageNet-1k is GATED and its label is a bare integer id;
    use food101 (ungated parquet) and map the ClassLabel id -> class name."""
    from datasets import load_dataset  # lazy
    last = None
    for repo, sp in [("ethz/food101", "validation"), ("food101", "validation")]:
        try:
            ds = load_dataset(repo, split=sp, streaming=True)
            try:
                names = ds.features["label"].names
            except Exception:
                names = None
            rows = list(itertools.islice(ds, n))
            out = []
            for i, r in enumerate(rows):
                img = r.get("image")
                lbl = r.get("label")
                name = (names[lbl].replace("_", " ")
                        if names and isinstance(lbl, int) and 0 <= lbl < len(names) else str(lbl))
                size = (img.width, img.height) if img is not None else (0, 0)
                out.append(GTSample(img, "Classify this image.", {"labels": [name]},
                                    "image_classification", f"cls_{i}", size))
            if out:
                return out
        except Exception as e:
            last = e
            continue
    raise RuntimeError(f"no classification dataset streamable (food101): {last}")


# COCO-80 class names in category-id order (confirmed from detection-datasets/coco
# ClassLabel features). objects.bbox is [x1,y1,x2,y2] in absolute pixels (xyxy).
COCO_CLASSES = [
    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat",
    "traffic light", "fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat",
    "dog", "horse", "sheep", "cow", "elephant", "bear", "zebra", "giraffe", "backpack",
    "umbrella", "handbag", "tie", "suitcase", "frisbee", "skis", "snowboard", "sports ball",
    "kite", "baseball bat", "baseball glove", "skateboard", "surfboard", "tennis racket",
    "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
    "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair",
    "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse",
    "remote", "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator",
    "book", "clock", "vase", "scissors", "teddy bear", "hair drier", "toothbrush",
]


def _coco_label(cid) -> str:
    if isinstance(cid, int) and 0 <= cid < len(COCO_CLASSES):
        return COCO_CLASSES[cid]
    return str(cid)


def load_coco_detection(n: int = 200, split: str = "val") -> list[GTSample]:
    """detection-datasets/coco: objects.bbox is xyxy pixels; category is a ClassLabel id."""
    rows = _hf_stream("detection-datasets/coco", split, n)
    out = []
    for i, r in enumerate(rows):
        img = r["image"]
        objs = r.get("objects", {})
        boxes = [{"label": _coco_label(c), "bbox": list(map(float, b)), "fmt": "xyxy"}
                 for c, b in zip(objs.get("category", []), objs.get("bbox", []))]
        out.append(GTSample(img, "Detect all objects in this image. Output only the raw JSON object.",
                            {"boxes": boxes}, "bbox_grounding", f"coco_{i}", (img.width, img.height)))
    return out


def load_textvqa(n: int = 200, split: str = "validation") -> list[GTSample]:
    """OCR GT. The script-based 'textvqa' repo is rejected by modern `datasets`; use
    parquet repos. GT = {"text": <gold answer>}; the model transcribes the image and
    the OCR scorer credits containment of the answer."""
    from datasets import load_dataset  # lazy
    last = None
    for repo, sp in [("lmms-lab/textvqa", "validation"), ("howard-hou/OCR-VQA", "test")]:
        try:
            ds = load_dataset(repo, split=sp, streaming=True)
            rows = list(itertools.islice(ds, n))
            out = []
            for i, r in enumerate(rows):
                img = r.get("image")
                ans = r.get("answers")
                if ans is None:
                    ans = r.get("answer") or r.get("questions")
                if isinstance(ans, (list, tuple)):
                    ans = next((str(a) for a in ans if str(a).strip()), "")
                size = (img.width, img.height) if img is not None else (0, 0)
                out.append(GTSample(img, "Read all the text in this image.",
                                    {"text": str(ans or "")}, "ocr_text", f"ocr_{i}", size))
            if out:
                return out
        except Exception as e:
            last = e
            continue
    raise RuntimeError(f"no OCR dataset streamable (textvqa/ocr-vqa): {last}")


# ──────────────────────────────────────────────────────────────────────────────
# Synthetic data-format images (self-contained: no external dataset). Renders a
# small record in several serialization formats to an image, with exact GT for
# both the format (data_type) and the normalized content. Tests whether a VLM can
# recognize a data format from a screenshot and re-serialize it to JSON.
# ──────────────────────────────────────────────────────────────────────────────

_DATATYPE_RECORDS = [
    {"name": "Alice", "age": "30", "city": "Paris"},
    {"id": "7", "title": "Widget", "price": "9"},
    {"user": "bob", "active": "true", "score": "42"},
    {"country": "Japan", "capital": "Tokyo", "pop": "14"},
]


def _datatype_font(sz=22):
    from PIL import ImageFont
    for name in ("DejaVuSansMono.ttf", "DejaVuSans.ttf"):
        try:
            return ImageFont.truetype(name, sz)
        except Exception:
            continue
    try:
        return ImageFont.load_default(size=sz)  # Pillow >= 10
    except Exception:
        return ImageFont.load_default()


def _render_text_image(text: str, size=(640, 360)) -> "object":
    from PIL import Image, ImageDraw
    img = Image.new("RGB", size, (255, 255, 255))
    d = ImageDraw.Draw(img)
    d.multiline_text((18, 18), text, fill=(0, 0, 0), font=_datatype_font(22), spacing=8)
    return img


def _serialize(rec: dict, fmt: str) -> str:
    if fmt == "json":
        import json as _j
        return _j.dumps(rec, indent=2)
    if fmt == "yaml":
        return "\n".join(f"{k}: {v}" for k, v in rec.items())
    if fmt == "toml":
        return "\n".join(f'{k} = "{v}"' for k, v in rec.items())
    if fmt == "xml":
        inner = "".join(f"<{k}>{v}</{k}>" for k, v in rec.items())
        return f"<record>{inner}</record>"
    if fmt == "csv":
        return ",".join(rec.keys()) + "\n" + ",".join(rec.values())
    if fmt == "markdown":
        return "# Record\n" + "\n".join(f"- **{k}**: {v}" for k, v in rec.items())
    raise ValueError(fmt)


_DATATYPE_FORMATS = ["json", "yaml", "toml", "xml", "csv", "markdown"]


def make_datatype_samples(n: int = 18, split=None) -> list[GTSample]:
    """Self-contained: render records across formats. GT = {data_type, content}."""
    out = []
    i = 0
    while len(out) < n:
        rec = _DATATYPE_RECORDS[i % len(_DATATYPE_RECORDS)]
        fmt = _DATATYPE_FORMATS[i % len(_DATATYPE_FORMATS)]
        text = _serialize(rec, fmt)
        # csv normalizes to a one-row list; everything else to the dict
        content = [rec] if fmt == "csv" else rec
        img = _render_text_image(text)
        out.append(GTSample(img, "Identify the data format and contents. Output only raw JSON.",
                            {"data_type": fmt, "content": content},
                            "data_type", f"dt_{i}", img.size))
        i += 1
    return out


# ──────────────────────────────────────────────────────────────────────────────
# Synthetic colored-shapes scenes (self-contained). One scene yields exact GT for
# spatial relations (by x-order), depth ordering (bigger circle = nearer), and
# subject fixation (largest circle = primary subject). Reliable + no download.
# ──────────────────────────────────────────────────────────────────────────────

import itertools as _it

_SHAPE_COLORS = {"red": (220, 30, 30), "green": (30, 170, 30), "blue": (40, 40, 220)}
_SHAPE_NAMES = ["red", "green", "blue"]
_SHAPE_SIZES = [110, 76, 46]  # diameters: big / medium / small (depth cue)


def _shape_scene(i: int):
    """Deterministic 3-circle scene. Returns ((W,H), [shape dicts]) sorted left→right."""
    W, H = 540, 320
    x_centers = [110, 270, 430]
    color_perm = list(_it.permutations(range(3)))[i % 6]      # which color in which column
    size_perm = list(_it.permutations(range(3)))[(i // 6) % 6]  # which color gets which size
    shapes = []
    for ci, color in enumerate(_SHAPE_NAMES):
        cx = x_centers[color_perm[ci]]
        d = _SHAPE_SIZES[size_perm[ci]]
        cy = H // 2
        shapes.append({"label": color, "cx": cx, "cy": cy, "d": d, "area": d * d,
                       "bbox": [cx - d / 2, cy - d / 2, cx + d / 2, cy + d / 2]})
    shapes.sort(key=lambda s: s["cx"])
    return (W, H), shapes


def _render_scene(size, shapes):
    from PIL import Image, ImageDraw
    img = Image.new("RGB", size, (245, 245, 245))
    d = ImageDraw.Draw(img)
    for s in shapes:
        d.ellipse(s["bbox"], fill=_SHAPE_COLORS[s["label"]])
    return img


def make_shapes_samples(n: int = 12, split=None) -> list[GTSample]:
    """Scenes carrying GT for spatial / depth / subject_fixation simultaneously."""
    out = []
    for i in range(n):
        size, shapes = _shape_scene(i)
        # spatial: left_of for every left→right pair
        triples = []
        for a, b in _it.combinations(shapes, 2):  # already x-sorted → a left of b
            triples.append([a["label"], "left_of", b["label"]])
        # depth: bigger area = nearer
        pairs = []
        for a, b in _it.combinations(shapes, 2):
            pairs.append({"a": a["label"], "b": b["label"],
                          "a_is": "nearer" if a["area"] >= b["area"] else "farther"})
        # subject: largest area
        subj = max(shapes, key=lambda s: s["area"])
        gt = {"triples": triples, "pairs": pairs,
              "label": subj["label"], "box": subj["bbox"], "fmt": "xyxy"}
        img = _render_scene(size, shapes)
        out.append(GTSample(img, "Analyze the colored shapes. Output only raw JSON.",
                            gt, "shapes", f"shapes_{i}", size))
    return out


def _circle_polygon(cx, cy, d, n=16):
    """Approximate a circle (diameter d, center cx,cy) as a flat pixel-coord
    polygon [x1,y1,x2,y2,...] with n vertices."""
    import math
    r = d / 2.0
    flat = []
    for k in range(n):
        ang = 2.0 * math.pi * k / n
        flat.append(cx + r * math.cos(ang))
        flat.append(cy + r * math.sin(ang))
    return flat


def make_segmentation_samples(n: int = 12, split=None) -> list[GTSample]:
    """Self-contained instance-segmentation GT: reuse the 3-circle shape scenes.
    Each colored circle becomes one mask whose polygon is the circle approximated
    by 16 vertices (label = color). Polygons are in PIXEL coords; the scorer
    converts model polygons from NORM_0_1000 to pixels."""
    out = []
    for i in range(n):
        size, shapes = _shape_scene(i)
        masks = [{"label": s["label"],
                  "polygon_pixels": _circle_polygon(s["cx"], s["cy"], s["d"], n=16)}
                 for s in shapes]
        img = _render_scene(size, shapes)
        out.append(GTSample(img, "Segment the colored shapes as labeled polygons. Output only raw JSON.",
                            {"masks": masks}, "segmentation", f"seg_{i}", size))
    return out


def make_outline_samples(n: int = 12, split=None) -> list[GTSample]:
    """Self-contained: reuse the 3-circle synthetic scene. GT outline = the largest
    circle approximated as a 16-point polygon (pixels), label = its color."""
    out = []
    for i in range(n):
        size, shapes = _shape_scene(i)
        main = max(shapes, key=lambda s: s["area"])        # largest = main object
        poly = _circle_polygon(main["cx"], main["cy"], main["d"], 16)
        gt = {"outline": poly, "label": main["label"], "bbox": main["bbox"], "fmt": "xyxy"}
        img = _render_scene(size, shapes)
        out.append(GTSample(img, "Trace the main object's outline. Output only raw JSON.",
                            gt, "outline_association", f"outline_{i}", size))
    return out


_BOX3D_COLORS = {"red": (220, 40, 40), "green": (40, 175, 40), "blue": (50, 50, 225)}
_BOX3D_NAMES = ["red", "green", "blue"]


def _box3d_scene(i: int):
    """Deterministic 2-3 colored boxes at known ground (x,z) positions.

    GT convention (normalized 0..1): bbox3d = [x, y, z, w, h, l, yaw] with
    x = left-right ground position, z = depth (0 near .. 1 far), y = 0 (on the
    floor), (w,h,l) the box footprint width / height / length, yaw = 0. The GT is
    exact-by-construction; the render is a simplified ground-plane 3D proxy.
    """
    import math
    import itertools
    W, H = 480, 360
    n_boxes = 2 + (i % 2)                       # 2 or 3 boxes
    names = _BOX3D_NAMES[:n_boxes]
    perm = list(itertools.permutations(range(n_boxes)))[i % math.factorial(n_boxes)]
    x_slots = [0.2, 0.5, 0.8][:n_boxes]
    z_slots = [0.25, 0.55, 0.85][:n_boxes]
    objects, draw = [], []
    for k, color in enumerate(names):
        x = x_slots[perm[k] % n_boxes]
        z = z_slots[k]                          # increasing depth per index
        w = 0.16 + 0.04 * ((i + k) % 3)         # footprint width
        l = 0.14
        h = 0.22 + 0.03 * (k % 2)               # box height
        objects.append({"class": color,
                        "bbox3d": [round(x, 4), 0.0, round(z, 4),
                                   round(w, 4), round(h, 4), round(l, 4), 0.0]})
        draw.append((color, x, z, w, h))
    return (W, H), objects, draw


def _render_box3d_scene(size, draw):
    """Perspective proxy: nearer (small z) boxes drawn lower in frame and larger."""
    from PIL import Image, ImageDraw
    W, H = size
    img = Image.new("RGB", size, (235, 235, 240))
    d = ImageDraw.Draw(img)
    d.rectangle([0, int(H * 0.5), W, H], fill=(205, 200, 190))   # ground band
    for color, x, z, w, h in sorted(draw, key=lambda t: t[2], reverse=True):  # far first
        scale = 1.0 - 0.45 * z                                  # nearer = bigger
        bw = w * W * scale
        bh = h * H * scale
        cx = x * W
        cy = (0.5 + 0.45 * z) * H                               # nearer = lower
        d.rectangle([cx - bw / 2, cy - bh, cx + bw / 2, cy],
                    fill=_BOX3D_COLORS[color], outline=(20, 20, 20))
    return img


def make_3d_samples(n: int = 12, split=None) -> list[GTSample]:
    """Self-contained synthetic 3D scenes. GT exact-by-construction (proxy)."""
    out = []
    for i in range(n):
        size, objects, draw = _box3d_scene(i)
        img = _render_box3d_scene(size, draw)
        out.append(GTSample(img, "Identify the 3D boxes. Output only raw JSON.",
                            {"objects": objects}, "geometric_3d_object_id",
                            f"box3d_{i}", size))
    return out


def make_camera_samples(n: int = 12, split=None) -> list[GTSample]:
    """Self-contained synthetic camera-roll set. A clear orientation cue (an upward
    arrow over a horizon line) is drawn upright, then the whole image is rotated by a
    KNOWN roll angle that varies by index; yaw=pitch=0 (a single 2D cue cannot
    disambiguate yaw/pitch). GT = {"rotation": [0, 0, roll_deg]}.

    NOTE: SIMPLIFIED proxy — this tests recovery of ROLL from a 2D cue only; it does
    not exercise yaw/pitch (which would need a 3D scene). Reliable, no download.
    """
    from PIL import Image, ImageDraw

    W, H = 480, 480
    cx, cy = W / 2.0, H / 2.0
    # deterministic spread of rolls across the wrapped range, indexed by sample
    roll_table = [0, 15, 30, 45, 60, 90, -15, -30, -45, -60, -90, 120,
                  -120, 150, 75, -75, 10, -10]
    out = []
    for i in range(n):
        roll = float(roll_table[i % len(roll_table)])
        base = Image.new("RGB", (W, H), (250, 250, 250))
        d = ImageDraw.Draw(base)
        d.line([(60, cy), (W - 60, cy)], fill=(60, 60, 60), width=6)               # horizon line
        d.line([(cx, cy), (cx, 90)], fill=(200, 40, 40), width=8)                  # arrow shaft (points up)
        d.polygon([(cx, 60), (cx - 22, 105), (cx + 22, 105)], fill=(200, 40, 40))  # arrow head
        # Rotate scene by -roll about the centre (expand=False keeps size + GT stable):
        # a positive camera roll (CW) rotates scene content CCW in the image.
        img = base.rotate(-roll, resample=Image.BICUBIC, center=(cx, cy),
                          fillcolor=(250, 250, 250), expand=False)
        gt = {"rotation": [0.0, 0.0, roll]}
        out.append(GTSample(img,
                            "Estimate the camera rotation [yaw, pitch, roll]. Output only raw JSON.",
                            gt, "camera_rotational_offset", f"camrot_{i}", (W, H)))
    return out


def make_gqa_samples(n: int = 200, split: str = "validation") -> list[GTSample]:
    """Grounded-VQA GT (REAL, best-effort). Streams a VQA dataset; one sample per
    (image, question, answers). The question is per-image and goes in
    GTSample.prompt; gt = {"answers": [<gold strings>]}. Image is row["image"]
    (a PIL image).

    Repo ids are BEST-EFFORT — the maintainer must verify id/config/split:
      primary : "lmms-lab/GQA"   (testdev_balanced / val splits; row has
                                   "question" + "answer"; image under "image")
      fallback: "HuggingFaceM4/VQAv2" (row has "question" + "answers"
                                       list-of-dicts or list-of-strings)
    The answer-field probing below tolerates both shapes.
    """
    # Script-based repos (HuggingFaceM4/VQAv2, lmms-lab/GQA) are rejected by modern
    # `datasets`. Use PARQUET repos (verified format:parquet on the Hub), in order.
    rows = None
    for repo, sp in [("lmms-lab/VQAv2", split), ("merve/vqav2-small", "validation"),
                     ("merve/vqav2-small", "train"), ("lmms-lab/OK-VQA", "val2014")]:
        try:
            rows = _hf_stream(repo, sp, n)
            if rows:
                break
        except Exception:
            continue
    if not rows:
        raise RuntimeError("no parquet VQA dataset streamable "
                           "(tried lmms-lab/VQAv2, merve/vqav2-small, lmms-lab/OK-VQA)")
    out = []
    for i, r in enumerate(rows):
        img = r.get("image")
        question = str(r.get("question") or r.get("question_str") or "What is in this image?")
        raw_ans = r.get("answers")
        if raw_ans is None:
            raw_ans = r.get("multiple_choice_answer") or r.get("answer")
        if isinstance(raw_ans, dict):                 # {"answer": "x"} or value-map
            raw_ans = raw_ans.get("answer") or list(raw_ans.values())
        if isinstance(raw_ans, (list, tuple)):
            answers = []
            for a in raw_ans:
                if isinstance(a, dict):               # VQAv2: [{"answer": "x"}, ...]
                    a = a.get("answer", "")
                if str(a).strip():
                    answers.append(str(a))
        elif raw_ans is not None and str(raw_ans).strip():
            answers = [str(raw_ans)]
        else:
            answers = []
        size = (img.width, img.height) if img is not None else (0, 0)
        out.append(GTSample(img, question, {"answers": answers},
                            "vit_accuracy_to_prompt", f"vqa_{i}", size))
    return out


def make_semantic_samples(n: int = 12, split=None) -> list[GTSample]:
    """Self-contained colored-shapes scenes carrying GT semantic-association triples.

    Reuses the deterministic 3-circle scene (`_shape_scene`). Associations are
    derived purely from geometry so they are exact and reproducible:
      * left->right ordering   -> (left, "left_of", right) AND (right, "right_of", left)
      * adjacency (consecutive) -> (a, "near", b) for neighbouring shapes
      * taxonomy                -> (color, "is_a", "circle") for every shape
    GT shape: {"triples": [[a, relation, b], ...]} -- read directly by score_triples,
    which does tolerant subject/object matching + normalized-exact predicate matching.
    Relations are chosen so they round-trip cleanly through metrics._norm_pred
    (left_of/right_of/near/is_a stay identical after normalization).
    """
    out = []
    for i in range(n):
        size, shapes = _shape_scene(i)  # sorted left->right
        triples: list[list] = []
        # ordering relations over every left->right pair (both directions)
        for a, b in _it.combinations(shapes, 2):  # a is left of b
            triples.append([a["label"], "left_of", b["label"]])
            triples.append([b["label"], "right_of", a["label"]])
        # adjacency ("near") for consecutive shapes in the x-ordering
        for a, b in zip(shapes, shapes[1:]):
            triples.append([a["label"], "near", b["label"]])
        # taxonomic: each colored shape is a circle
        for s in shapes:
            triples.append([s["label"], "is_a", "circle"])
        img = _render_scene(size, shapes)
        out.append(GTSample(
            img,
            "List semantic associations between the shapes. Output only raw JSON.",
            {"triples": triples}, "semantic_association", f"semassoc_{i}", size,
        ))
    return out


def _style_font(sz=28):
    from PIL import ImageFont
    for name in ("DejaVuSans.ttf", "DejaVuSansMono.ttf"):
        try:
            return ImageFont.truetype(name, sz)
        except Exception:
            continue
    try:
        return ImageFont.load_default(size=sz)  # Pillow >= 10
    except Exception:
        return ImageFont.load_default()


def _render_style_image(style: str, size=(320, 320)):
    """Render a controllable, visually-distinguishable exemplar for each coarse style.
    photo: smooth RGB gradient (photographic continuous tone). painting: soft color blobs
    on canvas. sketch: black outlines on white. 3d_render: lit/shaded sphere. anime:
    flat-shaded face with big eyes. other: a labelled fallback."""
    from PIL import Image, ImageDraw
    import math
    W, H = size
    cx, cy = W // 2, H // 2

    if style == "photo":
        img = Image.new("RGB", size, (0, 0, 0))
        px = img.load()
        for y in range(H):
            for x in range(W):
                px[x, y] = (int(40 + 180 * x / W), int(40 + 180 * y / H),
                            int(120 + 100 * ((x + y) % 50) / 50))
        return img

    if style == "painting":
        img = Image.new("RGB", size, (235, 225, 205))
        d = ImageDraw.Draw(img)
        for (bx, by), r, col in [((90, 90), 70, (200, 70, 60)),
                                 ((210, 120), 60, (70, 110, 190)),
                                 ((140, 220), 80, (90, 170, 90))]:
            d.ellipse([bx - r, by - r, bx + r, by + r], fill=col)
        return img

    if style == "sketch":
        img = Image.new("RGB", size, (255, 255, 255))
        d = ImageDraw.Draw(img)
        d.rectangle([cx - 70, cy - 70, cx + 70, cy + 70], outline=(0, 0, 0), width=3)
        d.line([cx - 70, cy - 70, cx + 70, cy + 70], fill=(0, 0, 0), width=2)
        d.line([cx + 70, cy - 70, cx - 70, cy + 70], fill=(0, 0, 0), width=2)
        d.ellipse([cx - 40, cy - 40, cx + 40, cy + 40], outline=(0, 0, 0), width=2)
        return img

    if style == "3d_render":
        img = Image.new("RGB", size, (245, 245, 250))
        d = ImageDraw.Draw(img)
        r = 90
        for yy in range(cy - r, cy + r):
            for xx in range(cx - r, cx + r):
                dx, dy = (xx - cx) / r, (yy - cy) / r
                if dx * dx + dy * dy <= 1.0:
                    lx, ly = -0.5, -0.6
                    nz = math.sqrt(max(0.0, 1.0 - dx * dx - dy * dy))
                    shade = max(0.12, (-dx * lx - dy * ly + nz) / 1.7)
                    v = int(60 + 195 * min(1.0, shade))
                    d.point((xx, yy), fill=(v, int(v * 0.7), int(v * 0.5)))
        return img

    if style == "anime":
        img = Image.new("RGB", size, (250, 240, 230))
        d = ImageDraw.Draw(img)
        d.ellipse([cx - 80, cy - 90, cx + 80, cy + 70], fill=(255, 224, 196),
                  outline=(40, 30, 30), width=3)
        for ex in (cx - 35, cx + 35):
            d.ellipse([ex - 18, cy - 10, ex + 18, cy + 30], fill=(255, 255, 255),
                      outline=(20, 20, 20), width=2)
            d.ellipse([ex - 10, cy + 2, ex + 10, cy + 26], fill=(60, 110, 200))
            d.ellipse([ex - 4, cy + 6, ex + 4, cy + 16], fill=(20, 20, 20))
        d.polygon([(cx - 90, cy - 90), (cx - 30, cy - 110), (cx, cy - 80)], fill=(90, 60, 40))
        return img

    # "other" fallback
    img = Image.new("RGB", size, (200, 200, 200))
    ImageDraw.Draw(img).text((20, H // 2), "other", fill=(0, 0, 0), font=_style_font(28))
    return img


# Each rendered style implies a controlled (layout, symmetry) GT pair.
_STYLE_LAYOUTS = {
    "photo": ("rule_of_thirds", "none"),
    "painting": ("scattered", "none"),
    "sketch": ("centered", "radial"),
    "3d_render": ("centered", "radial"),
    "anime": ("centered", "vertical"),
    "other": ("centered", "none"),
}
_STYLE_ORDER = ["photo", "painting", "sketch", "3d_render", "anime"]


def make_style_samples(n: int = 10, split=None) -> list[GTSample]:
    """Self-contained: render distinguishable styles we control. Cycles through
    photo/painting/sketch/3d_render/anime. GT = {style, layout, symmetry}."""
    out = []
    for i in range(n):
        style = _STYLE_ORDER[i % len(_STYLE_ORDER)]
        layout, symmetry = _STYLE_LAYOUTS[style]
        size = (320, 320)
        img = _render_style_image(style, size)
        gt = {"style": style, "layout": layout, "symmetry": symmetry}
        out.append(GTSample(img, "Classify the visual style and structure. Output only raw JSON.",
                            gt, "style_structural_awareness", f"style_{i}", size))
    return out


# ──────────────────────────────────────────────────────────────────────────────
# REAL COCO instance segmentation GT (for segmentation / outline / subject) — parses
# the official COCO annotations JSON directly (no script-dataset, no pycocotools) and
# pulls images by URL. Replaces the synthetic colored-shape GT with real images.
# ──────────────────────────────────────────────────────────────────────────────
_COCO_CACHE: dict = {}
_COCO_PERSON_CAT = 1     # COCO category_id for "person"


def _coco_ann_file(name: str) -> str:
    """Ensure `{cache_dir}/{name}` exists. ONE zip download extracts BOTH
    instances_val2017.json and captions_val2017.json (they ship in the same
    annotations_trainval2017.zip — extracting only one wastes the 241MB fetch)."""
    import io
    import os
    import urllib.request
    import zipfile

    cache_dir = os.environ.get("HF_HOME") or os.environ.get("TMPDIR") or "/tmp"
    os.makedirs(cache_dir, exist_ok=True)
    path = os.path.join(cache_dir, name)
    if not os.path.exists(path):
        print(f"  downloading COCO val2017 annotations (~241MB, one-time) for {name} …")
        zurl = "http://images.cocodataset.org/annotations/annotations_trainval2017.zip"
        zb = urllib.request.urlopen(zurl, timeout=600).read()
        with zipfile.ZipFile(io.BytesIO(zb)) as z:
            for member in ("instances_val2017.json", "captions_val2017.json"):
                target = os.path.join(cache_dir, member)
                if not os.path.exists(target):
                    with z.open(f"annotations/{member}") as f:
                        open(target, "wb").write(f.read())
    return path


def _coco_ann(kind: str = "instances") -> dict:
    """Parsed-JSON cache for the COCO annotation files. `kind` is "instances" or
    "captions". Keeps the existing _COCO_CACHE["ann"] key for instances."""
    import json as _json

    key = "ann" if kind == "instances" else f"ann_{kind}"
    if key not in _COCO_CACHE:
        with open(_coco_ann_file(f"{kind}_val2017.json"), encoding="utf-8") as f:
            _COCO_CACHE[key] = _json.load(f)
    return _COCO_CACHE[key]


def _coco_instances(n: int) -> list:
    """Returns [(image, (W,H), image_id, [{label, polygon_pixels, box_xyxy, area}])].
    Downloads + caches instances_val2017.json (~one-time) and the first `n` val images."""
    import io
    import urllib.request
    from collections import defaultdict
    from PIL import Image

    key = f"inst_{n}"
    if key in _COCO_CACHE:
        return _COCO_CACHE[key]
    data = _coco_ann("instances")
    cats = {c["id"]: c["name"] for c in data["categories"]}
    imgs = {im["id"]: im for im in data["images"]}
    anns = defaultdict(list)
    for a in data["annotations"]:
        anns[a["image_id"]].append(a)

    out = []
    for iid in list(imgs):
        if len(out) >= n:
            break
        info = imgs[iid]
        try:
            raw = urllib.request.urlopen(
                f"http://images.cocodataset.org/val2017/{info['file_name']}", timeout=60).read()
            img = Image.open(io.BytesIO(raw)).convert("RGB")
        except Exception:
            continue
        objs = []
        for a in anns[iid]:
            seg = a.get("segmentation")
            if a.get("iscrowd") or not isinstance(seg, list) or not seg:
                continue                                   # skip RLE / crowd
            poly = [float(v) for v in seg[0]]
            if len(poly) < 6:
                continue
            x, y, w, h = a["bbox"]
            objs.append({"label": cats.get(a["category_id"], "object"),
                         "polygon_pixels": poly, "box_xyxy": [x, y, x + w, y + h],
                         "area": float(a.get("area", w * h))})
        if objs:
            out.append((img, (img.width, img.height), f"coco_{iid}", objs))
    _COCO_CACHE[key] = out
    return out


def load_coco_segmentation(n: int = 24, split=None) -> list[GTSample]:
    return [GTSample(img, "Segment every object as a labeled polygon.",
                     {"masks": [{"label": o["label"], "polygon_pixels": o["polygon_pixels"]}
                                for o in objs]},
                     "segmentation", iid, size)
            for (img, size, iid, objs) in _coco_instances(n)]


def load_coco_outline(n: int = 24, split=None) -> list[GTSample]:
    out = []
    for (img, size, iid, objs) in _coco_instances(n):
        big = max(objs, key=lambda o: o["area"])
        out.append(GTSample(img, "Trace the main object's outline.",
                            {"outline": big["polygon_pixels"], "label": big["label"]},
                            "outline_association", iid, size))
    return out


def load_coco_subject(n: int = 24, split=None) -> list[GTSample]:
    out = []
    for (img, size, iid, objs) in _coco_instances(n):
        big = max(objs, key=lambda o: o["area"])
        out.append(GTSample(img, "Identify the primary subject.",
                            {"label": big["label"], "box": big["box_xyxy"], "fmt": "xyxy"},
                            "subject_fixation", iid, size))
    return out


# ── multi-person slice (fusion-tier validation GT) ────────────────────────────
def _select_multi_person_ids(ann: dict, *, min_persons: int = 2, max_persons: int = 6,
                             min_person_area_frac: float = 0.005,
                             require_nonperson: bool = False) -> list:
    """Image ids with TRUSTWORTHY multi-person GT: min..max non-crowd persons, no
    crowd-person annotation anywhere in the image (a crowd RLE blob means "many
    unlabeled people" — the count GT becomes untrustworthy), and no tiny background
    persons (< min_person_area_frac of the image). Deliberately a CLEAN slice; the
    bias is stated in every validation report. Pure filter over the parsed
    annotations — no network, testable with a fake ann dict."""
    from collections import defaultdict

    imgs = {im["id"]: im for im in ann["images"]}
    per_img = defaultdict(list)
    for a in ann["annotations"]:
        per_img[a["image_id"]].append(a)
    out = []
    for iid, image_anns in per_img.items():
        info = imgs.get(iid)
        if info is None:
            continue
        wh = float(info["width"] * info["height"]) or 1.0
        persons = [a for a in image_anns if a["category_id"] == _COCO_PERSON_CAT]
        if any(a.get("iscrowd") for a in persons):
            continue
        if not (min_persons <= len(persons) <= max_persons):
            continue
        if any(float(a.get("area", 0.0)) < min_person_area_frac * wh for a in persons):
            continue
        if require_nonperson and not any(
                a["category_id"] != _COCO_PERSON_CAT and not a.get("iscrowd")
                for a in image_anns):
            continue
        out.append(iid)
    out.sort()      # deterministic selection order
    return out


def _multi_person_gt(image_anns: list, cats: dict) -> dict:
    """Shape one image's annotations into the fusion GT. Keeps ALL instances and ALL
    polygon parts per annotation — occluded people are routinely split into 2+
    polygons; the first-polygon-only rule used by _coco_instances would corrupt
    person masks on exactly this slice."""
    persons, objects = [], []
    for a in image_anns:
        if a.get("iscrowd"):
            continue
        seg = a.get("segmentation")
        polys = ([[float(v) for v in part] for part in seg
                  if isinstance(part, list) and len(part) >= 6]
                 if isinstance(seg, list) else [])
        x, y, w, h = a["bbox"]
        rec = {"ann_id": a["id"], "box_xyxy": [x, y, x + w, y + h],
               "polygons": polys, "area": float(a.get("area", w * h))}
        if a["category_id"] == _COCO_PERSON_CAT:
            persons.append(rec)
        else:
            objects.append(dict(rec, label=cats.get(a["category_id"], "object")))
    return {"persons": persons, "objects": objects, "n_persons": len(persons)}


def load_coco_multi_person(n: int = 24, split=None, *, min_persons: int = 2,
                           max_persons: int = 6, min_person_area_frac: float = 0.005,
                           require_nonperson: bool = False) -> list[GTSample]:
    """Clean 2-6-person COCO slice for fusion validation. GT retains all instances +
    all polygon parts; the 5 human captions ride in meta["captions"]. Filtering runs
    over the cached annotations BEFORE any image download."""
    import io
    import urllib.request
    from collections import defaultdict
    from PIL import Image

    key = (f"multi_{n}_{min_persons}_{max_persons}_{min_person_area_frac}"
           f"_{require_nonperson}")
    if key in _COCO_CACHE:
        return _COCO_CACHE[key]
    ann = _coco_ann("instances")
    cap_ann = _coco_ann("captions")
    cats = {c["id"]: c["name"] for c in ann["categories"]}
    imgs = {im["id"]: im for im in ann["images"]}
    per_img = defaultdict(list)
    for a in ann["annotations"]:
        per_img[a["image_id"]].append(a)
    caps = defaultdict(list)
    for c in cap_ann["annotations"]:
        caps[c["image_id"]].append(str(c["caption"]).strip())

    out = []
    for iid in _select_multi_person_ids(
            ann, min_persons=min_persons, max_persons=max_persons,
            min_person_area_frac=min_person_area_frac,
            require_nonperson=require_nonperson):
        if len(out) >= n:
            break
        info = imgs[iid]
        try:
            raw = urllib.request.urlopen(
                f"http://images.cocodataset.org/val2017/{info['file_name']}",
                timeout=60).read()
            img = Image.open(io.BytesIO(raw)).convert("RGB")
        except Exception:
            continue
        gt = _multi_person_gt(per_img[iid], cats)
        out.append(GTSample(img, "Fuse the scene into entities, relations, and counts.",
                            gt, "fusion_scene", f"coco_{iid}",
                            (img.width, img.height), meta={"captions": caps.get(iid, [])}))
    _COCO_CACHE[key] = out
    return out


def load_coco_multi_person_rich(n: int = 24, split=None) -> list[GTSample]:
    """Multi-person images that ALSO contain a non-person object (relation richness)."""
    return load_coco_multi_person(n, split, require_nonperson=True)


DATASET_REGISTRY: dict[str, Callable[..., list[GTSample]]] = {
    "imagenet_val": load_imagenet_val,
    "coco_detection": load_coco_detection,
    "coco_segmentation": load_coco_segmentation,
    "coco_outline": load_coco_outline,
    "coco_subject": load_coco_subject,
    "coco_multi_person": load_coco_multi_person,
    "coco_multi_person_rich": load_coco_multi_person_rich,
    "textvqa": load_textvqa,
    "datatype_synth": make_datatype_samples,
    "shapes_synth": make_shapes_samples,
    "segmentation_synth": make_segmentation_samples,
    "outline_synth": make_outline_samples,
    "boxes3d_synth": make_3d_samples,
    "camera_rot_synth": make_camera_samples,
    "gqa": make_gqa_samples,
    "semantic_synth": make_semantic_samples,
    "style_synth": make_style_samples,
}


def load_gt(dataset_key: str, n: int = 200, split: str = "validation",
            dataset: str = "full") -> list[GTSample]:
    """Top-level GT loader. dataset='smoke' uses the packaged offline set."""
    if dataset == "smoke" or dataset_key in ("", "smoke"):
        # caller passes the category as dataset_key for smoke
        return smoke_samples(dataset_key, n)
    loader = DATASET_REGISTRY.get(dataset_key)
    if loader is None:
        raise KeyError(f"no loader for dataset {dataset_key!r}. known: {list(DATASET_REGISTRY)}")
    return loader(n=n, split=split)