import os
from glob import glob
import sys

from torchvision import transforms as T
from PIL import Image


def create_transform_function():
    transform = T.Compose([
        T.ColorJitter(0.5, 0.3, 0.3),
        T.GaussianBlur(kernel_size=(5, 9), sigma=(0.1, 5)),
        T.RandomHorizontalFlip(p=0.5),
        T.RandomVerticalFlip(p=0.5),
        # T.RandomEqualize(p=0.2),
        T.RandomAutocontrast(p=0.5)
    ])
    return transform

def transform_blur():
    transform = T.Compose([
        T.GaussianBlur(kernel_size=(5, 9), sigma=(0.1, 5)),
    ])
    return transform

def transform_blur_flip():
    transform = T.Compose([
        T.GaussianBlur(kernel_size=(5, 9), sigma=(0.1, 5)),
        T.RandomHorizontalFlip(p=0.5),
        T.RandomVerticalFlip(p=0.5),
    ])
    return transform

if __name__ == "__main__":
    input_image_dir = sys.argv[1]
    image_extension = sys.argv[2]
    output_image_dir = sys.argv[3]
    number_of_copy = sys.argv[4]

    transform = create_transform_function()

    if not os.path.isdir(output_image_dir):
        os.makedirs(output_image_dir, exist_ok=True)

    for image_path in glob(os.path.join(input_image_dir, f"*.{image_extension}")):
        im = Image.open(image_path)

        for index in range(int(number_of_copy)):
            trans_im = transform(im)

            image_basename = os.path.basename(image_path)
            new_image_path = os.path.join(
                output_image_dir, 
                os.path.splitext(image_basename)[0] + f"_{index}" + os.path.splitext(image_basename)[1])

            trans_im.save(new_image_path)

