"""
fetch_coco_subset.py - build train and eval sets from sayakpaul/coco-30-val-2014.

The stream of 30K (image, caption) rows is split by row index:
  rows 0..EVAL_N-1        -> eval set (real images for FID + captions for generation)
  rows TRAIN_START..end   -> training candidates

Training images whose content hash also appears in the eval set are dropped,
so the model is never trained on an image it is evaluated against (the same
image can appear in multiple rows with different captions). Training rows are
also deduped against each other by image hash.

Outputs (in --out dir):
  coco_train.npz               images uint8 (N, 64, 64, 3), center-cropped
  coco_train.captions.json     N caption strings
  eval_real/real_00000.jpg...  EVAL_N real images, 256x256 center-crop, q95
  eval_captions.json           EVAL_N caption strings (for generation + CLIP)

Usage:
  python fetch_coco_subset.py --out ../pm-work
  # much faster: pre-download the 10 parquet shards, then
  python fetch_coco_subset.py --out ../pm-work --parquet-dir ../pm-work/shards
"""

import argparse
import glob
import hashlib
import json
import os

import numpy as np
from datasets import load_dataset
from PIL import Image

EVAL_N      = 5000
TRAIN_START = 10000     # comfortable gap after the eval rows
TRAIN_RES   = 64
EVAL_RES    = 256


def center_crop_resize(img: Image.Image, size: int) -> Image.Image:
    w, h = img.size
    side = min(w, h)
    left, top = (w - side) // 2, (h - side) // 2
    return img.crop((left, top, left + side, top + side)).resize(
        (size, size), Image.BICUBIC
    )


def img_hash(img: Image.Image) -> str:
    """Content hash of the image, robust to caption-duplicated rows."""
    return hashlib.md5(
        img.convert("RGB").resize((64, 64), Image.BILINEAR).tobytes()
    ).hexdigest()


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--out", required=True)
    p.add_argument("--parquet-dir", default=None,
                   help="dir with the dataset's parquet shards already downloaded "
                        "(sorted filename order == hub streaming row order)")
    args = p.parse_args()

    real_dir = os.path.join(args.out, "eval_real")
    os.makedirs(real_dir, exist_ok=True)

    if args.parquet_dir:
        files = sorted(glob.glob(os.path.join(args.parquet_dir, "*.parquet")))
        assert files, f"no parquet files in {args.parquet_dir}"
        ds = load_dataset("parquet", data_files=files, split="train", streaming=True)
    else:
        ds = load_dataset("sayakpaul/coco-30-val-2014", split="train", streaming=True)

    eval_hashes = set()
    eval_captions = []
    train_images, train_captions = [], []
    train_hashes = set()
    skipped_eval_overlap = skipped_dupe = 0

    for i, row in enumerate(ds):
        img, caption = row["image"], row["caption"]

        if i < EVAL_N:
            eval_hashes.add(img_hash(img))
            eval_captions.append(caption)
            center_crop_resize(img.convert("RGB"), EVAL_RES).save(
                os.path.join(real_dir, f"real_{i:05d}.jpg"), quality=95
            )
            if (i + 1) % 1000 == 0:
                print(f"  eval: {i + 1}/{EVAL_N}", flush=True)
            continue

        if i < TRAIN_START:
            continue

        h = img_hash(img)
        if h in eval_hashes:
            skipped_eval_overlap += 1
            continue
        if h in train_hashes:
            skipped_dupe += 1
            continue
        train_hashes.add(h)
        thumb = center_crop_resize(img.convert("RGB"), TRAIN_RES)
        train_images.append(np.array(thumb, dtype=np.uint8))
        train_captions.append(caption)
        if len(train_images) % 2000 == 0:
            print(f"  train: {len(train_images)} (row {i})", flush=True)

    images = np.stack(train_images)
    np.savez_compressed(os.path.join(args.out, "coco_train.npz"), images=images)
    with open(os.path.join(args.out, "coco_train.captions.json"), "w", encoding="utf-8") as f:
        json.dump(train_captions, f)
    with open(os.path.join(args.out, "eval_captions.json"), "w", encoding="utf-8") as f:
        json.dump(eval_captions, f)

    print(f"\neval:  {len(eval_captions)} pairs @ {EVAL_RES}x{EVAL_RES} -> {real_dir}")
    print(f"train: {len(train_captions)} pairs @ {TRAIN_RES}x{TRAIN_RES} "
          f"(skipped: {skipped_eval_overlap} eval-overlap, {skipped_dupe} dupes)")


if __name__ == "__main__":
    main()