#!/usr/bin/env python3
"""Отфильтровать готовый датасет (`make_dataset.py`) в новый, не пересчитывая латенты.

Латенты VAE — самая дорогая часть нарезки (часы GPU), а фильтр работает по тексту капшена.
Поэтому здесь ничего не кодируется заново: читаем arrow-файлы, отбрасываем строки
по правилам из `train.py` (FILTER_RULES) и пишем урезанный датасет в `--save`.

Структуру сохраняем один-в-один: каждый исходный `chunk_*/sub_*` даёт один новый
`chunk_*/sub_*` (только без отброшенных строк). Поэтому группы уже однородны по размеру
картинки — `make_dataset.py` писал их именно так, — и батчи в `train.py` остаются
одинаковыми по HxW. Никакой перегруппировки и, значит, никаких копий латентов в RAM.

Запуск:
    python3 src/filter_dataset.py --src src/work/ds_testg_full --save src/work/ds_testg_filtered
    python3 src/filter_dataset.py --src ... --save ... --filter views,text,multi,long
"""
import argparse
import os
import shutil
import sys

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import filters as F  # noqa: E402  (правила отбраковки капшенов)
import train as T  # noqa: E402  (только find_dataset_chunks — как train.py ищет чанки)


def parse_args():
    p = argparse.ArgumentParser(description="Датасет -> отфильтрованный датасет (без пересчёта латентов)")
    p.add_argument("--src", required=True, help="папка с готовым датасетом (chunk_*/sub_*)")
    p.add_argument("--save", required=True, help="куда сохранить отфильтрованный")
    p.add_argument("--filter", default="all",
                   help=f"группы через запятую ({', '.join(F.FILTER_GROUPS)}) или 'all'")
    return p.parse_args()


def main():
    args = parse_args()
    groups = tuple(g.strip() for g in args.filter.split(",") if g.strip())
    if "all" in groups:
        groups = F.FILTER_GROUPS
    bad = [g for g in groups if g not in F.FILTER_GROUPS]
    if bad:
        raise SystemExit(f"--filter: неизвестные группы {bad}; есть {list(F.FILTER_GROUPS)}")

    from datasets import load_from_disk
    srcs = sorted(p for p in T.find_dataset_chunks(args.src) if p.startswith(
        os.path.join(args.src, "chunk_")))
    if not srcs:
        raise SystemExit(f"в {args.src} нет чанков chunk_*/sub_*")

    print(f"группы: {', '.join(groups)}", flush=True)
    os.makedirs(args.save, exist_ok=True)

    dropped = {g: 0 for g in F.FILTER_GROUPS}
    total = kept = 0
    for src in srcs:
        ds = load_from_disk(src)
        keep = [i for i, text in enumerate(ds["text"])
                if not _is_musor(text, groups, dropped)]
        total += len(ds)
        if not keep:
            print(f"  {os.path.relpath(src, args.src)}: 0 из {len(ds)} — группа пустая", flush=True)
            continue
        rel = os.path.relpath(src, args.src)
        dst = os.path.join(args.save, rel)
        if os.path.isdir(dst):
            shutil.rmtree(dst)  # пересборка с нуля
        os.makedirs(os.path.dirname(dst), exist_ok=True)
        ds.select(keep).save_to_disk(dst)
        open(os.path.join(os.path.dirname(dst), ".done"), "w").close()
        kept += len(keep)
        print(f"  {rel}: {len(keep)} из {len(ds)}", flush=True)

    total_dropped = total - kept
    print(f"\nвсего {total}, отброшено {total_dropped} "
          f"({total_dropped / max(total, 1):.1%}), осталось {kept}", flush=True)
    for g, c in dropped.items():
        if c:
            print(f"  {g:8} {c:6d}  {c / max(total, 1):5.1%}", flush=True)
    print(f"готово -> {args.save}", flush=True)


def _is_musor(text, groups, dropped):
    reason = F.filter_reason(text, groups)
    if reason is not None:
        dropped[reason] += 1
    return reason is not None


if __name__ == "__main__":
    main()