"""
Creates folder with augmented data.
Author: Martin Tomasovic
"""
from PIL import Image, ImageDraw 

import albumentations as A
import cv2
import matplotlib.pyplot as plt


# Gentle augmentations
transform = A.Compose([
    A.HorizontalFlip(p=0.5),
    A.RandomBrightnessContrast(brightness_limit=0.1, contrast_limit=0.1, p=0.5),
    A.Rotate(limit=10, p=0.7),  # Small rotation
])

transform2 = A.Compose([
    A.GaussianBlur(blur_limit=1, p=0.3),
    A.Affine(scale=(0.95, 1.05), translate_percent=(0.01, 0.03), rotate=(-5, 5), p=0.5),
    A.HueSaturationValue(hue_shift_limit=5, sat_shift_limit=10, val_shift_limit=10, p=0.4),
    A.HorizontalFlip(p=0.3),
])

transform3 = A.Compose([
    A.RandomBrightnessContrast(p=0.3),
    A.Rotate(limit=5, p=0.5),
])

import os

parent_dir = 'Augumented_data_2_all/dev/images'

for root, dirs, files in os.walk(parent_dir):
    # get class name
    class_name = os.path.basename(root)
    print(class_name)
    # skip the root directory
    if root == parent_dir:
        continue
    
    for file in files:
        # load image
        img_pth = 'Augumented_data_2_all/dev/images/' + class_name + "/" + file
        img = cv2.imread(img_pth)  # load image as NumPy array
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # convert to RGB for matplotlib
        plt.axis('off')
        plt.imshow(img)

        # augument the image
        augmented_images = []
        for i in range(0,5):
            augmented_images.append(transform(image=img)['image'])

        for i in range(0,5):
            augmented_images.append(transform2(image=img)['image'])

        #for i in range(0,3):
        #    augmented_images.append(transform3(image=img)['image'])
        
        # save the 
        for idx, aimg in enumerate(augmented_images):
            aimg_pth = img_pth.replace(".png", f"_aug_{idx}.png") 
            plt.imsave(aimg_pth, aimg)