"""RefCOCO (jxu124/refcoco): referring expressions with masks and boxes on COCO train2014 photos.

The images are not in the dataset: point COCO_DIR at COCO train2014 (or at the subset that
recipes/fetch_refcoco.sh downloads: only the ~20k images RefCOCO uses).
"""

import json, os, random

COCO = os.environ.get("COCO_DIR", "/workspace/data/coco_ref")


def rows(split):
    from datasets import load_dataset
    return load_dataset("jxu124/refcoco", split=split)


def examples(split_rows, n, seed, one_sentence=True):
    """Refs as dict(image, size, polys, bbox (x1, y1, x2, y2 pixels), ref_id, sentence). Training: one
    random sentence per ref. Evaluation with one_sentence=False: every expression, as the RefCOCO
    literature scores it."""
    idx = list(range(len(split_rows)))
    random.Random(seed).shuffle(idx)
    out = []
    for i in idx:
        r = split_rows[i]
        ann, info = json.loads(r["raw_anns"]), json.loads(r["raw_image_info"])
        sents = [s["sent"] for s in r["sentences"]]
        base = dict(image=os.path.join(COCO, info["file_name"]), size=(info["width"], info["height"]),
                    polys=ann["segmentation"], bbox=list(r["bbox"]), ref_id=r["ref_id"])
        if one_sentence:
            out.append(dict(base, sentence=random.Random(seed + i).choice(sents)))
        else:
            out += [dict(base, sentence=s) for s in sents]
        if n and len(out) >= n:
            break
    return out[:n] if n else out


def photo_boxes(n, split, seed):
    """Referents as box examples (one expression per object): general objects in photos."""
    out = []
    for ex in examples(rows(split), n, seed, one_sentence=True):
        (W, H), (x1, y1, x2, y2) = ex["size"], ex["bbox"]
        out.append(dict(image=ex["image"], size=ex["size"], instruction=ex["sentence"],
                        box=[x1 / W, y1 / H, x2 / W, y2 / H], norm=True, source="refcoco"))
    return out