#!/usr/bin/env python3
"""Convert WIDER FACE (HF parquet, CUHK-CSE/wider_face @ refs/convert/parquet) to single-class YOLO format.

- Skips boxes flagged `invalid` and degenerate boxes (<1px).
- Keeps small faces (no min-size filter) since blurring needs high recall.
- Drops images whose only boxes are invalid (they would teach the model to ignore real faces).
- Downscales images so the long side is <= --max-side (matches training imgsz; boxes are normalized).

Usage:
  python convert_wider_to_yolo.py --raw raw/default --out dataset
"""
import argparse
import io
import os
from concurrent.futures import ProcessPoolExecutor
from pathlib import Path

import pyarrow.parquet as pq
from PIL import Image

SPLITS = {"train": "train", "validation": "val"}


def convert_row(args):
    name, img_bytes, bboxes, invalid, out_img, out_lbl, max_side = args
    valid = [b for b, inv in zip(bboxes, invalid) if not inv and b[2] >= 1 and b[3] >= 1]
    if not valid:
        return "dropped", 0, 0
    im = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    w, h = im.size
    scale = min(1.0, max_side / max(w, h))
    if scale < 1.0:
        im = im.resize((max(1, round(w * scale)), max(1, round(h * scale))), Image.LANCZOS)
    stem = Path(name).stem
    im.save(out_img / f"{stem}.jpg", quality=95)
    lines = []
    for x, y, bw, bh in valid:
        x1, y1 = max(0.0, x), max(0.0, y)
        x2, y2 = min(float(w), x + bw), min(float(h), y + bh)
        if x2 - x1 < 1 or y2 - y1 < 1:
            continue
        lines.append(f"0 {(x1 + x2) / 2 / w:.6f} {(y1 + y2) / 2 / h:.6f} {(x2 - x1) / w:.6f} {(y2 - y1) / h:.6f}")
    (out_lbl / f"{stem}.txt").write_text("\n".join(lines) + "\n")
    return "ok", len(lines), len(bboxes) - len(valid)


def rows(files, out_img, out_lbl, max_side):
    for f in files:
        pf = pq.ParquetFile(f)
        for batch in pf.iter_batches(batch_size=64):
            for r in batch.to_pylist():
                yield (r["image"]["path"], r["image"]["bytes"], r["faces"]["bbox"],
                       r["faces"]["invalid"], out_img, out_lbl, max_side)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--raw", default="raw/default")
    ap.add_argument("--out", default="dataset")
    ap.add_argument("--max-side", type=int, default=640)
    ap.add_argument("--limit", type=int, default=0, help="max images per split (0 = all; for testing)")
    ap.add_argument("--workers", type=int, default=os.cpu_count())
    a = ap.parse_args()

    out = Path(a.out)
    for src, dst in SPLITS.items():
        files = sorted((Path(a.raw) / src).glob("*.parquet"))
        out_img, out_lbl = out / "images" / dst, out / "labels" / dst
        out_img.mkdir(parents=True, exist_ok=True)
        out_lbl.mkdir(parents=True, exist_ok=True)
        gen = rows(files, out_img, out_lbl, a.max_side)
        if a.limit:
            import itertools
            gen = itertools.islice(gen, a.limit)
        ok = dropped = faces = 0
        with ProcessPoolExecutor(a.workers) as ex:
            for status, n, _ in ex.map(convert_row, gen, chunksize=16):
                if status == "ok":
                    ok += 1
                    faces += n
                else:
                    dropped += 1
        print(f"{dst}: {ok} images, {faces} valid faces, {dropped} images dropped (no valid boxes)")

    (out / "data.yaml").write_text(
        f"path: {out.resolve()}\ntrain: images/train\nval: images/val\nnames:\n  0: face\n")
    print("wrote", out / "data.yaml")


if __name__ == "__main__":
    main()