import numpy as np
from Augmentor.Operations import Skew, Distort, Rotate, Shear, Flip, Zoom
from PIL import Image
import cv2

from utils.augmentation.Cloner import Clone
from utils.augmentation.Colorizer import Colorize
from utils.augmentation.Skitcher import Skitch

import random


def do_operation(opt, image, padding=50):
    image_cv = cv2.cvtColor((image * 255).astype(np.uint8), cv2.IMREAD_COLOR)  # [..., ::-1]
    dim = image_cv.shape[:2]

    if str(opt) == 'Skew':
        color = [0, 0, 0]
        top, bottom = padding, padding
        left, right = padding, padding
        image_cv = cv2.copyMakeBorder(image_cv, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)

    image = opt.perform_operation([Image.fromarray(image_cv)])[0]

    return np.array(image.resize(dim)) / 255.0


operations = {'skew': lambda: Skew(probability=1, skew_type="RANDOM", magnitude=0.7),
              'distort': lambda: Distort(probability=1, grid_width=random.randint(1, 50), grid_height=random.randint(1, 50),
                                 magnitude=5),
              'rotate': lambda: Rotate(probability=1, rotation=random.randint(1, 360)),
              'shear': lambda: Shear(probability=1, max_shear_left=0, max_shear_right=random.randint(5, 15)) \
                  if random.randint(0, 1) == 1 else Shear(probability=1, max_shear_left=random.randint(5, 15),
                                                          max_shear_right=0),
              'zoom': lambda: Zoom(probability=1, min_factor=random.randint(2, 10) / 10,
                              max_factor=random.randint(10, 12) / 10),
              'colorize': lambda: Colorize(probability=1),
              'skitch': lambda: Skitch(probability=1),
              'flip': lambda: Flip(probability=1, top_bottom_left_right="RANDOM"),
              'clone': lambda: Clone(probability=1)
              }
