#!/usr/bin/env python3
"""Правила отбраковки мусорных капшенов (теги gelbooru).

Живут отдельно от обучения: `filter_dataset.py` применяет их один раз к готовому датасету,
а `train.py` о них ничего не знает — обучается на том, что осталось.

Основано на просмотре выборки: `multiple views` — это character sheet и спрайт-листы
(модель учится рисовать двух похожих персонажей рядом), `speech bubble`/`english text` —
источники каши вместо букв в кадре.
"""
import re

FILTER_RULES = {
    # листы персонажа, ракурсы, комиксы
    "views": ("multiple views", "character sheet", "concept art", "reference sheet",
              "4koma", "comic", "manga cover", "collage", "split screen", "multiple angles"),
    # речь, надписи, подписи авторов
    "text": ("speech bubble", "thought bubble", "english text", "japanese text",
             "chinese text", "korean text", "mixed-language text", "text focus",
             "signature", "artist name", "copyright name", "character name", "watermark",
             "letterboxed", "logo", "body writing", "clothes writing", "writing", "engrish"),
    # монохром и незаконченные рисунки
    "mono": ("monochrome", "greyscale", "sketch", "lineart"),
    # несколько персонажей в кадре
    "multi": ("multiple girls", "multiple boys", "2girls", "3girls", "4girls", "5girls",
              "6+girls", "2boys", "3boys", "4boys", "5boys", "6+boys", "group"),
    # низкое качество и битая анатомия: модель не должна учиться на браке
    "quality": ("lowres", "jpeg artifacts", "bad anatomy", "bad hands", "bad proportions",
                "bad quality", "worst quality", "low quality", "poorly drawn", "blurry",
                "partially visible"),
    # обложки, официальный арт, сканы, референсы
    "ref": ("official art", "cover", "promotional", "reference", "scan", "artbook", "poster"),
    # размонтированные тела и лишние конечности
    "body": ("disembodied", "amputee", "multiple heads", "extra arms", "extra legs",
             "extra limbs", "conjoined", "vore", "tentacles"),
    # обрезанные кадры: фигура не влезает в кадр
    "crop": ("out of frame", "out-of-frame", "cropped"),
    # без людей: пейзажи, предметы, еда
    "nohuman": ("no humans", "background only", "still life", "scenery", "object focus",
                "food focus"),
}

FILTER_GROUPS = tuple(FILTER_RULES) + ("long",)

# Капшены длиннее этого числа тегов — скорее всего нестандартная сцена или каша
# (медиана в датасете 16, p99 — 54): 14 % строк.
MAX_TAGS = 30

# Тег `sign` ловится только как отдельный тег: подстрокой он срабатывает на `design`
# (91 строка из 70k). Регулярка дороже сравнения подстрокой, поэтому выполняется только
# для группы text.
_SIGN_RE = re.compile(r"(^|[\s,])sign([\s,]|$)")


def filter_reason(text, groups):
    """Причина отбраковки строки или None. `text` — капшен как есть из датасета.

    Сравнение по подстроке в нижнем регистре: капшен — это список тегов через запятую,
    поэтому `"2girls"` не поймает `"12girls"` (такого тега нет), и ложных срабатываний
    на подстроках внутри других слов практически нет.
    """
    low = text.lower()
    for group in groups:
        if group == "long":  # не тег, а ограничение на число тегов
            if len(text.split(",")) > MAX_TAGS:
                return group
            continue
        for key in FILTER_RULES[group]:
            if key in low:
                return group
        if group == "text" and _SIGN_RE.search(low):
            return group
    return None