#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Split estratificado por GRUPO com **val/test só do ORIGINAL** e
garantia de NÃO VAZAMENTO entre splits, agora para estrutura RAW:

Lê de:
  MODELO/dataset/<WxH>/group/<grupo>/{previews,raws,masks,(masks2)}

Escreve em:
  MODELO/dataset/split/<split>/group/<grupo>/{previews,raws,(masks),(masks2)}

Definições:
- "Família" = todas as variações da MESMA base original:
    original_<base>.* e augmented_<base>_aug_XX.*
- Val/Test: só **original_<base>** (sem augmented)
- Train: original_<base> **e** todos augmented_<base>_aug_XX

Baseado no _7_split.py original (versão images/masks). :contentReference[oaicite:1]{index=1}
"""
import os
import re
import json
import shutil
import random
import argparse

# ⚙️ Configurações
with open("config.json", "r", encoding="utf-8") as f:
    config = json.load(f)
MODELO = config.get("camera")
USE_MASKS2 = config.get("dual_head", False)
RESOLUCAO = tuple(config.get("resolucao"))

# Pastas (ajustadas para PREVIEWS/RAWS)
pasta_origem = os.path.join(MODELO, "dataset", f"{RESOLUCAO[0]}x{RESOLUCAO[1]}", "group")
pasta_destino = os.path.join(MODELO, "dataset", "split")

IMG_EXTS = (".jpg", ".jpeg", ".png")
MSK_EXT  = ".png"   # máscaras normalizadas em PNG
RAW_EXT  = ".raw"

# Regex para identificar famílias
RE_ORIGINAL_PREFIX  = re.compile(r'^original_(.+)$', re.IGNORECASE)
RE_AUGMENTED_FAMILY = re.compile(r'^augmented_(.+?)(?:_aug[a-zA-Z0-9]*_\d+)?$', re.IGNORECASE)
RE_AUG_SUFFIX       = re.compile(r'_aug[a-zA-Z0-9]*_\d+$', re.IGNORECASE)


def garantir(p):
    os.makedirs(p, exist_ok=True)


def lista_grupos(root):
    """Lista grupos que têm previews + masks (raws opcional)."""
    if not os.path.isdir(root):
        return []
    out = []
    for g in sorted(os.listdir(root)):
        gdir = os.path.join(root, g)
        if not os.path.isdir(gdir):
            continue
        if os.path.isdir(os.path.join(gdir, "previews")) and os.path.isdir(os.path.join(gdir, "masks")):
            out.append(g)
    return out


def listar_previews(prev_dir):
    if not os.path.isdir(prev_dir):
        return []
    fs = []
    for f in os.listdir(prev_dir):
        ext = os.path.splitext(f.lower())[1]
        if ext in IMG_EXTS:
            fs.append(f)
    return sorted(fs)


def mask_from_image_name(img_name):
    base, _ = os.path.splitext(img_name)
    return base + MSK_EXT


def mask2_from_image_name(img_name):
    base, _ = os.path.splitext(img_name)
    return base + MSK_EXT  # masks2 também normalizadas em PNG


def raw_from_image_name(img_name):
    base, _ = os.path.splitext(img_name)
    return base + RAW_EXT


def classify_source_and_family(filename_no_ext):
    """
    Retorna (source, family_key)
      source ∈ {"original", "augmented", "unknown"}
      family_key = base associada ao original (sem prefixo/sufixos), ex: "foo_001"
    """
    m = RE_ORIGINAL_PREFIX.match(filename_no_ext)
    if m:
        return "original", m.group(1)

    m = RE_AUGMENTED_FAMILY.match(filename_no_ext)
    if m:
        return "augmented", m.group(1)

    # legado: tenta deduzir se é augmented por sufixo, e família é o próprio nome sem sufixo
    if RE_AUG_SUFFIX.search(filename_no_ext):
        fam = RE_AUG_SUFFIX.sub("", filename_no_ext)
        return "augmented", fam

    return "unknown", filename_no_ext


def build_family_index(prev_dir, msk_dir):
    """
    Constroi índice de famílias a partir de previews/masks.
    Retorna: dict family -> {"original": str|None, "augmented": [str], "all": [str]}
    (strings são NOMES DE ARQUIVO, não paths completos; assumem que a máscara existe)
    """
    familias = {}
    imgs = listar_previews(prev_dir)
    for img_name in imgs:
        base_no_ext, ext = os.path.splitext(img_name)
        mask_name = mask_from_image_name(img_name)
        if not os.path.exists(os.path.join(msk_dir, mask_name)):
            continue  # garante pareamento preview/mask

        source, fam = classify_source_and_family(base_no_ext)
        d = familias.setdefault(fam, {"original": None, "augmented": [], "all": []})
        d["all"].append(img_name)
        if source == "original":
            d["original"] = img_name
        elif source == "augmented":
            d["augmented"].append(img_name)
        else:
            # trata como original desconhecido para não perder dado
            if d["original"] is None:
                d["original"] = img_name
            else:
                d["augmented"].append(img_name)
    return familias


def allocate_counts(n, p_train, p_val, p_test, min_train, min_val, min_test):
    """
    Mesmo esquema do script original: calcula quantas FAMÍLIAS vão para train/val/test.
    """
    n_train = int(round(n * p_train))
    n_val   = int(round(n * p_val))
    n_test  = n - n_train - n_val

    if n_test < 0:
        excesso = -n_test
        take_train = min(excesso, max(0, n_train))
        n_train -= take_train
        excesso -= take_train
        if excesso > 0:
            take_val = min(excesso, max(0, n_val))
            n_val -= take_val
            excesso -= take_val
        n_test = 0

    min_sum = min_train + min_val + min_test
    if n >= min_sum:
        n_train = max(n_train, min_train)
        n_val   = max(n_val,   min_val)
        n_test  = max(n_test,  min_test)

        total = n_train + n_val + n_test
        while total > n:
            if n_test > min_test:
                n_test -= 1
            elif n_val > min_val:
                n_val -= 1
            elif n_train > min_train:
                n_train -= 1
            else:
                break
            total = n_train + n_val + n_test
        while total < n:
            if n_train - min_train <= n_val - min_val:
                n_train += 1
            else:
                n_val += 1
            total = n_train + n_val + n_test
    else:
        n_train = min(n, max(1, min_train))
        resto = n - n_train
        n_val = max(0, min(resto, min_val))
        n_test = max(0, resto - n_val)

    # ajuste final
    diff = n - (n_train + n_val + n_test)
    if diff != 0:
        if diff > 0:
            # adiciona em train, depois val
            take = min(diff, n - n_train)
            n_train += take
            diff -= take
            if diff > 0:
                n_val += diff
        else:
            diff = -diff
            # tira de test, depois val
            take = min(diff, n_test)
            n_test -= take
            diff -= take
            if diff > 0:
                n_val -= diff

    return n_train, n_val, n_test


def copiar(
    nomes,
    src_prev_dir,
    src_msk_dir,
    dst_prev_dir,
    dst_msk_dir,
    src_msk2_dir=None,
    dst_msk2_dir=None,
    src_raw_dir=None,
    dst_raw_dir=None,
):
    """
    Copia preview + mask (+ mask2) (+ raw) para o split.
    """
    garantir(dst_prev_dir)
    garantir(dst_msk_dir)

    use_msk2 = bool(src_msk2_dir and dst_msk2_dir and os.path.isdir(src_msk2_dir))
    use_raw  = bool(src_raw_dir and dst_raw_dir and os.path.isdir(src_raw_dir))

    if use_msk2:
        garantir(dst_msk2_dir)
    if use_raw:
        garantir(dst_raw_dir)

    moved = 0
    for nome in nomes:
        mask_name = mask_from_image_name(nome)
        src_prev = os.path.join(src_prev_dir, nome)
        src_msk  = os.path.join(src_msk_dir, mask_name)

        if not (os.path.exists(src_prev) and os.path.exists(src_msk)):
            continue

        shutil.copy2(src_prev, os.path.join(dst_prev_dir, nome))
        shutil.copy2(src_msk,  os.path.join(dst_msk_dir,  mask_name))

        if use_msk2:
            m2_name = mask2_from_image_name(nome)
            src_m2 = os.path.join(src_msk2_dir, m2_name)
            if os.path.exists(src_m2):
                shutil.copy2(src_m2, os.path.join(dst_msk2_dir, m2_name))

        if use_raw:
            raw_name = raw_from_image_name(nome)
            src_raw = os.path.join(src_raw_dir, raw_name)
            if os.path.exists(src_raw):
                shutil.copy2(src_raw, os.path.join(dst_raw_dir, raw_name))

        moved += 1
    return moved


def split_group(group_name, p_train, p_val, p_test, seed, mins, caps_map=None):
    src_prev_dir = os.path.join(pasta_origem, group_name, "previews")
    src_msk_dir  = os.path.join(pasta_origem, group_name, "masks")
    src_msk2_dir = os.path.join(pasta_origem, group_name, "masks2")
    src_raw_dir  = os.path.join(pasta_origem, group_name, "raws")

    use_msk2 = USE_MASKS2 and os.path.isdir(src_msk2_dir)
    use_raw  = os.path.isdir(src_raw_dir)

    familias = build_family_index(src_prev_dir, src_msk_dir)

    # apenas famílias que têm ORIGINAL para participar de val/test
    familias_originais = [fam for fam, d in familias.items() if d["original"] is not None]
    total_familias = len(familias_originais)
    if total_familias == 0:
        print(f"[{group_name}] 0 famílias com original, pulando.")
        return {"train": 0, "val": 0, "test": 0, "familias": 0}

    rng = random.Random(seed)
    rng.shuffle(familias_originais)

    n_tr, n_va, n_te = allocate_counts(
        total_familias, p_train, p_val, p_test,
        mins["train"], mins["val"], mins["test"]
    )

    fam_train = set(familias_originais[:n_tr])
    fam_val   = set(familias_originais[n_tr:n_tr+n_va])
    fam_test  = set(familias_originais[n_tr+n_va: n_tr+n_va+n_te])

    # CAP por grupo (apenas no TRAIN)
    if caps_map and group_name in caps_map:
        cap = caps_map[group_name]
        if len(fam_train) > cap:
            fam_list = list(fam_train)
            rng.shuffle(fam_list)
            kept    = set(fam_list[:cap])
            dropped = set(fam_list[cap:])
            fam_train = kept
            print(f"[{group_name}] cap-train-families={cap} → mantidas {len(kept)} famílias, descartadas {len(dropped)} do TRAIN")

    nomes_train, nomes_val, nomes_test = [], [], []

    for fam, d in familias.items():
        if fam in fam_train:
            if d["original"]:
                nomes_train.append(d["original"])
            if d["augmented"]:
                nomes_train.extend(d["augmented"])
        elif fam in fam_val:
            if d["original"]:
                nomes_val.append(d["original"])
        elif fam in fam_test:
            if d["original"]:
                nomes_test.append(d["original"])
        else:
            pass

    # dest dirs
    dest_train_prev = os.path.join(pasta_destino, "train", "group", group_name, "previews")
    dest_train_msk  = os.path.join(pasta_destino, "train", "group", group_name, "masks")
    dest_val_prev   = os.path.join(pasta_destino, "val",   "group", group_name, "previews")
    dest_val_msk    = os.path.join(pasta_destino, "val",   "group", group_name, "masks")
    dest_test_prev  = os.path.join(pasta_destino, "test",  "group", group_name, "previews")
    dest_test_msk   = os.path.join(pasta_destino, "test",  "group", group_name, "masks")

    dest_train_msk2 = os.path.join(pasta_destino, "train", "group", group_name, "masks2") if use_msk2 else None
    dest_val_msk2   = os.path.join(pasta_destino, "val",   "group", group_name, "masks2") if use_msk2 else None
    dest_test_msk2  = os.path.join(pasta_destino, "test",  "group", group_name, "masks2") if use_msk2 else None

    dest_train_raw  = os.path.join(pasta_destino, "train", "group", group_name, "raws") if use_raw else None
    dest_val_raw    = os.path.join(pasta_destino, "val",   "group", group_name, "raws") if use_raw else None
    dest_test_raw   = os.path.join(pasta_destino, "test",  "group", group_name, "raws") if use_raw else None

    m_train = copiar(
        nomes_train,
        src_prev_dir, src_msk_dir,
        dest_train_prev, dest_train_msk,
        src_msk2_dir, dest_train_msk2,
        src_raw_dir,  dest_train_raw
    )
    m_val   = copiar(
        nomes_val,
        src_prev_dir, src_msk_dir,
        dest_val_prev, dest_val_msk,
        src_msk2_dir, dest_val_msk2,
        src_raw_dir,  dest_val_raw
    )
    m_test  = copiar(
        nomes_test,
        src_prev_dir, src_msk_dir,
        dest_test_prev, dest_test_msk,
        src_msk2_dir, dest_test_msk2,
        src_raw_dir,  dest_test_raw
    )

    print(f"[{group_name}] famílias={total_familias}  →  train(imgs)={m_train}, val(imgs)={m_val}, test(imgs)={m_test}")
    return {"train": m_train, "val": m_val, "test": m_test, "familias": total_familias}


def main():
    ap = argparse.ArgumentParser(description="Split estratificado por grupo SEM vazamento (RAW: previews/raws/masks).")
    ap.add_argument("--train", type=float, default=0.70, help="Proporção de treino (default=0.70).")
    ap.add_argument("--val",   type=float, default=0.29, help="Proporção de validação (default=0.29).")
    ap.add_argument("--test",  type=float, default=0.01, help="Proporção de teste (default=0.01).")
    ap.add_argument("--seed",  type=int,   default=42,   help="Seed do embaralhamento (default=42).")

    ap.add_argument("--min-train", type=int, default=1, help="Mínimo de FAMÍLIAS por grupo em train (default=1).")
    ap.add_argument("--min-val",   type=int, default=1, help="Mínimo de FAMÍLIAS por grupo em val (default=1).")
    ap.add_argument("--min-test",  type=int, default=0, help="Mínimo de FAMÍLIAS por grupo em test (default=0).")

    ap.add_argument("--modelo", type=str, default=None, help="Sobrescreve MODELO do config.json.")
    ap.add_argument("--resolucao", type=str, default=None, help="Sobrescreve resolução no formato WxH (ex: 960x544).")

    ap.add_argument("--cap-train-families", type=str, default="",
        help="Mapa 'grupo:cap,...' p/ limitar número de FAMÍLIAS no TRAIN. Ex.: 'chao:350'")

    args = ap.parse_args()

    modelo = args.modelo or MODELO
    if args.resolucao:
        try:
            w, h = args.resolucao.lower().split("x")
            resolucao = (int(w), int(h))
        except Exception:
            resolucao = RESOLUCAO
    else:
        resolucao = RESOLUCAO

    def parse_cap_map(s):
        caps = {}
        if not s:
            return caps
        for item in s.split(","):
            k, v = item.strip().split(":")
            caps[k.strip()] = int(v)
        return caps

    caps_map = parse_cap_map(args.cap_train_families)

    global pasta_origem, pasta_destino
    pasta_origem = os.path.join(modelo, "dataset", f"{resolucao[0]}x{resolucao[1]}", "group")
    pasta_destino = os.path.join(modelo, "dataset", "split")

    soma = args.train + args.val + args.test
    if soma <= 0:
        raise ValueError("Soma de proporções deve ser > 0.")
    p_train = args.train / soma
    p_val   = args.val   / soma
    p_test  = args.test  / soma

    mins = {
        "train": max(0, args.min_train),
        "val":   max(0, args.min_val),
        "test":  max(0, args.min_test),
    }

    garantir(pasta_destino)

    grupos = lista_grupos(pasta_origem)
    if not grupos:
        print(f"[WARN] Nenhum grupo encontrado em: {pasta_origem}")
        return

    random.seed(args.seed)

    total_global = {"train": 0, "val": 0, "test": 0, "familias": 0}
    print(f"Grupos: {', '.join(grupos)}")
    print(f"Proporções normalizadas: train={p_train:.3f}, val={p_val:.3f}, test={p_test:.3f}")
    print(f"Mínimos por grupo (famílias): train={mins['train']} val={mins['val']} test={mins['test']}")

    for g in grupos:
        res = split_group(g, p_train, p_val, p_test, args.seed, mins, caps_map=caps_map)
        for k in total_global.keys():
            total_global[k] += res.get(k, 0)

    print("\nResumo global (imagens copiadas):")
    print(f"  train: {total_global['train']}")
    print(f"  val:   {total_global['val']}")
    print(f"  test:  {total_global['test']}")
    print(f"  famílias (total): {total_global['familias']}")
    print("\n✅ Split RAW sem vazamento concluído!")


if __name__ == "__main__":
    main()