import os
import shutil
import tqdm
import cv2
import numpy as np
import imgaug as ia
import imgaug.augmenters as iaa
from imgaug.augmentables.segmaps import SegmentationMapsOnImage
from parse_args import parse_args

ia.seed(1)
# Define our augmentation pipeline.
seq = iaa.Sequential([
    # iaa.Dropout([0.05, 0.2]),      # drop 5% or 20% of all pixels
    # iaa.Sharpen((0.0, 1.0)),       # sharpen the image
    # iaa.Affine(rotate=(-45, 45)),  # rotate by -45 to 45 degrees (affects segmaps)
    # iaa.ElasticTransformation(alpha=50, sigma=5),  # apply water effect (affects segmaps)
    # iaa.ChangeColorTemperature((1100, 10000)),
    iaa.Resize((224, 224)),
    iaa.Rotate((0, 360)),
    iaa.Sometimes(0.5, iaa.Rotate((0, 360))),
    iaa.Sometimes(0.5, iaa.AddToSaturation((-10, 10))),
    iaa.Sometimes(0.5, iaa.LinearContrast((0.6, 1.4))),
    iaa.Sometimes(0.5, iaa.AddToHue((-10, 10))),
    # iaa.AddToSaturation((-10, 10)),
    # iaa.LinearContrast((0.6, 1.4)),
    # iaa.AddToHue((-10, 10)),
    iaa.Fliplr(0.5),
    iaa.Flipud(0.5)
], random_order=True)


def seg_augmentation(args, mode='train'):
    print("*" * 20)
    print("Seg augmentation.")
    print("*" * 20)
    input_path = os.path.join(args.dataset_dir, 'data')
    train_input_dir = os.path.join(input_path, mode)
    augmentation_output_dir = os.path.join(input_path, 'augmentation_' + mode)

    if os.path.exists(augmentation_output_dir):
        shutil.rmtree(augmentation_output_dir)
    os.mkdir(augmentation_output_dir)
    os.mkdir(os.path.join(augmentation_output_dir, 'image'))
    os.mkdir(os.path.join(augmentation_output_dir, 'mask'))

    for path in tqdm.tqdm(os.listdir(os.path.join(train_input_dir, 'image'))):
        base = path.split('.')[0]
        image_path = os.path.join(os.path.join(train_input_dir, 'image'), path)
        mask_path = image_path.replace('image', 'mask')
        image = cv2.imread(image_path)
        mask = cv2.imread(mask_path)
        shutil.copy(image_path, os.path.join(augmentation_output_dir, 'image',
                                             base + '[ORIGIN].png'))
        shutil.copy(mask_path, os.path.join(augmentation_output_dir, 'mask',
                                            base + '[ORIGIN][MASK].png'))
        mask = SegmentationMapsOnImage(mask, shape=image.shape)
        for i in range(args.aug_nums):
            image_aug, mask_aug = seq(image=image, segmentation_maps=mask)
            cv2.imwrite(os.path.join(augmentation_output_dir, 'image',
                                     base + '[AUG][{}][IMAGE].png'.format(i)), image_aug)
            cv2.imwrite(os.path.join(augmentation_output_dir, 'mask',
                                     base + '[AUG][{}][MASK].png'.format(i)), mask_aug.get_arr())


if __name__ == '__main__':
    args = parse_args()
    seg_augmentation(args, 'train')
    seg_augmentation(args, 'test')