import json, os, cv2
from PIL import Image
import albumentations as A

# ⚙️ Configurações
with open("config.json", "r") as f:
    config = json.load(f)
MODELO = config["camera"]

# Pastas
dataset_path = os.path.join(MODELO, "dataset", "original", "images")
masks_path   = os.path.join(MODELO, "dataset", "original", "masks")
aug_img_out  = os.path.join(MODELO, "dataset", "augmented", "images")
aug_msk_out  = os.path.join(MODELO, "dataset", "augmented", "masks")
os.makedirs(aug_img_out, exist_ok=True)
os.makedirs(aug_msk_out, exist_ok=True)

# Pipeline de augmentations
train_tf = A.Compose([
    A.HorizontalFlip(p=0.5),

    # Geométricas (aplicam em imagem e máscara)
    A.ShiftScaleRotate(
        shift_limit=0.01, 
        scale_limit=0.10, 
        rotate_limit=5,
        border_mode=cv2.BORDER_REFLECT_101, 
        #value=(255,255,255),
        #mask_value=(255,255,255),
        interpolation=cv2.INTER_LINEAR,
        p=0.3
    ),

    # Fotométricas (somente imagem)
    A.OneOf([
        A.RandomBrightnessContrast(0.2, 0.2, p=1),
        A.HueSaturationValue(hue_shift_limit=5, sat_shift_limit=20, val_shift_limit=15, p=1),
        A.RandomGamma(gamma_limit=(90,110), p=1),
    ], p=0.7),

    A.OneOf([
        A.MotionBlur(blur_limit=3, p=1),
        A.GaussianBlur(blur_limit=3, p=1),
    ], p=0.20),

    A.OneOf([
        A.GaussNoise(var_limit=(5.0, 15.0), p=1),
        A.ImageCompression(quality_lower=50, quality_upper=85, p=1),
    ], p=0.20),

    A.RandomShadow(p=0.1),
    A.RandomSunFlare(p=0.1),
    A.ChannelShuffle(p=0.05),
    A.CoarseDropout(max_holes=6, max_height=16, max_width=16, p=0.1)

    # Resize final (img=LINEAR, mask=NEAREST)
    #A.Resize(height=H, width=W, interpolation=cv2.INTER_LINEAR, mask_interpolation=cv2.INTER_NEAREST),
], additional_targets={'mask':'mask'})

def load_rgb(path):
    # cv2 lê BGR → converte pra RGB (Albumentations usa RGB por padrão)
    im = cv2.imread(path, cv2.IMREAD_COLOR)
    if im is None:
        raise FileNotFoundError(path)
    return cv2.cvtColor(im, cv2.COLOR_BGR2RGB)

def save_rgb(path, arr_rgb):
    # Salva em RGB mantendo cores corretas
    Image.fromarray(arr_rgb).save(path)

def augment_images_and_masks(n_copies=6):
    # Faz pareamento por nome base (sem extensão)
    imgs = sorted([f for f in os.listdir(dataset_path) if os.path.isfile(os.path.join(dataset_path,f))])
    msks = sorted([f for f in os.listdir(masks_path)   if os.path.isfile(os.path.join(masks_path,f))])

    # Mapeia máscaras por nome-base
    msk_map = {os.path.splitext(m)[0]: m for m in msks}

    total = 0
    for img_file in imgs:
        base, ext = os.path.splitext(img_file)
        if base not in msk_map:
            print(f"[WARN] Máscara não encontrada para {img_file}, pulando.")
            continue

        img_path = os.path.join(dataset_path, img_file)
        msk_path = os.path.join(masks_path,   msk_map[base])

        # Carrega RGB (máscara como RGB também — mantemos as cores exatas)
        img = load_rgb(img_path)
        msk = load_rgb(msk_path)

        for i in range(n_copies):
            # Aplica aug; máscara recebe só geométricas
            aug = train_tf(image=img, mask=msk)
            img_aug = aug["image"]
            msk_aug = aug["mask"]

            # Salva
            out_img = os.path.join(aug_img_out, f"{base}_aug_{i:02d}{ext}")
            out_msk = os.path.join(aug_msk_out, f"{base}_aug_{i:02d}{os.path.splitext(msk_map[base])[1]}")
            save_rgb(out_img, img_aug)
            save_rgb(out_msk, msk_aug)
            total += 1

    print(f"Augmentation completed! {total} pares gerados.")

if __name__ == "__main__":
    augment_images_and_masks(n_copies=5)