"""Сборка набора из папки «картинка + .txt»: бакеты (кратны --multiple) и центр-кроп.

  python3 src/make_dataset.py --src ~/dataset/butterflies/butterfly \
          --out ~/dataset/butterflies_256 --long-side 256 --eval-n 30
"""
from __future__ import annotations

import argparse
import csv
import math
import random
from pathlib import Path

from PIL import Image, ImageOps

BUCKETS = {}


def buckets(long_side: int, multiple: int = 16) -> dict:
    L = long_side
    def r16(x): return max(multiple, int(round(x / multiple)) * multiple)
    return {
        "1:1": (L, L),
        "4:3": (L, r16(L * 3 / 4)),
        "3:4": (r16(L * 3 / 4), L),
        "3:2": (L, r16(L * 2 / 3)),
        "2:3": (r16(L * 2 / 3), L),
        "16:9": (L, r16(L * 9 / 16)),
        "9:16": (r16(L * 9 / 16), L),
    }


def plan_size(w, h, max_upscale=1.4):
    ratio = w / h
    label, (tw, th) = min(BUCKETS.items(), key=lambda kv: abs(math.log(ratio / (kv[1][0] / kv[1][1]))))
    if ratio > tw / th:
        cw, ch = round(h * tw / th), h
    else:
        cw, ch = w, round(w * th / tw)
    cw, ch = min(cw, w), min(ch, h)
    if cw * max_upscale < tw or ch * max_upscale < th:
        return None
    return label, tw, th, cw, ch


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--src", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--long-side", type=int, default=256)
    ap.add_argument("--eval-n", type=int, default=30)
    ap.add_argument("--max-upscale", type=float, default=1.4)
    ap.add_argument("--multiple", type=int, default=64,
                    help="кратность сторон: три даунсэмпла тела требуют латент, кратный 8")
    ap.add_argument("--seed", type=int, default=0)
    args = ap.parse_args()

    global BUCKETS
    BUCKETS = buckets(args.long_side, args.multiple)
    src, out = Path(args.src), Path(args.out)
    for sub in ("train", "eval"):
        (out / sub).mkdir(parents=True, exist_ok=True)

    pairs = []
    for img in sorted(list(src.glob("*.png")) + list(src.glob("*.jpg"))):
        txt = img.with_suffix(".txt")
        if txt.exists():
            pairs.append((img, txt))
    random.Random(args.seed).shuffle(pairs)
    print(f"пар: {len(pairs)}")

    rows, kept = [], 0
    for i, (img, txt) in enumerate(pairs):
        im = ImageOps.exif_transpose(Image.open(img)).convert("RGB")
        plan = plan_size(im.width, im.height, args.max_upscale)
        if plan is None:
            continue
        label, tw, th, cw, ch = plan
        left, top = (im.width - cw) // 2, (im.height - ch) // 2
        im = im.crop((left, top, left + cw, top + ch)).resize((tw, th), Image.LANCZOS)
        split = "eval" if i < args.eval_n else "train"
        name = f"{img.stem}"
        im.save(out / split / f"{name}.jpg", quality=95)
        (out / split / f"{name}.txt").write_text(txt.read_text().strip() + "\n")
        rows.append({"split": split, "name": name, "w": tw, "h": th, "bucket": label})
        kept += 1
    with open(out / "index.csv", "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=["split", "name", "w", "h", "bucket"])
        w.writeheader(); w.writerows(rows)
    print(f"готово: {kept} (train {kept-args.eval_n}, eval {args.eval_n}) -> {out}")


if __name__ == "__main__":
    main()