import os
from torchvision import transforms
from PIL import ImageOps, Image, ImageEnhance, ImageFilter
import torchvision.transforms.functional as TF
from torchvision.transforms.functional import to_pil_image
import torchvision.transforms as T
import cv2
import numpy as np

# ⚙️ Configurações
MODELO = "oak-1"

# Caminhos para os diretórios onde suas imagens e máscaras originais estão armazenadas
dataset_path = os.path.join(MODELO, "dataset", "original", "images")
masks_path = os.path.join(MODELO, "dataset", "original", "masks")

# Caminhos para os diretórios onde as imagens e máscaras aumentadas serão salvas
augmented_images_path = os.path.join(MODELO, "dataset", "augmented", "images")
augmented_masks_path = os.path.join(MODELO, "dataset", "augmented", "masks")

class ComposeWithSeed(object):
    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, i, img, mask):
        transforms.RandomHorizontalFlip(p=0.5)
        apply_mask, t = self.transforms[i]
        img = t(img)
        if apply_mask:
            mask = t(mask)
        return img, mask
    

def gaussian_blur(img, radius=1):
    return img.filter(ImageFilter.GaussianBlur(radius))

def add_gaussian_noise(img, mean=0, std=10):
    arr = np.array(img).astype(np.float32)
    noise = np.random.normal(mean, std, arr.shape)
    arr_noisy = np.clip(arr + noise, 0, 255).astype(np.uint8)
    return Image.fromarray(arr_noisy)

def lighting_more_sun(img):
    t = T.ColorJitter(brightness=0.25, contrast=0.25, saturation=0.15, hue=0.02)
    return t(img)

def lighting_less_sun(img):
    t = T.ColorJitter(brightness=0.25, contrast=0.25, saturation=0.05, hue=0.02)
    img = t(img)
    # Inverter o efeito de "mais sol" — clareia reduzindo brilho/contraste
    enhancer_b = ImageEnhance.Brightness(img)
    img = enhancer_b.enhance(0.75)  # < 1.0 escurece
    enhancer_c = ImageEnhance.Contrast(img)
    img = enhancer_c.enhance(0.85)  # < 1.0 reduz contraste
    return img

def flip_horizontal(img):
    return ImageOps.mirror(img)

def flip_vertical(img):
    return ImageOps.flip(img)  # ou TF.vflip(img)

def rotate_image(img, angle):
    return TF.rotate(img, angle, fill=(255,255,255))

def perspective_image(img, magnitude=0.5):
    width, height = img.size
    
    # Pontos de origem
    points_orig = np.float32([
        [0, 0],
        [width, 0],
        [0, height],
        [width, height]
    ])
    
    # Pontos de destino, deslocados com base na magnitude
    points_dest = np.float32([
        [int(magnitude * width), int(magnitude * height)],
        [int((1 - magnitude) * width), 0],
        [0, int((1 - magnitude) * height)],
        [width, height]
    ])
    
    # Calcula a matriz de transformação e aplica a transformação de perspectiva
    matrix = cv2.getPerspectiveTransform(points_orig, points_dest)
    img_transformed = cv2.warpPerspective(np.array(img), matrix, (width, height), borderValue=(255,255,255))
    
    return Image.fromarray(img_transformed)

# Definindo as transformações
transform_list = [
    (True, T.Lambda(lambda img: flip_horizontal(img))),  # Aplica flip na horizontal
    #(True, T.Lambda(lambda img: flip_vertical(img))),  # Aplica flip na vertical
    #(True, T.Lambda(lambda img: rotate_image(img, 90))),  # Rotação de 90 graus
    #(True, T.Lambda(lambda img: rotate_image(img, -90))),  # Rotação de -90 graus
    #(True, T.Lambda(lambda img: perspective_image(img, magnitude=0.2))),
    #(True, T.Lambda(lambda img: perspective_image(img, magnitude=0.1))),
    (False, T.Lambda(lambda img: lighting_more_sun(img))),
    (False, T.Lambda(lambda img: lighting_less_sun(img))),
    (False, T.Lambda(lambda img: gaussian_blur(img, radius=1))),
    (False, T.Lambda(lambda img: add_gaussian_noise(img, std=8))),
    #T.ToTensor(),  # Converte as imagens PIL para tensores PyTorch
]

# Agora, definimos a transformação composta com a classe personalizada
transform = ComposeWithSeed(transform_list)

# Verifica se os diretórios de destino existem, caso contrário, cria os diretórios
os.makedirs(augmented_images_path, exist_ok=True)
os.makedirs(augmented_masks_path, exist_ok=True)

# Função para aplicar a transformação e salvar as imagens e máscaras transformadas
def augment_images_and_masks(dataset_path, masks_path, augmented_images_path, augmented_masks_path, transform, num_copies):
    # Lista todos os arquivos nos diretórios do dataset de imagens e máscaras
    image_files = [f for f in os.listdir(dataset_path) if os.path.isfile(os.path.join(dataset_path, f))]
    mask_files = [f for f in os.listdir(masks_path) if os.path.isfile(os.path.join(masks_path, f))]

    for image_file, mask_file in zip(image_files, mask_files):
        image_path = os.path.join(dataset_path, image_file)
        mask_path = os.path.join(masks_path, mask_file)
        image = Image.open(image_path).convert('RGB')
        mask = Image.open(mask_path).convert('RGB')

        for i in range(num_copies):
            # Aplica a transformação de maneira consistente em ambos, imagem e máscara
            transformed_image, transformed_mask = transform(i, image, mask)

            # Salva a imagem e a máscara transformadas
            image_save_path = os.path.join(augmented_images_path, f"{os.path.splitext(image_file)[0]}_aug_{i}{os.path.splitext(image_file)[1]}")
            mask_save_path = os.path.join(augmented_masks_path, f"{os.path.splitext(mask_file)[0]}_aug_{i}{os.path.splitext(mask_file)[1]}")

            #transformed_image_pil = to_pil_image(transformed_image)
            transformed_image.save(image_save_path)
            #transformed_mask_pil = to_pil_image(transformed_mask)
            transformed_mask.save(mask_save_path)

# Chama a função para iniciar o processo de aumento de dados
augment_images_and_masks(dataset_path, masks_path, augmented_images_path, augmented_masks_path, transform, num_copies=len(transform_list))

print("Augmentation completed!")