from pathlib import Path
import random
import shutil
from collections import defaultdict

# =========================
# CONFIG
# =========================
BASE_DIR = Path(r".")  # <- altere

IMAGES_IN = BASE_DIR / "images" / "original"
LABELS_IN = BASE_DIR / "labels" / "original"

# percentuais (0..1). Tem que somar 1.0
TRAIN_P = 0.70
VAL_P   = 0.30
TEST_P  = 0.00  # se não quiser test, deixa 0

SEED = 42
COPY_MODE = True  # True=copia, False=move

# extensões de imagem aceitas
IMG_EXTS = [".jpg", ".jpeg", ".png"]

# se você quiser limitar a split por classes específicas, informe aqui (opcional)
# ex: {0, 1} para erva/cana. Se None, pega todas que aparecerem.
ALLOWED_CLASSES = None
# =========================


def ensure_dirs(*dirs: Path):
    for d in dirs:
        d.mkdir(parents=True, exist_ok=True)

def parse_classes_from_label(txt_path: Path):
    """
    Retorna um set de classes presentes no arquivo label YOLO-seg.
    Se vazio/sem arquivo -> set() (background).
    """
    if not txt_path.exists():
        return set()

    content = txt_path.read_text(encoding="utf-8", errors="ignore").strip()
    if not content:
        return set()

    classes = set()
    for line in content.splitlines():
        parts = line.strip().split()
        if not parts:
            continue
        try:
            cls = int(float(parts[0]))
            classes.add(cls)
        except:
            continue

    return classes

def key_from_classes(classes: set):
    """
    Cria chave de estrato por combinação de classes.
    background => "BG"
    ex: {0} => "C0"
        {1} => "C1"
        {0,1} => "C0_C1"
    """
    if not classes:
        return "BG"
    return "_".join([f"C{c}" for c in sorted(classes)])

def proportional_counts(n, p_train, p_val, p_test):
    """
    Converte n e percentuais em contagens inteiras que somam n.
    Faz arredondamento e ajusta no final.
    """
    t = int(round(n * p_train))
    v = int(round(n * p_val))
    s = int(round(n * p_test))
    # Ajusta para somar n
    total = t + v + s
    while total != n:
        if total > n:
            # tira de quem tem mais
            if t >= v and t >= s and t > 0:
                t -= 1
            elif v >= t and v >= s and v > 0:
                v -= 1
            elif s > 0:
                s -= 1
        else:
            # adiciona em quem tem menos
            if t <= v and t <= s:
                t += 1
            elif v <= t and v <= s:
                v += 1
            else:
                s += 1
        total = t + v + s
    return t, v, s

def copy_or_move(src: Path, dst: Path, copy_mode=True):
    if not src.exists():
        return
    if copy_mode:
        shutil.copy2(src, dst)
    else:
        shutil.move(src, dst)

def main():
    # valida percentuais
    psum = TRAIN_P + VAL_P + TEST_P
    if abs(psum - 1.0) > 1e-6:
        raise ValueError(f"TRAIN_P+VAL_P+TEST_P precisa somar 1.0. Atual: {psum}")

    if not IMAGES_IN.exists():
        raise FileNotFoundError(f"Pasta não encontrada: {IMAGES_IN}")
    if not LABELS_IN.exists():
        raise FileNotFoundError(f"Pasta não encontrada: {LABELS_IN}")

    # destinos
    images_train = BASE_DIR / "images" / "train"
    images_val   = BASE_DIR / "images" / "val"
    images_test  = BASE_DIR / "images" / "test"

    labels_train = BASE_DIR / "labels" / "train"
    labels_val   = BASE_DIR / "labels" / "val"
    labels_test  = BASE_DIR / "labels" / "test"

    ensure_dirs(images_train, images_val, images_test, labels_train, labels_val, labels_test)

    # lista imagens (baseadas em IMAGES_IN)
    imgs = []
    for ext in IMG_EXTS:
        imgs.extend(IMAGES_IN.glob(f"*{ext}"))
    imgs = sorted(imgs)

    if not imgs:
        print("[INFO] Nenhuma imagem encontrada em:", IMAGES_IN)
        return

    # monta itens com (stem, img_path, raw_path, label_path, classes, stratum)
    items = []
    all_classes = set()

    for img_path in imgs:
        stem = img_path.stem
        raw_path = IMAGES_IN / f"{stem}.raw"
        label_path = LABELS_IN / f"{stem}.txt"

        classes = parse_classes_from_label(label_path)
        if ALLOWED_CLASSES is not None:
            classes = set(c for c in classes if c in ALLOWED_CLASSES)

        all_classes |= classes
        stratum = key_from_classes(classes)

        items.append({
            "stem": stem,
            "img": img_path,
            "raw": raw_path if raw_path.exists() else None,
            "label": label_path if label_path.exists() else None,
            "classes": classes,
            "stratum": stratum
        })

    # agrupa por estrato
    groups = defaultdict(list)
    for it in items:
        groups[it["stratum"]].append(it)

    random.seed(SEED)

    train_set, val_set, test_set = [], [], []

    # split estratificado por estrato
    for stratum, group_items in groups.items():
        random.shuffle(group_items)
        n = len(group_items)
        nt, nv, ns = proportional_counts(n, TRAIN_P, VAL_P, TEST_P)

        train_set.extend(group_items[:nt])
        val_set.extend(group_items[nt:nt+nv])
        test_set.extend(group_items[nt+nv:nt+nv+ns])

    # embaralha cada split (opcional)
    random.shuffle(train_set)
    random.shuffle(val_set)
    random.shuffle(test_set)

    def report(split_name, split_items):
        counts = defaultdict(int)
        for it in split_items:
            if not it["classes"]:
                counts["BG"] += 1
            else:
                for c in it["classes"]:
                    counts[f"C{c}"] += 1
        return counts

    # Copia/move arquivos
    def place(items_list, img_out: Path, lbl_out: Path):
        for it in items_list:
            # imagem
            copy_or_move(it["img"], img_out / it["img"].name, COPY_MODE)

            # raw (se existir)
            if it["raw"] is not None:
                copy_or_move(it["raw"], img_out / it["raw"].name, COPY_MODE)

            # label: se não existir, cria vazio (background)
            out_label = lbl_out / f"{it['stem']}.txt"
            if it["label"] is not None:
                copy_or_move(it["label"], out_label, COPY_MODE)
            else:
                out_label.write_text("", encoding="utf-8")

    place(train_set, images_train, labels_train)
    place(val_set, images_val, labels_val)
    place(test_set, images_test, labels_test)

    # relatório
    print("========================================")
    print("Split concluído ✅")
    print(f"Total imagens: {len(items)}")
    print(f"Train: {len(train_set)} | Val: {len(val_set)} | Test: {len(test_set)}")
    print("----------------------------------------")
    print("Classes detectadas no dataset:", sorted(all_classes) if all_classes else "Nenhuma (só BG)")
    print("----------------------------------------")
    print("Distribuição (contagem de imagens que contém a classe):")
    print("Train:", dict(report("train", train_set)))
    print("Val  :", dict(report("val", val_set)))
    print("Test :", dict(report("test", test_set)))
    print("----------------------------------------")
    print("Pastas destino:")
    print(" -", images_train)
    print(" -", images_val)
    print(" -", images_test)
    print(" -", labels_train)
    print(" -", labels_val)
    print(" -", labels_test)
    print("========================================")

if __name__ == "__main__":
    main()