import os
import random
import shutil

# === CONFIG ===
PASTA_ORIGEM_IMG = "dataset/1280x720/images"
PASTA_ORIGEM_LBL = "dataset/1280x720/labels"
PASTA_SAIDA = "yolov7/train"

SPLIT = {
    "train": 0.7,
    "val": 0.2,
    "test": 0.1
}

# === PREPARA PASTAS ===
for tipo in SPLIT.keys():
    os.makedirs(os.path.join(PASTA_SAIDA, tipo, "images"), exist_ok=True)
    os.makedirs(os.path.join(PASTA_SAIDA, tipo, "labels"), exist_ok=True)

# === LISTA E EMBARALHA ===
arquivos = [f for f in os.listdir(PASTA_ORIGEM_IMG) if f.endswith(".jpg")]
random.shuffle(arquivos)

total = len(arquivos)
qt_train = int(SPLIT["train"] * total)
qt_val = int(SPLIT["val"] * total)

splits = {
    "train": arquivos[:qt_train],
    "val": arquivos[qt_train:qt_train+qt_val],
    "test": arquivos[qt_train+qt_val:]
}

# === COPIA ===
for tipo, lista in splits.items():
    for nome_img in lista:
        nome_lbl = nome_img.replace(".jpg", ".txt")

        origem_img = os.path.join(PASTA_ORIGEM_IMG, nome_img)
        origem_lbl = os.path.join(PASTA_ORIGEM_LBL, nome_lbl)

        destino_img = os.path.join(PASTA_SAIDA, tipo, "images", nome_img)
        destino_lbl = os.path.join(PASTA_SAIDA, tipo, "labels", nome_lbl)

        shutil.copy2(origem_img, destino_img)

        if os.path.exists(origem_lbl):
            shutil.copy2(origem_lbl, destino_lbl)
        else:
            print(f"[!] Label ausente: {nome_lbl}")

    print(f"[✔] {tipo.upper()} - {len(lista)} imagens")

print("\n✅ Split do dataset para YOLOv7 concluído!")