import os
import cv2
import albumentations as A
from albumentations.pytorch import ToTensorV2
import numpy as np
import random

# === CONFIG ===
PASTA_ORIGINAL = "dataset/original"
PASTA_AUGMENTED = "dataset/augmented"
NUM_AUGMENTACOES = 3  # quantas imagens gerar por imagem original

os.makedirs(os.path.join(PASTA_AUGMENTED, "images"), exist_ok=True)
os.makedirs(os.path.join(PASTA_AUGMENTED, "labels"), exist_ok=True)

transform = A.Compose([
    A.HorizontalFlip(p=0.5),
    A.VerticalFlip(p=0.1),
    A.RandomBrightnessContrast(p=0.3),
    A.Rotate(limit=10, p=0.4),
    A.RandomScale(scale_limit=0.1, p=0.3),
], bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels']))

imagens = [f for f in os.listdir(os.path.join(PASTA_ORIGINAL, "images")) if f.endswith(('.jpg', '.jpeg', '.png'))]

for nome_img in imagens:
    caminho_img = os.path.join(PASTA_ORIGINAL, "images", nome_img)
    caminho_lbl = os.path.join(PASTA_ORIGINAL, "labels", nome_img.replace(".jpg", ".txt").replace(".jpeg", ".txt").replace(".png", ".txt"))

    if not os.path.exists(caminho_lbl):
        print(f"[!] Label ausente: {nome_img}")
        continue

    # Carrega imagem e label
    image = cv2.imread(caminho_img)
    height, width = image.shape[:2]

    with open(caminho_lbl, 'r') as f:
        linhas = f.readlines()

    bboxes = []
    class_labels = []
    for linha in linhas:
        parts = linha.strip().split()
        if len(parts) != 5:
            continue
        cls, x, y, w, h = map(float, parts)
        bboxes.append([x, y, w, h])
        class_labels.append(int(cls))

    for i in range(NUM_AUGMENTACOES):
        augmented = transform(image=image, bboxes=bboxes, class_labels=class_labels)
        img_aug = augmented['image']
        bboxes_aug = augmented['bboxes']
        labels_aug = augmented['class_labels']

        nome_base = os.path.splitext(nome_img)[0]
        nome_img_out = f"{nome_base}_aug{i}.jpg"
        nome_lbl_out = f"{nome_base}_aug{i}.txt"

        cv2.imwrite(os.path.join(PASTA_AUGMENTED, "images", nome_img_out), img_aug)

        with open(os.path.join(PASTA_AUGMENTED, "labels", nome_lbl_out), 'w') as f:
            for cls, bbox in zip(labels_aug, bboxes_aug):
                x, y, w, h = bbox
                f.write(f"{cls} {x:.6f} {y:.6f} {w:.6f} {h:.6f}\n")

        print(f"[+] Augmentado: {nome_img_out}")

print("\n✅ Augmentation finalizado!")