import json
import os
import shutil
import random

# ⚙️ Configurações
with open("config.json", "r") as f:
    config = json.load(f)
MODELO = config["camera"]
RESOLUCAO = config["resolucao"]
pasta_origem = os.path.join(MODELO, "dataset", f"{RESOLUCAO[0]}x{RESOLUCAO[1]}")
pasta_destino = os.path.join(MODELO, "dataset", "split")

percent_train = 0.9
percent_val = 0.09
percent_test = 0.01

seed = 42
random.seed(seed)

# === Coleta imagens ===
pasta_rgb = os.path.join(pasta_origem, "images")
pasta_masks = os.path.join(pasta_origem, "masks")

arquivos = sorted([f for f in os.listdir(pasta_rgb) if f.endswith(".jpg") or f.endswith(".jpeg")])

# Embaralha
random.shuffle(arquivos)

# Divide
total = len(arquivos)
n_train = int(total * percent_train)
n_val = int(total * percent_val)

arquivos_train = arquivos[:n_train]
arquivos_val = arquivos[n_train:n_train+n_val]
arquivos_test = arquivos[n_train+n_val:]

conjuntos = {
    "train": arquivos_train,
    "val": arquivos_val,
    "test": arquivos_test
}

# === Função auxiliar ===
def copiar(imagens, conjunto):
    path_img_dest = os.path.join(pasta_destino, conjunto, "images")
    path_mask_dest = os.path.join(pasta_destino, conjunto, "masks")
    os.makedirs(path_img_dest, exist_ok=True)
    os.makedirs(path_mask_dest, exist_ok=True)

    for nome in imagens:
        nome_mask = nome.replace(".jpg", ".png").replace(".jpeg", ".png")
        if not os.path.exists(os.path.join(pasta_masks, nome_mask)):
            continue
        shutil.copy2(os.path.join(pasta_rgb, nome), os.path.join(path_img_dest, nome))
        shutil.copy2(os.path.join(pasta_masks, nome_mask), os.path.join(path_mask_dest, nome_mask))

# === Executa cópia ===
for conjunto, lista in conjuntos.items():
    print(f"[{conjunto}] {len(lista)} arquivos")
    copiar(lista, conjunto)

print("\n✅ Dataset dividido com sucesso!")