"""
WD14 harvest: convierte los tags crudos en una biblioteca de POSE / ESCENA / ROPA.

Objetivo: quedarse SOLO con lo que se trasplanta a tus personajes (pose, estado
NSFW, ropa y entorno) y tirar todo lo que describe a la chica concreta.

    "chica rubia, ojos azules, piel tostada, cola de gato, playa, bikini rojo,
     signo de paz"
      -> DROP  blonde_hair, blue_eyes, tan, cat_tail     (descripcion/personaje)
      -> KEEP  peace_sign, red_bikini, beach            (pose + ropa + entorno)

Salida en Archivos/wildcards_harvest/:
    pose.txt scene.txt outfit.txt state.txt camera.txt   (frecuencia, dedupeado)
    por_imagen.txt                                       (una linea por imagen)
    frecuencia.txt                                       (ranking + que se echó y por qué)
    otros.txt                                            (lo no clasificado, a revisar)

    python Engine/wd14_harvest.py
    python Engine/wd14_harvest.py --min-count 2
"""
import argparse
import csv
import re
from collections import Counter, defaultdict
from datetime import datetime
from pathlib import Path

BASE_DIR = Path(__file__).resolve().parent.parent
DANBOORU = BASE_DIR / "Archivos" / "Danbooru tags.csv"

RATING_CATEGORY = "9"

DROP_COUNTS = {
    "1girl", "1boy", "solo", "solo_focus", "multiple_girls", "2girls", "3girls",
    "4girls", "5girls", "6+girls", "1other", "everyone", "hetero", "couple",
    "2boys", "3boys", "4boys", "5boys", "6+boys", "male_focus", "no_humans",
}

# Emoticonos, texto, marcas de agua y calidad de imagen. Nunca sirven.
DROP_JUNK = re.compile(
    r"^[:;=8xXA><*_\-]+$|"
    r"text$|_text$|speech_bubble|thought_bubble|comic|monochrome|greyscale|"
    r"grayscale|border$|logo|watermark|signature|username|patreon|twitter|"
    r"web_address|artist|meme|lowres|pixelated|traditional_media|comment|"
    r"english|korean|chinese|japanese_text|oekaki|drawn_focus|commentary|"
    r"^\?$|^!$|^v$|^w$|"
    r"motion_lines|speech_lines|focus_comic|panel$|spotlight|"
    r"spoken_question_mark|spoken_heart|spoken_ellipsis|emphasis_lines|"
    r"^heart$|heart_(hands|censor)$|^broken_heart$|"
    r"folded$|^wrinkled$|bubble$|speech$|note$",
    re.I,
)

# Rastro de personaje y DESCRIPCION del cuerpo. El LoRA del personaje ya aporta
# pelo, ojos, piel, senos y complexion: esto solo estorba.
DROP_IDENTITY = re.compile(
    r"hair|eye|pupil|sclera|bangs|ponytail|braid|twintail|ahoge|sidelock|"
    r"fringe|ears?$|animal_ear|horn|tails?$|wings?$|"
    r"face|facial|freckle|mole|marking|makeup|lipstick|nail|tips?$|toenail|"
    r"tongue|teeth|tooth|soles$|soles|"
    r"skin|_skin|tan$|colored_skin|veins|veiny|body_fur|fur$|"
    r"breasts?$|boobs|tits|breast_size|large_breasts|small_breasts|"
    r"medium_breasts|huge_breasts|flat_chest|busty|bouncing_breasts|"
    r"hanging_breasts|breasts_squeezed_together|breasts_apart|"
    r"thick_thighs|thin_legs|slender|petite|tall$|curvy|plump|muscular$|"
    r"body_curve|wide_hips|narrow_waist|hourglass|thigh_gap|cleavage|"
    r"child$|loli|shota|minigirl|toddler|aged_|schoolgirl|"
    r"furry|anthropomorphic|anthro|interspecies|elf$|demon_girl$|android|"
    r"robot$|alien$|monster_girl$|dragon_girl$|fox_girl$|wolf_girl$|"
    r"rabbit_girl$|cat_girl$|horse_girl$|cow_girl$|bunny_girl$|sheep_girl$|"
    r"piercing$|barbell_piercing|labret|ear_gauntlet|forehead_jewel|"
    r"beads$|ring$|necklace$|saliva$|sweat$|wet$|wet_skin|oil$|"
    r"gloss$|armpits$|stomach$|back$|"
    r"animal$|creature$|personification$|magical_girl$|idol$",
    re.I,
)

# Encuadre. No es postura, va a su propio bucket.
CAMERA = re.compile(
    r"^(cowboy_shot|upper_body|lower_body|full_body|bust|portrait|close_up|"
    r"wide_shot|from_above|from_below|from_front|from_side|from_behind|"
    r"multiple_(views|images)|pov|pov_crotch|point_of_view|depth_of_field|"
    r"bokeh|foreshortening|perspective|zoomed|^cropped$|letterboxed|"
    r"first_person|third_person|dutch_angle|low_angle|high_angle|"
    r"extreme_close_up|full_moon|starry_sky|"
    r".*_focus$|^focus$|^focused$)",
    re.I,
)

# Poses posturales y acciones sobre el cuerpo. Es lo que nos interesa.
POSE = re.compile(
    r"standing|sitting|lying|kneeling|crouching|squatting|leaning|bending|"
    r"bent_over|bent_forward|hugging|holding|arms?_up|arms?_down|arm_up|"
    r"arm_down|arm_held|arms_behind|arms_around|arm_around|arms_in_front|"
    r"arms_crossed|crossed_arms|arms_raised|raised_arms|extended_arms|"
    r"outstretched_arm|hand_(on|over|under|between|behind|to_own|in_own)|"
    r"hands_(on|over|under|between|behind|to_own|in_own|folded|in_front|"
    r"together|clasped)|leg_up|leg_lift|legs_up|legs_spread|legs_together|"
    r"legs_apart|legs_over|crossed_legs|knee_up|knees_up|knee_bent|"
    r"bent_knee|feet_up|feet_on|feet_together|knees_together|"
    r"on_(back|side|stomach|all_fours|bed|chair|couch|table|floor|wall|"
    r"ground|knees|shoulders|hands|feet|top|desk)|"
    r"spread_|arched_back|arch_back|back_arched|twisting|turning|bowing|"
    r"jumping|walking|running|walking_away|head_tilt|leaning_back|"
    r"leaning_forward|reclining|perched|straddling|riding|mounting|^mount$|"
    r"cowgirl_position|reverse_cowgirl|squatting_cowgirl|standing_dildo|"
    r"missionary|doggystyle|doggy_style|upright_straddle|all_fours|"
    r"on_all_fours|prone_bone|presenting|pulling_up|clothing_aside|"
    r"clothes_lift|shirt_lift|skirt_lift|top_lift|undressing|"
    r"facing_away|facing_viewer|grabbing|reaching|touching|pointing|"
    r"head_grab|chin_grab|arm_grab|leg_grab|hair_grab|ass_grab|"
    r"pulling_hair|grabbing_own|grabbing_another|hand_on_own|hands_on_own|"
    r"hands_in_own_hair|hands_together|wrists_cuffed|cuffed|hogtied|"
    r"suspended_congress|full_nelson|leaning_on|perched_on|sitting_on|"
    r"kneeling_on|lying_on|standing_on|looking_back$|squatting$|"
    r"lifting_person|carrying$|"
    r".*_position$|.*_penetration$|holding_",
    re.I,
)

# Estado NSFW y fluidos: define el contenido explicito.
STATE = re.compile(
    r"^nude$|naked|topless|bottomless|undressed|partially_nude|completely_nude|"
    r"nipple|areola|pussy|clitoris|vagina|vaginal|anus|(?:^|_)anal(?:$|_)|oral_sex|"
    r"penis|^cock$|(?:^|_)dick(?:$|_)|testicle|erection|^balls$|"
    r"^cum(?:$|_)|cum_?(on|in|over|inside|face|mouth|pussy|ass|body|breasts|tits|hand|"
    r"wall|floor|drip)|cumdrip|(?:^|_)semen(?:$|_)|ejaculation|creampie|pussy_juice|"
    r"pussy_peek|vaginal_fluid|blowjob|fellatio|handjob|paizuri|after_paizuri|"
    r"titjob|footjob|deepthroat|^sex$|sex_from_|clothed_sex|group_sex|"
    r"threesome|spread_pussy|spread_ass|ahegao|orgasm|climax|orgasmic|"
    r"moaning|orgasmic|heavy_breathing|open_mouth|closed_mouth|parted_lips|"
    r"tongue_out|cross-section|x-ray|internal_view|doll_view|"
    r"censored|mosaic_censoring|bar_censor|uncensored|convenient_censoring|"
    r"after_sex|not_worn|breasts_out|between_breasts|underboob|sideboob|"
    r"netorare|panties_aside|covering_breasts|covered_nipples|"
    r"female_pubic_hair|male_pubic_hair|pubic_hair|fetish$|"
    r"^oral$|cameltoe|internal_cumshot|creampie_on_face|"
    r"clothed_(female|male)_nude|bathing$|shared_bathing|"
    r"^afterward$|postsex|post_orgasm",
    re.I,
)

# Prendas y accesorios de ropa.
OUTFIT = re.compile(
    r"dress|shirt|skirt|panties|^bra$|bra_|thighhighs|kneehighs|stockings|"
    r"pantyhose|uniform|gloves|socks|shoes|boots|jacket|coat|sweater|apron|"
    r"lingerie|bikini|swimsuit|leotard|bodysuit|harness|collar|^tie$|necktie|"
    r"bow|ribbon|scarf|^belt$|jewelry|earrings|necklace|neckwear|bracelet|"
    r"anklet|wristband|armband|robe|kimono|miko|armor|top$|top_|bottom$|"
    r"legwear|underwear|^pants$|panty|corset|garter|stocking|camisole|"
    r"blouse|cardigan|turtleneck|blazer|vest|waistband|headband|ear_cover|"
    r"choker|maid_apron|maid_headdress|frills?|lace|satin|leather|fur_trim|"
    r"highleg|cropped|crop_top|midriff|open_clothes|see-through|nightgown|"
    r"overalls|bag$|babydoll|playboy_bunny|^thong$|sleeves|"
    r"detached_sleeves|juliet_sleeves|puffy_sleeves|white_sleeves|"
    r"bloomers|buttondown|dress_shirt|clothing|costume|^hat$|^cap$|"
    r"^maid$|^maid_|witch_hat|witch$|floral|_print$|^sleeveless$|sleeve|"
    r"^sunglasses$|^glasses$|.*_bra$|barefoot|"
    r"uniform$|bride_",
    re.I,
)

# Escenario: lugar, atrezo, luz y clima.
SCENE = re.compile(
    r"bedroom|bathroom|beach|outdoor|indoor|background|^sky$|cloud|^sun$|"
    r"^moon$|night|^day$|^room$|^bed$|window|door|^floor$|^wall$|^tree$|"
    r"flower|grass|^water$|ocean|^sea$|mountain|forest|^field$|^street$|"
    r"^city$|school|classroom|gymnasium|^pool$|poolside|shower|bathtub|"
    r"kitchen|living_room|balcony|garden|park|snow|^rain$|sunset|sunrise|"
    r"^dark$|^light$|lighting|^lamp$|candle|neon|^stage$|rooftop|^hill$|"
    r"canyon|desert|savanna|rice_field|temple|^shrine$|castle|^tower$|ruins|"
    r"bridge|harbor|dock|^pier$|silhouette|shadow|sparkle|petal|blossom|"
    r"tile_wall|tiled_floor|curtain|bed_sheet|^pillow$|blanket|bookshelf|"
    r"^desk$|^monitor$|computer_keyboard|laptop|^fire$|fireplace|window_light|"
    r"^nature$|scenery|landscape|horizon|mountain_range|onsen|hot_spring|"
    r"sauna|changing_room|locker_room|white_background|simple_background|"
    r"gradient_background|blue_sky|cloudy|rainy|^office$|hair_bow|hair_ornament|"
    r"^hat$|^cap$|^hood$|^scarf$|^rose$|tulip|sakura|^tv$|remote|"
    r"computer_keyboard|^bottle$|^cup$|^glass$|^plate$|^food$|mirror|"
    r"^stairs$|railing|^fence$|^plant$|potted_plant|houseplant|^vase$|"
    r"^frame$|picture$|bookshelf|desk|monitor|laptop|^pot$|^lamp$",
    re.I,
)

# Expresion facial: util como detalle, va aparte.
EXPRESSION = re.compile(
    r"^smile$|smirk|grin|expressionless|expression|^blush$|blushing|grimace|"
    r"^crying$|^tears?$|yawning|^pout$|embarrassed|surprised|serious$|"
    r"neutral|arrogant|naughty_face|smug|relieved|satisfied|focused$|dazed|"
    r"^wink$|aroused|horny|sleepy|serious$|confident|worried|annoyed|"
    r"doubtful|disappointed|shy$|curious$|humbled",
    re.I,
)

DROP_META = {"the_pose", "posing", "pose", "realistic", "photorealistic",
             "no_outline", "no_shadow", "x-ray_film", "netorare"}

BODY_PARTS = re.compile(
    r"^(feet|toes|foot|legs?|knees?|thighs?|arms?|hands?|fingers?|elbows?|"
    r"shoulders?|hips?|waist|chest|navel|butt|ass|breasts?|shoulder_blades)$",
    re.I,
)


def load_danbooru():
    """-> (tags_de_personaje, tags_de_copyright, alias->canonico)"""
    chars, copies, alias = set(), set(), {}
    if not DANBOORU.exists():
        return chars, copies, alias
    with DANBOORU.open(newline="", encoding="utf-8") as f:
        for row in csv.DictReader(f):
            tag, cat = row["tag"].lower(), row["type"]
            if cat == "4":
                chars.add(tag)
            elif cat == "3":
                copies.add(tag)
            if row["aliases"]:
                for a in row["aliases"].split(","):
                    a = a.strip().lower()
                    if a and a != tag:
                        alias[a] = tag
    return chars, copies, alias


def classify(tag, chars, copies):
    t = tag.lower()
    if t in chars or t in copies:
        return None, "personaje/copyright"
    if t in DROP_COUNTS:
        return None, "conteo"
    if t in DROP_META:
        return None, "basura"
    if DROP_JUNK.search(t):
        return None, "basura"
    if DROP_IDENTITY.search(t):
        return None, "identidad"
    if CAMERA.match(t):
        return "camera", "encuadre"
    if BODY_PARTS.match(t):
        return None, "descripcion anatomica"
    # POSE antes que STATE: 'cowgirl_position' es postura, no estado.
    if POSE.search(t):
        return "pose", "pose"
    if STATE.search(t):
        return "state", "estado NSFW"
    if EXPRESSION.search(t):
        return "expression", "expresion"
    if OUTFIT.search(t):
        return "outfit", "ropa"
    if SCENE.search(t):
        return "scene", "escena"
    return "otros", "sin clasificar"


BUCKETS = ("pose", "scene", "outfit", "state", "expression", "camera", "otros")


def collect(folder, limit=None, recursive=True):
    folder = Path(folder)
    if not folder.exists():
        sys_exit(f"[!] No existe la carpeta: {folder}")
    it = folder.rglob("*") if recursive else folder.iterdir()
    files = sorted(p for p in it
                   if p.is_file() and p.suffix.lower() in IMAGE_EXTS)
    if limit:
        files = files[:limit]
    if not files:
        sys_exit(f"[!] No hay imagenes en {folder}")
    return files


def sys_exit(msg):
    raise SystemExit(msg)


IMAGE_EXTS = (".png", ".jpg", ".jpeg", ".webp", ".bmp")


def main():
    p = argparse.ArgumentParser(description="Cosecha pose/escena/ropa de los tags WD14")
    p.add_argument("--tags", default=str(BASE_DIR / "Archivos" / "wd14" / "tags"))
    p.add_argument("--out", default=str(BASE_DIR / "Archivos" / "wildcards_harvest"))
    p.add_argument("--min-count", type=int, default=1)
    p.add_argument("--spaces", action="store_true", help="guiones bajos -> espacios")
    args = p.parse_args()

    src = Path(args.tags)
    if not src.exists():
        sys_exit(f"[!] No existe {src}. Ejecuta antes wd14_tagger.py")
    files = sorted(src.rglob("*.txt"))
    if not files:
        sys_exit(f"[!] No hay .txt en {src}")

    chars, copies, alias = load_danbooru()
    print(f"[-] {len(files)} imagenes | Danbooru: {len(chars)} personaje, {len(copies)} copyright")

    buckets = defaultdict(Counter)
    dropped = Counter()
    per_image = []
    canonicalized = Counter()

    for f in files:
        raw = [t.strip() for t in f.read_text(encoding="utf-8").split(",") if t.strip()]
        keep, seen = [], set()
        for tag in raw:
            canon = alias.get(tag.lower(), tag.lower())
            if canon != tag.lower():
                canonicalized[f"{tag.lower()} -> {canon}"] += 1
            bucket, reason = classify(canon, chars, copies)
            if bucket is None:
                dropped[reason] += 1
                continue
            if canon in seen:
                continue
            seen.add(canon)
            buckets[bucket][canon] += 1
            keep.append(canon)
        per_image.append((f.relative_to(src).as_posix(), keep, len(raw)))

    out = Path(args.out)
    out.mkdir(parents=True, exist_ok=True)
    fmt = (lambda t: t.replace("_", " ")) if args.spaces else (lambda t: t)

    total_kept = sum(len(k) for _, k, _ in per_image)
    total_raw = sum(r for _, _, r in per_image)

    for bucket in BUCKETS:
        items = [(t, n) for t, n in buckets[bucket].most_common()
                 if n >= args.min_count]
        if not items:
            continue
        (out / f"{bucket}.txt").write_text(
            "".join(f"{fmt(t)}\n" for t, _ in items), encoding="utf-8")

    with (out / "por_imagen.txt").open("w", encoding="utf-8", newline="\n") as f:
        f.write(f"# {len(per_image)} imagenes | {total_kept}/{total_raw} tags conservados\n")
        f.write("# formato: <imagen>  (<conservados>/<originales>)\n#   <tags sin descripcion>\n\n")
        for name, keep, raw_n in per_image:
            f.write(f"{name}  ({len(keep)}/{raw_n})\n  {', '.join(fmt(t) for t in keep)}\n\n")

    with (out / "frecuencia.txt").open("w", encoding="utf-8", newline="\n") as f:
        f.write(f"# generado {datetime.now():%Y-%m-%d %H:%M} | {len(per_image)} imagenes\n")
        f.write(f"# {total_kept}/{total_raw} tags conservados "
                f"({total_kept/total_raw*100:.0f}%)\n\n")
        f.write("ECHADOS:\n")
        for reason, n in dropped.most_common():
            f.write(f"  {reason:24} {n}\n")
        if canonicalized:
            f.write("\nCANONICALIZADOS (alias -> canonico):\n")
            for k, v in canonicalized.most_common():
                f.write(f"  {v:3}x  {k}\n")
        for bucket in BUCKETS:
            if not buckets[bucket]:
                continue
            f.write(f"\n{'='*70}\n{bucket.upper()}  "
                    f"({len(buckets[bucket])} distintos)\n{'='*70}\n")
            for t, n in buckets[bucket].most_common():
                f.write(f"  {n:3}x  {fmt(t)}\n")

    print(f"[-] {total_kept}/{total_raw} tags conservados ({total_kept/total_raw*100:.0f}%)")
    print("\nECHADOS:")
    for reason, n in dropped.most_common():
        print(f"  {reason:24} {n}")
    print("\nCONSERVADOS:")
    for b in BUCKETS:
        if buckets[b]:
            print(f"  {b:12} {len(buckets[b]):4} tags distintos")
    print(f"\n[OK] -> {out}")


if __name__ == "__main__":
    main()