import os
from torchvision import transforms
from PIL import ImageOps, Image
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

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

    def __call__(self, i, img, mask):
        transforms.RandomHorizontalFlip(p=0.5)
        t = self.transforms[i]
        img = t(img)
        mask = t(mask)
        return img, mask
    
# Função para aplicar flip horizontal e vertical
def flip_image(img):
    return ImageOps.mirror(ImageOps.flip(img))

# Função para rotacionar a imagem em ângulos fixos
def rotate_image(img, angle):
    return TF.rotate(img, angle)

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))
    
    return Image.fromarray(img_transformed)

# Definindo as transformações
transform_list = [
    T.Lambda(lambda img: flip_image(img)),  # Aplica flip nos dois eixos
    T.Lambda(lambda img: rotate_image(img, 90)),  # Rotação de 90 graus
    T.Lambda(lambda img: rotate_image(img, -90)),  # Rotação de -90 graus
    T.Lambda(lambda img: perspective_image(img, magnitude=0.2)),
    T.Lambda(lambda img: perspective_image(img, magnitude=0.1)),
    # Aqui, você pode adicionar outras transformações fixas conforme necessário
    #T.Resize((256, 256)),  # Redimensionamento para o tamanho desejado
    T.ToTensor(),  # Converte as imagens PIL para tensores PyTorch
    # Normalização pode ser adicionada aqui, se necessário
]

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

# Caminhos para os diretórios onde suas imagens e máscaras originais estão armazenadas
dataset_path = 'dataset/original/images'
masks_path = 'dataset/original/masks'

# Caminhos para os diretórios onde as imagens e máscaras aumentadas serão salvas
augmented_images_path = 'dataset/augmented/images'
augmented_masks_path = 'dataset/augmented/masks'

# 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=5):
    # 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=5)

print("Augmentation completed!")