import os
import cv2
import numpy as np

# === CONFIGURAÇÕES ===
PASTA_BASE = "dataset"
FONTE_DADOS = ["original", "augmented"]  # <<-- novas fontes
LABELMAP_PATH = os.path.join(PASTA_BASE, "labelmap.txt")

RESOLUCOES = {
    "512x512": (512, 512),
    "768x768": (768, 768),
    "384x384": (384, 384)
}

# === Função para ler o labelmap ===
def carregar_labelmap_completo(caminho):
    cor_para_id = {}
    id_para_nome = {}
    cores_bgr = []

    with open(caminho, 'r') as arquivo:
        idx = 0
        for linha in arquivo:
            if linha.startswith("#") or not linha.strip():
                continue
            partes = linha.strip().split(':')
            if len(partes) >= 2:
                nome_classe, cor_rgb_str = partes[0], partes[1]
                r, g, b = map(int, cor_rgb_str.split(','))
                cor_bgr = (b, g, r)  # Corrige para BGR

                if nome_classe.lower() == "ignore":
                    ignore_bgr = cor_bgr
                    continue  # NÃO adiciona ignore no LUT de classes

                cor_para_id[cor_bgr] = idx
                cores_bgr.append(cor_bgr)
                id_para_nome[idx] = nome_classe
                idx += 1

    return cor_para_id, cores_bgr, id_para_nome, ignore_bgr

def converter_mask_rgb_para_ids(img_rgb, mapa, ignore_bgr):
    h, w, _ = img_rgb.shape
    mask = np.ones((h, w), dtype=np.uint8) * 255  # Inicializa como ignore
    for cor, classe_id in mapa.items():
        r, g, b = cor
        cond = (img_rgb[:,:,0]==r) & (img_rgb[:,:,1]==g) & (img_rgb[:,:,2]==b)
        mask[cond] = classe_id
    # Pixels brancos (ou ignore_bgr) continuam como 255
    return mask

# === Início do processamento ===
cor_para_id, cores_bgr, id_para_nome, ignore_bgr = carregar_labelmap_completo(LABELMAP_PATH)

# Cria pastas de saída
for nome_res, dim in RESOLUCOES.items():
    os.makedirs(os.path.join(PASTA_BASE, nome_res, "images"), exist_ok=True)
    os.makedirs(os.path.join(PASTA_BASE, nome_res, "masks"), exist_ok=True)

# Processa cada fonte de dados (original + augmented)
for fonte in FONTE_DADOS:
    if (not os.path.exists(os.path.join(PASTA_BASE, fonte))):
        continue
    pasta_rgb = os.path.join(PASTA_BASE, fonte, "images")
    pasta_masks = os.path.join(PASTA_BASE, fonte, "masks")

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

    for i, nome in enumerate(nomes_arquivos, 1):
        caminho_rgb = os.path.join(pasta_rgb, nome)
        caminho_mask = os.path.join(pasta_masks, nome.replace(".jpg", ".png").replace(".jpeg", ".png"))

        img_rgb = cv2.imread(caminho_rgb)
        if img_rgb is None:
            print(f"[!] Erro ao ler imagem {nome}")
            continue

        # Tenta carregar a máscara RGB (se existir)
        if os.path.exists(caminho_mask):
            img_mask_rgb = cv2.imread(caminho_mask)
            #img_mask_rgb = cv2.cvtColor(img_mask_rgb, cv2.COLOR_BGR2RGB)  # ← CORRIGE isso!
            if img_mask_rgb is not None:
                mask_ids = converter_mask_rgb_para_ids(img_mask_rgb, cor_para_id, ignore_bgr)
            else:
                print(f"[!] Erro ao ler máscara {caminho_mask}, ignorando.")
                mask_ids = None
        else:
            mask_ids = None

        for nome_res, dim in RESOLUCOES.items():
            # Cria nomes únicos baseados na fonte
            nome_saida_img = f"{fonte}_{nome}"
            nome_saida_mask = nome_saida_img.replace(".jpg", ".png").replace(".jpeg", ".png")

            # Redimensiona e salva imagem
            img_resized = cv2.resize(img_rgb, dim, interpolation=cv2.INTER_AREA)
            path_img_saida = os.path.join(PASTA_BASE, nome_res, "images", nome_saida_img)
            cv2.imwrite(path_img_saida, img_resized)

            # Redimensiona e salva máscara (se existir)
            if mask_ids is not None:
                mask_resized = cv2.resize(mask_ids, dim, interpolation=cv2.INTER_NEAREST)
                path_mask_saida = os.path.join(PASTA_BASE, nome_res, "masks", nome_saida_mask)
                cv2.imwrite(path_mask_saida, mask_resized)

        print(f"[{fonte}] [{i}/{total}] Redimensionado: {nome}")

print("\n✅ Concluído com sucesso! Todas as fontes foram processadas.")