"""Сборка набора по классам (woman/man/cat) из ~/dataset/allphoto_train — тест анатомии.

Зачем: бабочки однородны и анатомию не показывают. Нужен набор, где модель обязана выучить
тело: разные люди, позы, ракурсы, и капшены, которые называют, кто это и что делает.

Классы: woman / man / cat — по N уникальных капшенов на класс (дедуп по первым 60 символам,
иначе один photoshoot даёт десятки почти одинаковых кадров).

  python3 src/make_dataset_people.py --n 60 --out ~/dataset/testset_anatomy
"""
from __future__ import annotations

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

from PIL import Image, ImageOps

SRC = Path.home() / "dataset" / "allphoto_train"

CLASSES = {
    "woman": ("woman", "women", "girl"),
    "man": ("man", "men", "boy"),
    "cat": ("cat",),
}
# бакеты как в moss/src/make_dataset.py: обе стороны кратны 16 (латент /8, потом patch-2),
# соотношение выбирается ближайшим по |log(ratio_источника / ratio_бакета)|, затем ЦЕНТР-КРОП —
# растягивать нельзя, иначе лица и тела деформируются
def buckets(long_side: int = 512, multiple: int = 64) -> dict:
    """Бакеты: длинная сторона long_side, обе стороны кратны 16 (латент /8, затем /2)."""
    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),
    }
BUCKETS = buckets(512)

# для людей требуем, чтобы в кадре было тело, а не только лицо: бабочки-портреты анатомию не учат
BODY_WORDS = ("standing", "walking", "sitting", "seated", "jeans", "shorts", "dress", "skirt",
              "shirt", "legs", "arms", "hands", "wearing", "full body", "walks", "stands")


def classify(cap: str) -> str | None:
    low = cap.lower()
    for cls, words in CLASSES.items():
        if any(re.search(rf"\b{w}\b", low) for w in words):
            # «man» внутри «woman» уже отсечён \b; но «woman» не должна становиться «man»
            if cls == "man" and re.search(r"\bwoman\b", low):
                return "woman"
            return cls
    return None


def trim_black(im: Image.Image, thr: float = 12.0, var: float = 8.0) -> Image.Image:
    """Срезает однородные чёрные поля по краям (в allphoto_train их полно: 71 из 171 кадра)."""
    import numpy as np
    a = np.asarray(im.convert("L"), dtype=np.float32)
    def uniform_edge(line):        # чёрная и без структуры -> рамка, а не ночная сцена
        return line.mean() < thr and line.std() < var
    top, bot, left, right = 0, a.shape[0], 0, a.shape[1]
    while top < bot - 8 and uniform_edge(a[top]): top += 1
    while bot > top + 8 and uniform_edge(a[bot - 1]): bot -= 1
    while left < right - 8 and uniform_edge(a[top:bot, left]): left += 1
    while right > left + 8 and uniform_edge(a[top:bot, right - 1]): right -= 1
    return im.crop((left, top, right, bot)) if (left or top or right != a.shape[1] or bot != a.shape[0]) else im


def plan_size(w: int, h: int, max_upscale: float = 1.4):
    """(w,h) -> (label, tw, th, cw, ch) или None, если источника не хватает (как в moss)."""
    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 load_bucketed(path: Path, max_upscale: float = 1.4):
    """EXIF-поворот -> срез чёрных полей -> ближайший бакет -> центр-кроп -> ресайз."""
    im = ImageOps.exif_transpose(Image.open(path)).convert("RGB")
    im = trim_black(im)
    plan = plan_size(im.width, im.height, max_upscale)
    if plan is None:
        return None, None
    label, tw, th, cw, ch = plan
    left, top = (im.width - cw) // 2, (im.height - ch) // 2
    return im.crop((left, top, left + cw, top + ch)).resize((tw, th), Image.LANCZOS), label


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--n", type=int, default=60, help="картинок на класс")
    ap.add_argument("--eval-n", type=int, default=4, help="сколько на класс отложить для показа")
    ap.add_argument("--max-upscale", type=float, default=1.4)
    ap.add_argument("--long-side", type=int, default=512)
    ap.add_argument("--multiple", type=int, default=64)
    ap.add_argument("--seed", type=int, default=0)
    ap.add_argument("--out", default=str(Path.home() / "dataset" / "testset_anatomy"))
    args = ap.parse_args()

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

    caps = sorted(SRC.glob("*.txt"))
    rng = random.Random(args.seed)
    rng.shuffle(caps)

    seen_prefix: set[str] = set()
    picked: dict[str, list[tuple[Path, str]]] = {c: [] for c in CLASSES}
    need = args.n + args.eval_n
    for f in caps:
        cap = " ".join(f.read_text().split())
        if len(cap) < 40:
            continue
        cls = classify(cap)
        if cls is None or len(picked[cls]) >= need:
            continue
        low = cap.lower()
        if cls == "cat" and re.search(r"\b(dog|puppy|kitten|rabbit)\b", low) and "cat" not in low:
            continue
        if cls in ("woman", "man") and not any(w in low for w in BODY_WORDS):
            continue
        key = cap[:60].lower()
        if key in seen_prefix:
            continue
        img = f.with_suffix(".jpg")
        if not img.exists():
            continue
        seen_prefix.add(key)
        picked[cls].append((img, cap))
        if all(len(v) >= need for v in picked.values()):
            break

    rows = []
    for cls, items in picked.items():
        for i, (img, cap) in enumerate(items):
            split = "eval" if i < args.eval_n else "train"
            im, bucket = load_bucketed(img)
            if im is None:
                continue
            name = f"{cls}_{i:03d}"
            im.save(out / split / f"{name}.jpg", quality=92)
            (out / split / f"{name}.txt").write_text(cap + "\n")
            rows.append({"split": split, "class": cls, "name": name, "w": im.width, "h": im.height,
                         "bucket": bucket, "caption": cap})
            print(f"  {split}/{name} {im.width}x{im.height} | {cap[:80]}")

    with open(out / "index.csv", "w", newline="") as fh:
        w = csv.DictWriter(fh, fieldnames=["split", "class", "name", "w", "h", "bucket", "caption"])
        w.writeheader(); w.writerows(rows)
    per = {c: sum(1 for r in rows if r["class"] == c and r["split"] == "train") for c in CLASSES}
    print(f"\nитого: {sum(per.values())} train {per} + {sum(1 for r in rows if r['split']=='eval')} eval"
          f" -> {out}")


if __name__ == "__main__":
    main()