import numpy as np
import cv2
from random import randint
from tqdm import tqdm


def flatten_array(arr):
    return arr.reshape(-1, *arr.shape[2:])

def ImageDataAugmentation(img, mask, num_augmentations):
    augmented_images = []
    augmented_masks = []

    for i in range(num_augmentations):
        
        # Apply a random rotation to the images
        angle = randint(-15, 15)
        M = cv2.getRotationMatrix2D((img.shape[1] / 2, img.shape[0] / 2), angle, 1)
        rotated_img = cv2.warpAffine(img, M, (img.shape[1], img.shape[0]))
        rotated_mask = cv2.warpAffine(mask, M, (mask.shape[1], mask.shape[0]))

        # Apply a random horizontal flip to the images
        if randint(0, 1):
            flipped_img = cv2.flip(rotated_img, 1)
            flipped_mask = cv2.flip(rotated_mask, 1)
        else:
            flipped_img = rotated_img
            flipped_mask = rotated_mask

        # Append the augmented images and masks to the lists
        augmented_images.append(flipped_img)
        augmented_masks.append(flipped_mask)

    # Convert the lists of augmented images and masks to NumPy arrays
    augmented_images = np.array(augmented_images)
    augmented_masks = np.array(augmented_masks)

    return augmented_images, augmented_masks

def DataAugmentation (imgs_train, imgs_mask_train, num_augmentations):
    augmented_images_train = []
    augmented_masks_train = []

    for i in tqdm(range(len(imgs_train))):
        augmented_images, augmented_masks = ImageDataAugmentation(imgs_train[i,:,:,0], imgs_mask_train[i,:,:,0], num_augmentations)
        augmented_images_train.append(augmented_images)
        augmented_masks_train.append(augmented_masks)

    augmented_images_train = np.array(augmented_images_train)
    augmented_masks_train = np.array(augmented_masks_train)

    augmented_images_train = flatten_array(augmented_images_train)
    augmented_masks_train = flatten_array(augmented_masks_train)

    return augmented_images_train, augmented_masks_train