"""Evaluation sets outside the training data.

  screenspot   ScreenSpot (bevaya/ScreenSpot, test, 1,272), shuffled with seed 0: clicks and element boxes
  coco17       COCO val2017 photos (RefCOCO is all train2014): objects prompted by category name, for
               categories with exactly one instance in the image, at least 1% of it, polygon masks only
  fashion      Fashionpedia val: one garment per image, by name (a new domain; boxes only)

Point COCO17_DIR at a folder with annotations/instances_val2017.json and val2017/
(recipes/fetch_coco17.sh).
"""

import json, os, random

COCO17 = os.environ.get("COCO17_DIR", "/workspace/data/coco17")
MIN_AREA = 0.01
GARMENTS = 27  # Fashionpedia classes 0-26 are whole garments and accessories; the rest are parts


def screenspot(n=0):
    from datasets import load_dataset
    ds = load_dataset("bevaya/ScreenSpot", split="test").shuffle(seed=0)
    ds = ds.select(range(n)) if n else ds
    return [dict(image=r["image"], instruction=r["instruction"], box=list(r["bbox"]), norm=True,
                 group=(r["data_source"], r["data_type"])) for r in ds]


def coco17(n, seed=0):
    d = json.load(open(f"{COCO17}/annotations/instances_val2017.json"))
    cats = {c["id"]: c["name"] for c in d["categories"]}
    imgs = {i["id"]: i for i in d["images"]}
    by = {}
    for a in d["annotations"]:
        by.setdefault((a["image_id"], a["category_id"]), []).append(a)
    out = []
    for (iid, cid), anns in sorted(by.items()):
        if len(anns) != 1:
            continue
        a, im = anns[0], imgs[iid]
        W, H = im["width"], im["height"]
        if a["iscrowd"] or not isinstance(a["segmentation"], list) or a["area"] < MIN_AREA * W * H:
            continue
        x, y, w, h = a["bbox"]
        out.append(dict(image=f"{COCO17}/val2017/{im['file_name']}", size=(W, H), sentence=cats[cid],
                        instruction=cats[cid], polys=a["segmentation"], bbox=[x, y, x + w, y + h],
                        box=[x / W, y / H, (x + w) / W, (y + h) / H], norm=True, ref_id=a["id"], source="coco17"))
    random.Random(seed).shuffle(out)
    return out[:n] if n else out


def fashion(n, seed=0):
    from datasets import load_dataset
    ds = load_dataset("detection-datasets/fashionpedia", split="val")
    names = ds.features["objects"]["category"].feature.names
    idx = list(range(len(ds)))
    random.Random(seed).shuffle(idx)
    out = []
    for i in idx:
        r = ds[i]
        W, H, o = r["width"], r["height"], r["objects"]
        counts = {}
        for c in o["category"]:
            counts[c] = counts.get(c, 0) + 1
        for c, (x1, y1, x2, y2) in zip(o["category"], o["bbox"]):  # x1, y1, x2, y2 in pixels
            if c < GARMENTS and counts[c] == 1 and (x2 - x1) * (y2 - y1) >= MIN_AREA * W * H:
                out.append(dict(image=r["image"], size=(W, H), instruction=f"the {names[c].split(',')[0]}",
                                box=[x1 / W, y1 / H, x2 / W, y2 / H], norm=True, source="fashionpedia"))
                break  # one object per image
        if n and len(out) >= n:
            break
    return out