#!/usr/bin/env python3
"""Микро-датасет для обучения MOSS-Image: латенты VAE + условие адаптера + маска.

Формат строки (HF arrow dataset, `save_to_disk`):
    vae          float16 [128, H/16, W/16]  — латент, уже упакованный и bn-нормированный,
                                              ровно то, что ждёт MOSSImageTransformer
    text_embeds  float16 [L, 4096]          — выход адаптерного ТЕ, без паддинга (как в пайплайне)
                                              (только с --te)
    text_mask    bool    [L]                — маска условия (для нашей нарезки вся True;
                                              нужна, когда батч склеивают паддингом)  (только с --te)
    text         str                        — исходный капшен (для глаз и отладки)
    width/height int                        — размер картинки, из которой сделан латент
    aspect       str                        — '1:1' / '3:4' / '4:3'
    source       str                        — путь исходного файла

С `--te` рядом с датасетом кладётся `empty_condition.npz` (embeds float16 [L,4096] + mask bool [L]) —
условие пустого промпта, то есть негатив CFG. Он считается тем же адаптером один раз здесь,
чтобы обучению (`src/train.py`) текстовый энкодер был не нужен вообще. Без `--te` текстовый
энкодер не загружается вообще: в датасете только латенты и мета (текст, размеры, источник).

Что делает с картинкой:
  * EXIF-поворот применяется до всего остального;
  * выбирается ближайшее по |log(w/h / target)| из набора соотношений (см. ASPECTS ниже);
  * центр-кроп до выбранного соотношения, затем LANCZOS-ресайз до целевого размера. Апскейл
    разрешён до --max-upscale (1.3 по умолчанию): мельче 1/1.3 от цели — отбрасывается.

Латент: vae.encode(x).latent_dist.mode() -> [32,H/8,W/8]; упаковка 2x2 -> [128,H/16,W/16];
нормировка (z - bn.running_mean) / sqrt(bn.running_var + eps). Обратная операция — в
`MOSSImagePipeline.__call__`, round-trip проверен (MSE 3e-4).

Текст (только с --te): тот же путь, что `MOSSImagePipeline.encode_prompt` — chat-шаблон с
add_generation_prompt=True/enable_thinking=True, обрезка по --text-len, отдельный forward на
промпт без паддинга (паддинг меняет цель на ~0.5 %, см. adapter/src/cache_targets.py).

Инфраструктура — как в sdxs-micro/dataset.py: accelerate делит файлы между GPU, каждый процесс
пишет чанки-поддатасеты <save>/chunk_*/sub_* с маркером .done (перезапуск продолжает с места).
Склейки в один датасет нет: src/train.py сам находит все чанки в папке и конкатенирует их при
чтении (лениво, arrow-файлы не переписываются).

Запуск:
    python3 src/make_dataset.py --dry-run --limit 20          # только отчёт по размерам, без GPU
    python3 src/make_dataset.py --te --limit 20               # 20 картинок + условие ТЕ
    python3 src/make_dataset.py --folder ~/dataset/all          # только латенты (ТЕ не грузится)
    accelerate launch --multi_gpu --num_processes 2 src/make_dataset.py --te --folder ~/dataset/allphoto_train

Как читать в обучении:
    row = ds[i]                       # np.array() даёт float64: приводим явно
    lat = torch.from_numpy(np.asarray(row["vae"], dtype=np.float16))[None]          # [1,128,H/16,W/16]
    cond = torch.from_numpy(np.asarray(row["text_embeds"], dtype=np.float16))[None]  # [1,L,4096]
    mask = torch.from_numpy(np.asarray(row["text_mask"]))[None]                      # [1,L] bool
"""
import argparse
import gc
import math
import os
import random
import shutil
import sys
import time

import numpy as np
import torch
from PIL import Image, ImageOps
from accelerate import Accelerator
from datasets import Dataset, load_from_disk

MODEL = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))  # папка модели
if MODEL not in sys.path:
    sys.path.insert(0, MODEL)  # diffusers ищет pipeline_moss_image как модуль

accelerator = Accelerator()
device = accelerator.device
process_index = accelerator.process_index
num_processes = accelerator.num_processes
is_main = accelerator.is_main_process

# Соотношения, у которых обе стороны кратны 16 (H/16 и W/16 должны быть целыми). Набор держит
# длинную сторону 1152, кроме 16:9/9:16: у 16:9 высота делится на 16 только при ширине, кратной 256,
# поэтому там 1280x720 вместо 1152x648.
ASPECTS = {
    "1:1": (1152, 1152),   # 1.33 Mpx — самый тяжёлый бакет
    "4:3": (1152, 864),
    "3:4": (864, 1152),
    "3:2": (1152, 768),
    "2:3": (768, 1152),
    "16:9": (1280, 720),
    "9:16": (720, 1280),
}
DEFAULT_FOLDER = os.path.expanduser("~/dataset")
DEFAULT_SAVE = os.path.join(MODEL, "datasets", "dataset")
# fp16 проверен на этом VAE: round-trip MSE тот же, что в fp32 (8.99e-04 на трёх фото),
# нечисел нет, зато вдвое быстрее (0.7 против 1.3 с на 3 картинки) и меньше VRAM (1.5 против 2.5 ГиБ)
VAE_DTYPE = torch.float16
# pyarrow держит офсеты списков в int32 — не больше 1.5 млрд элементов на под-датасет
_MAX_ELEMS_PER_DATASET = int(1_500_000_000)


def parse_args():
    p = argparse.ArgumentParser(description="MOSS-Image: картинки -> латенты + условие + маска")
    p.add_argument("--folder", default=DEFAULT_FOLDER, help="папка с *.jpg/*.png и парными *.txt")
    p.add_argument("--save", default=DEFAULT_SAVE, help="куда сохранить датасет")
    p.add_argument("--limit", type=int, default=0,
                   help="взять только N случайных пар (сид --seed); 0 — все")
    p.add_argument("--seed", type=int, default=0, help="сид случайного отбора для --limit")
    p.add_argument("--batch-size", type=int, default=4, help="картинок на один forward VAE")
    p.add_argument("--chunk-size", type=int, default=10000, help="файлов на чанк (на процесс)")
    p.add_argument("--text-len", type=int, default=1024, help="макс. токенов условия (как в пайплайне)")
    p.add_argument("--max-upscale", type=float, default=1.3,
                   help="во сколько раз разрешено увеличивать картинку до целевой; 1.0 — только "
                        "даунскейл (всё мельче целевого размера отбрасывается)")
    p.add_argument("--strip-words", default="*tags*",
                   help="через запятую: подстроки, вырезаемые из капшена (по умолчанию '*tags*', "
                        "то есть и '*tags* , 1girl', и '*tags* 1girl'); "
                        "пустая строка — ничего не вырезать. После вырезки капшен сжимается по "
                        "пробелам и strip-ится по ' ,;|'; пустой — строка отбрасывается")
    p.add_argument("--dry-run", action="store_true",
                   help="только отчёт: что выбрано, сколько отброшено; без GPU и без записи")
    p.add_argument("--te", action="store_true",
                   help="считать условие адаптерным ТЕ (text_embeds/text_mask + empty_condition.npz); "
                        "без флага текстовый энкодер не грузится, в датасете только латенты и мета")
    return p.parse_args()


# ---------------------------------------------------------------- VAE и текст
def load_vae():
    from diffusers import AutoencoderKLFlux2
    vae = AutoencoderKLFlux2.from_pretrained(os.path.join(MODEL, "vae"), torch_dtype=VAE_DTYPE)
    return vae.to(device).eval()


def load_text_encoder():
    from transformers import AutoTokenizer
    from pipeline_moss_image import MOSSAdapterTextEncoder
    path = os.path.join(MODEL, "text_encoder")
    tok = AutoTokenizer.from_pretrained(path)
    # from_pretrained адаптера сам поднимает голову из ../adapter и ученика из --path
    te = MOSSAdapterTextEncoder.from_pretrained(path).to(device).eval()
    return te, tok


def render(tok, text):
    """Ровно то, что делает MOSSImagePipeline.encode_prompt (сырая строка, не id)."""
    return tok.apply_chat_template([{"role": "user", "content": text}], tokenize=False,
                                   add_generation_prompt=True, enable_thinking=True)


@torch.no_grad()
def encode_text(te, tok, text, text_len):
    """Промпт -> (embeds float16 [L,4096], mask bool [L]). Без паддинга, как в пайплайне."""
    ids = tok(render(tok, text), truncation=True, max_length=text_len)["input_ids"]
    if not ids:
        raise ValueError("пустой промпт после токенизации")
    ids = torch.tensor([ids], dtype=torch.long, device=device)
    out = te(input_ids=ids, attention_mask=torch.ones_like(ids, dtype=torch.bool), use_cache=False)
    emb = out.last_hidden_state[0].to(torch.float16)
    if not torch.isfinite(emb).all():
        raise RuntimeError(f"условие не влезло в float16 (промпт: {text[:60]!r})")
    mask = np.ones(emb.shape[0], dtype=bool)
    return emb.cpu().numpy(), mask


@torch.no_grad()
def encode_latents(vae, images):
    """list[PIL] (одинакового размера) -> np.float16 [B,128,H/16,W/16]."""
    x = torch.stack([torch.from_numpy(np.array(im)).permute(2, 0, 1).float().div(127.5).sub(1)
                     for im in images]).to(device, VAE_DTYPE)
    lat = vae.encode(x).latent_dist.mode()
    b, c, h, w = lat.shape
    if h % 2 or w % 2:
        raise ValueError(f"латент {h}x{w} не делится на 2 — картинка не кратна 16?")
    # упаковка 2x2: [B,c,h,w] -> [B,4c,h/2,w/2]; обратная операция в MOSSImagePipeline.__call__
    z = lat.reshape(b, c, h // 2, 2, w // 2, 2).permute(0, 1, 3, 5, 2, 4).reshape(
        b, c * 4, h // 2, w // 2)
    bn = vae.bn
    mean = bn.running_mean.to(z.dtype).view(1, -1, 1, 1)
    std = (bn.running_var + vae.config.batch_norm_eps).sqrt().to(z.dtype).view(1, -1, 1, 1)
    return ((z - mean) / std).half().cpu().numpy()


# ---------------------------------------------------------------- выбор размера
def plan_size(width, height, max_upscale=1.0):
    """(w,h) исходника -> (label, tw, th, cw, ch) кропа или None, если картинка слишком мала."""
    ratio = width / height
    label, (tw, th) = min(ASPECTS.items(),
                          key=lambda kv: abs(math.log(ratio / (kv[1][0] / kv[1][1]))))
    if ratio > tw / th:                     # шире цели -> режем ширину
        cw, ch = round(height * tw / th), height
    else:                                   # выше цели -> режем высоту
        cw, ch = width, round(width * th / tw)
    cw, ch = min(cw, width), min(ch, height)
    if cw * max_upscale < tw or ch * max_upscale < th:  # мельче цели -> отбрасываем
        return None
    return label, tw, th, cw, ch


def header_size(path):
    """(w,h) с учётом EXIF-поворота, без декодирования пикселей."""
    with Image.open(path) as im:
        w, h = im.size
        if im.getexif().get(0x0112, 1) in (5, 6, 7, 8):
            w, h = h, w
    return w, h


def load_image(path, plan):
    label, tw, th, cw, ch = plan
    im = ImageOps.exif_transpose(Image.open(path)).convert("RGB")
    left, top = (im.width - cw) // 2, (im.height - ch) // 2
    return im.crop((left, top, left + cw, top + ch)).resize((tw, th), Image.LANCZOS)


# ---------------------------------------------------------------- обход папки
def process_folder(folder, max_upscale=1.0):
    """Пары картинка+капшен, детерминированный порядок (для resume и шардинга)."""
    rows, skipped = [], {"unpaired": 0, "broken": 0, "small": 0}
    seen = 0
    for root, dirs, files in os.walk(folder):
        dirs.sort()
        for fn in sorted(files):
            if not fn.lower().endswith((".jpg", ".jpeg", ".png", ".webp")):
                continue
            seen += 1  # обход 20k+ файлов занимает десятки секунд — печатаем, что живы
            if seen % 2000 == 0:
                print(f"process_folder: просмотрено {seen} картинок, пар {len(rows)}", flush=True)
            path = os.path.join(root, fn)
            txt = os.path.splitext(path)[0] + ".txt"
            if not os.path.exists(txt):
                skipped["unpaired"] += 1
                continue
            try:
                w, h = header_size(path)
            except Exception:
                skipped["broken"] += 1
                continue
            plan = plan_size(w, h, max_upscale)
            if plan is None:
                skipped["small"] += 1
                continue
            rows.append((path, txt, plan))
    print(f"process_folder: {len(rows)} пар, отброшено {skipped['unpaired']} без капшена + "
          f"{skipped['broken']} битых + {skipped['small']} мелких", flush=True)
    return rows


def pick_limit(rows, limit, seed):
    if limit <= 0 or limit >= len(rows):
        return rows
    chosen = sorted(random.Random(seed).sample(range(len(rows)), limit))
    return [rows[i] for i in chosen]


def read_caption(path, strip_words=()):
    try:
        text = open(path, encoding="utf-8", errors="replace").read()
    except Exception:
        return None
    for word in strip_words:
        text = text.replace(word, " ")
    text = " ".join(text.split())  # пробелы/переносы по краям и двойные после вырезки
    text = text.strip(" ,;|\t")     # висячие разделители от вырезанного токена
    return text or None


def clear_cuda_memory():
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    gc.collect()


def dry_run(rows):
    print(f"\n{'aspect':6} {'crop':>12} {'target':>11}  файл")
    for path, _, (label, tw, th, cw, ch) in rows[:20]:
        print(f"{label:6} {cw:5}x{ch:<6} {tw:5}x{th:<5}  {os.path.basename(path)}")
    hist = {}
    for _, _, (label, *_rest) in rows:
        hist[label] = hist.get(label, 0) + 1
    print("\nвсего", len(rows), "по соотношениям:", dict(sorted(hist.items())),
          f"(из них показано {min(20, len(rows))})")


# ---------------------------------------------------------------- кодирование
def process_group(vae, te, tok, args, rows, group_dir, lat_shape, strip_words):
    """Одна size-группа (в батч можно класть только одинаковые кадры) -> под-датасеты.

    С `--te` в строку добавляются text_embeds/text_mask; без него эти ключи не создаются вовсе.
    """
    keys = ["vae", "text", "width", "height", "aspect", "source"]
    if args.te:
        keys += ["text_embeds", "text_mask"]
    marker = os.path.join(group_dir, ".done")
    if os.path.exists(marker):
        print(f"[GPU {process_index}] пропуск готовой группы {os.path.basename(group_dir)}",
              flush=True)
        return 0
    if os.path.isdir(group_dir):
        shutil.rmtree(group_dir)  # недописанная группа с прошлого запуска -> пересчитать
    os.makedirs(group_dir, exist_ok=True)

    # элементов на строку: латент + условие (L*4096, L оцениваем в 512 токенов); держим
    # под-датасет ниже лимита pyarrow
    elems = lat_shape[0] * lat_shape[1] * lat_shape[2]
    if args.te:
        elems += 512 * 4096
    save_batch = max(64, _MAX_ELEMS_PER_DATASET // max(elems, 1))

    acc = {k: [] for k in keys}
    saved = sub_n = skipped = 0

    def flush():
        nonlocal saved, sub_n
        if not acc["vae"]:
            return
        ds = Dataset.from_dict(dict(acc))
        path = os.path.join(group_dir, f"sub_{sub_n:05d}")
        ds.save_to_disk(path)
        print(f"[GPU {process_index}] сохранено {len(acc['vae'])} -> {path}", flush=True)
        saved += len(acc["vae"])
        sub_n += 1
        for v in acc.values():
            v.clear()

    for start in range(0, len(rows), args.batch_size):
        batch = []  # (image, text, plan, path) — только целиком прочитанные сэмплы
        for path, txt_path, plan in rows[start:start + args.batch_size]:
            try:
                img = load_image(path, plan)
            except Exception as e:
                print(f"[GPU {process_index}] битая картинка {path}: {e}", flush=True)
                skipped += 1
                continue
            text = read_caption(txt_path, strip_words)
            if text is None:
                print(f"[GPU {process_index}] пустой/нечитаемый капшен {txt_path}", flush=True)
                skipped += 1
                continue
            batch.append((img, text, plan, path))
        if not batch:
            continue

        lat = encode_latents(vae, [im for im, _, _, _ in batch])
        embeds = None
        if args.te:  # паддинга нет, L у промптов разная -> только фиксированный батч
            embeds = [encode_text(te, tok, text, args.text_len)
                      for _, text, _, _ in batch]
        for i, (row_lat, (_, text, (label, tw, th, _, _), path)) in enumerate(zip(lat, batch)):
            acc["vae"].append(row_lat)
            if embeds is not None:
                emb, mask = embeds[i]
                acc["text_embeds"].append(emb)
                acc["text_mask"].append(mask)
            acc["text"].append(text)
            acc["width"].append(tw)
            acc["height"].append(th)
            acc["aspect"].append(label)
            acc["source"].append(path)
        del lat, batch, embeds
        if len(acc["vae"]) >= save_batch:
            flush()

    flush()
    open(marker, "w").close()
    print(f"[GPU {process_index}] группа {group_dir.rsplit('_', 1)[-1]}: "
          f"{saved} сэмплов (skip {skipped})", flush=True)
    return saved


def process_in_chunks(vae, te, tok, args, rows, root, strip_words):
    total = len(rows)
    t0 = time.time()
    if is_main and args.te:
        # негатив CFG один раз: дальше обучению текстовый энкодер не нужен
        emb, mask = encode_text(te, tok, "", args.text_len)
        np.savez(os.path.join(root, "empty_condition.npz"), embeds=emb, mask=mask)
        print(f"пустой промпт: {emb.shape} -> {root}/empty_condition.npz", flush=True)
    saved = 0
    for chunk_idx, start in enumerate(range(0, total, args.chunk_size), 1):
        end = min(start + args.chunk_size, total)
        groups = {}
        for path, txt_path, plan in rows[start:end]:
            key = (plan[1], plan[2])
            groups.setdefault(key, []).append((path, txt_path, plan))
        for (tw, th), group_rows in sorted(groups.items()):
            lat_shape = (128, th // 16, tw // 16)
            group_dir = os.path.join(
                root, f"chunk_{chunk_idx:05d}_{tw}x{th}_proc_{process_index}")
            saved += process_group(vae, te, tok, args, group_rows, group_dir, lat_shape,
                                   strip_words)
        clear_cuda_memory()
        print(f"[GPU {process_index}] чанк {chunk_idx} готов ({end - start} файлов, "
              f"{time.time() - t0:.0f} c всего)", flush=True)
    return saved


# ---------------------------------------------------------------- сборка
def count_chunks(root):
    """Сколько под-датасетов и строк уже на месте (то, что прочтёт src/train.py)."""
    n_ds = n_rows = 0
    for name in sorted(os.listdir(root)):
        group_dir = os.path.join(root, name)
        if not (os.path.isdir(group_dir) and name.startswith("chunk_")):
            continue
        if not os.path.exists(os.path.join(group_dir, ".done")):
            print(f"[main] ВНИМАНИЕ: {name} без маркера .done — обучение увидит его "
                  f"неполным; запустите скрипт ещё раз, чтобы досчитать", flush=True)
        for sub in sorted(os.listdir(group_dir)):
            if sub.startswith("sub_"):
                n_ds += 1
                n_rows += len(load_from_disk(os.path.join(group_dir, sub)))
    return n_ds, n_rows


def main():
    args = parse_args()
    strip_words = [w for w in (x.strip() for x in args.strip_words.split(",")) if w]
    rows = pick_limit(process_folder(args.folder, args.max_upscale), args.limit, args.seed)
    by_aspect = {}
    for _, _, (label, *_rest) in rows:
        by_aspect[label] = by_aspect.get(label, 0) + 1
    if is_main:
        print(f"к обработке: {len(rows)}, по соотношениям: {dict(sorted(by_aspect.items()))}",
              flush=True)
    if not rows:
        raise SystemExit("нечего кодировать: ни одной подходящей пары картинка+капшен")
    if args.dry_run:
        if is_main:
            dry_run(rows)
        return

    # без CUDA accelerate молча уходит в CPU и каждый процесс берёт все файлы — падаем сразу
    if not torch.cuda.is_available():
        raise SystemExit("CUDA недоступна (torch.cuda.is_available() == False)")
    world_size = int(os.environ.get("WORLD_SIZE", num_processes))
    if world_size != num_processes:
        raise SystemExit(f"accelerate не распределил процессы: WORLD_SIZE={world_size}, "
                         f"num_processes={num_processes}")

    if is_main:
        print(f"processes: {num_processes}, folder: {args.folder}, save: {args.save}", flush=True)

    # стабильный порядок -> детерминированный шардинг и resume
    rows.sort(key=lambda r: (r[2][1], r[2][2], r[0]))
    mine = list(range(process_index, len(rows), num_processes))
    rows = [rows[i] for i in mine]
    print(f"[GPU {process_index}] обрабатываю {len(rows)} файлов", flush=True)

    os.makedirs(args.save, exist_ok=True)

    vae = load_vae()
    te, tok = (load_text_encoder() if args.te else (None, None))
    print(f"[GPU {process_index}] VAE{'' if args.te else ' (без ТЕ: --te не указан)'} на "
          f"{device}, VRAM {torch.cuda.memory_allocated() / 2**30:.1f} ГиБ", flush=True)
    t0 = time.time()
    try:
        process_in_chunks(vae, te, tok, args, rows, args.save, strip_words)
    finally:
        del vae, te, tok
        clear_cuda_memory()
    accelerator.wait_for_everyone()

    if is_main:
        n_ds, n_rows = count_chunks(args.save)
        print(f"[main] готово за {time.time() - t0:.0f} c: {n_rows} сэмплов в {n_ds} чанках "
              f"-> {args.save} (склейки нет, src/train.py читает чанки сам)", flush=True)


if __name__ == "__main__":
    main()