import cv2
import numpy as np

try:
    import scipy
    # scipy.ndimage cannot be accessed until explicitly imported
    from scipy import ndimage
except ImportError:
    scipy = None


def flip_axis(x, axis):
    x = np.asarray(x).swapaxes(axis, 0)
    x = x[::-1, ...]
    x = x.swapaxes(0, axis)
    return x


def random_rotation(x, rg, row_axis=0, col_axis=1, channel_axis=2,
                    fill_mode='nearest', cval=0., interpolation_order=1):
    """Performs a random rotation of a Numpy image tensor.
    # Arguments
        x: Input tensor. Must be 3D.
        rg: Rotation range, in degrees.
        row_axis: Index of axis for rows in the input tensor.
        col_axis: Index of axis for columns in the input tensor.
        channel_axis: Index of axis for channels in the input tensor.
        fill_mode: Points outside the boundaries of the input
            are filled according to the given mode
            (one of `{'constant', 'nearest', 'reflect', 'wrap'}`).
        cval: Value used for points outside the boundaries
            of the input if `mode='constant'`.
        interpolation_order int: order of spline interpolation.
            see `ndimage.interpolation.affine_transform`
    # Returns
        Rotated Numpy image tensor.
    """
    theta = np.random.uniform(-rg, rg)
    x = apply_affine_transform(x, theta=theta, channel_axis=channel_axis,
                               fill_mode=fill_mode, cval=cval,
                               order=interpolation_order)
    return x


def random_shift(x, wrg, hrg, row_axis=0, col_axis=1, channel_axis=2,
                 fill_mode='nearest', cval=0., interpolation_order=1):
    """Performs a random spatial shift of a Numpy image tensor.
    # Arguments
        x: Input tensor. Must be 3D.
        wrg: Width shift range, as a float fraction of the width.
        hrg: Height shift range, as a float fraction of the height.
        row_axis: Index of axis for rows in the input tensor.
        col_axis: Index of axis for columns in the input tensor.
        channel_axis: Index of axis for channels in the input tensor.
        fill_mode: Points outside the boundaries of the input
            are filled according to the given mode
            (one of `{'constant', 'nearest', 'reflect', 'wrap'}`).
        cval: Value used for points outside the boundaries
            of the input if `mode='constant'`.
        interpolation_order int: order of spline interpolation.
            see `ndimage.interpolation.affine_transform`
    # Returns
        Shifted Numpy image tensor.
    """
    h, w = x.shape[row_axis], x.shape[col_axis]
    tx = np.random.uniform(-hrg, hrg) * h
    ty = np.random.uniform(-wrg, wrg) * w
    x = apply_affine_transform(x, tx=tx, ty=ty, channel_axis=channel_axis,
                               fill_mode=fill_mode, cval=cval,
                               order=interpolation_order)
    return x


def random_shear(x, intensity, row_axis=0, col_axis=1, channel_axis=2,
                 fill_mode='nearest', cval=0., interpolation_order=1):
    """Performs a random spatial shear of a Numpy image tensor.
    # Arguments
        x: Input tensor. Must be 3D.
        intensity: Transformation intensity in degrees.
        row_axis: Index of axis for rows in the input tensor.
        col_axis: Index of axis for columns in the input tensor.
        channel_axis: Index of axis for channels in the input tensor.
        fill_mode: Points outside the boundaries of the input
            are filled according to the given mode
            (one of `{'constant', 'nearest', 'reflect', 'wrap'}`).
        cval: Value used for points outside the boundaries
            of the input if `mode='constant'`.
        interpolation_order int: order of spline interpolation.
            see `ndimage.interpolation.affine_transform`
    # Returns
        Sheared Numpy image tensor.
    """
    shear = np.random.uniform(-intensity, intensity)
    x = apply_affine_transform(x, shear=shear, channel_axis=channel_axis,
                               fill_mode=fill_mode, cval=cval,
                               order=interpolation_order)
    return x


def random_zoom(x, zoom_range, row_axis=0, col_axis=1, channel_axis=2,
                fill_mode='nearest', cval=0., interpolation_order=1):
    """Performs a random spatial zoom of a Numpy image tensor.
    # Arguments
        x: Input tensor. Must be 3D.
        zoom_range: Tuple of floats; zoom range for width and height.
        row_axis: Index of axis for rows in the input tensor.
        col_axis: Index of axis for columns in the input tensor.
        channel_axis: Index of axis for channels in the input tensor.
        fill_mode: Points outside the boundaries of the input
            are filled according to the given mode
            (one of `{'constant', 'nearest', 'reflect', 'wrap'}`).
        cval: Value used for points outside the boundaries
            of the input if `mode='constant'`.
        interpolation_order int: order of spline interpolation.
            see `ndimage.interpolation.affine_transform`
    # Returns
        Zoomed Numpy image tensor.
    # Raises
        ValueError: if `zoom_range` isn't a tuple.
    """
    if len(zoom_range) != 2:
        raise ValueError('`zoom_range` should be a tuple or list of two'
                         ' floats. Received: %s' % (zoom_range,))

    if zoom_range[0] == 1 and zoom_range[1] == 1:
        zx, zy = 1, 1
    else:
        zx, zy = np.random.uniform(zoom_range[0], zoom_range[1], 2)
    x = apply_affine_transform(x, zx=zx, zy=zy, channel_axis=channel_axis,
                               fill_mode=fill_mode, cval=cval,
                               order=interpolation_order)
    return x


def random_channel_shift(x, intensity, channel_axis=0):
    x = np.rollaxis(x, channel_axis, 0)
    min_x, max_x = np.min(x), np.max(x)
    channel_images = [np.clip(x_channel + np.random.uniform(-intensity, intensity), min_x, max_x)
                      for x_channel in x]
    x = np.stack(channel_images, axis=0)
    x = np.rollaxis(x, 0, channel_axis + 1)
    return x


# For curving soybean pods. L.C.Uzal
def random_curves_transform(x, strength=0.1, range=(0.,255.)):
    low, high = range
    delta = (high - low) * strength / 2.
    xp = np.random.uniform(low=low + delta, high=high - delta)
    yp = np.random.uniform(low=xp-delta, high=xp+delta)
    xp = np.asarray([low, xp, high])
    yp = np.asarray([low, yp, high])
    return np.interp(x,xp,yp)


def transform_matrix_offset_center(matrix, x, y):
    o_x = float(x) / 2 + 0.5
    o_y = float(y) / 2 + 0.5
    offset_matrix = np.array([[1, 0, o_x], [0, 1, o_y], [0, 0, 1]])
    reset_matrix = np.array([[1, 0, -o_x], [0, 1, -o_y], [0, 0, 1]])
    transform_matrix = np.dot(np.dot(offset_matrix, matrix), reset_matrix)
    return transform_matrix


def apply_affine_transform(x, theta=0, tx=0, ty=0, shear=0, zx=1, zy=1,
                           row_axis=0, col_axis=1, channel_axis=2,
                           fill_mode='nearest', cval=0., order=1):
    """Applies an affine transformation specified by the parameters given.
    # Arguments
        x: 2D numpy array, single image.
        theta: Rotation angle in degrees.
        tx: Width shift.
        ty: Heigh shift.
        shear: Shear angle in degrees.
        zx: Zoom in x direction.
        zy: Zoom in y direction
        row_axis: Index of axis for rows in the input image.
        col_axis: Index of axis for columns in the input image.
        channel_axis: Index of axis for channels in the input image.
        fill_mode: Points outside the boundaries of the input
            are filled according to the given mode
            (one of `{'constant', 'nearest', 'reflect', 'wrap'}`).
        cval: Value used for points outside the boundaries
            of the input if `mode='constant'`.
        order int: order of interpolation
    # Returns
        The transformed version of the input.
    """
    if scipy is None:
        raise ImportError('Image transformations require SciPy. '
                          'Install SciPy.')
    transform_matrix = None
    if theta != 0:
        theta = np.deg2rad(theta)
        rotation_matrix = np.array([[np.cos(theta), -np.sin(theta), 0],
                                    [np.sin(theta), np.cos(theta), 0],
                                    [0, 0, 1]])
        transform_matrix = rotation_matrix

    if tx != 0 or ty != 0:
        shift_matrix = np.array([[1, 0, tx],
                                 [0, 1, ty],
                                 [0, 0, 1]])
        if transform_matrix is None:
            transform_matrix = shift_matrix
        else:
            transform_matrix = np.dot(transform_matrix, shift_matrix)

    if shear != 0:
        shear = np.deg2rad(shear)
        shear_matrix = np.array([[1, -np.sin(shear), 0],
                                 [0, np.cos(shear), 0],
                                 [0, 0, 1]])
        if transform_matrix is None:
            transform_matrix = shear_matrix
        else:
            transform_matrix = np.dot(transform_matrix, shear_matrix)

    if zx != 1 or zy != 1:
        zoom_matrix = np.array([[zx, 0, 0],
                                [0, zy, 0],
                                [0, 0, 1]])
        if transform_matrix is None:
            transform_matrix = zoom_matrix
        else:
            transform_matrix = np.dot(transform_matrix, zoom_matrix)

    if transform_matrix is not None:
        h, w = x.shape[row_axis], x.shape[col_axis]
        transform_matrix = transform_matrix_offset_center(
            transform_matrix, h, w)
        x = np.rollaxis(x, channel_axis, 0)
        final_affine_matrix = transform_matrix[:2, :2]
        final_offset = transform_matrix[:2, 2]

        channel_images = [ndimage.interpolation.affine_transform(
            x_channel,
            final_affine_matrix,
            final_offset,
            order=order,
            mode=fill_mode,
            cval=cval) for x_channel in x]
        x = np.stack(channel_images, axis=0)
        x = np.rollaxis(x, 0, channel_axis + 1)
    return x


def random_transform(x, rotation_range=0,
                     width_shift_range=0.,
                     height_shift_range=0.,
                     shear_range=0.,
                     zoom_range=0.,
                     channel_shift_range=0.,
                     horizontal_flip=False,
                     vertical_flip=False,
                     random_curves_strength=0.):
    # Generate params
    if rotation_range:
        theta = np.random.uniform(-rotation_range, rotation_range)
    else:
        theta = 0
    
    h, w = x.shape[0], x.shape[1]
    if height_shift_range:
        tx = np.random.uniform(-height_shift_range, height_shift_range) * h
    else:
        tx = 0
    if width_shift_range:
        ty = np.random.uniform(-width_shift_range, width_shift_range) * w
    else:
        ty = 0

    if shear_range:
        shear = np.random.uniform(-shear_range, shear_range)
    else:
        shear = 0

    if np.isscalar(zoom_range):
        zoom_range = [1 - zoom_range, 1 + zoom_range]
    elif len(zoom_range) == 2:
        zoom_range = [zoom_range[0], zoom_range[1]]
    else:
        raise ValueError('`zoom_range` should be a float or '
                         'a tuple or list of two floats. '
                         'Received arg: ', zoom_range)
    if zoom_range[0] == 1 and zoom_range[1] == 1:
        zx, zy = 1, 1
    else:
        zx, zy = np.random.uniform(zoom_range[0], zoom_range[1], 2)

    # Apply transforms
    x = apply_affine_transform(x,
                               theta,
                               tx, ty,
                               shear,
                               zx, zy)

    if channel_shift_range != 0:
        x = random_channel_shift(x, channel_shift_range, 2)

    if horizontal_flip:
        if np.random.random() < 0.5:
            x = flip_axis(x, 1)

    if vertical_flip:
        if np.random.random() < 0.5:
            x = flip_axis(x, 0)

    if random_curves_strength > 0.:
        x = random_curves_transform(x, random_curves_strength)

    return x


if __name__ == "__main__":
    import argparse
    from PIL import Image
    parser = argparse.ArgumentParser()
    parser.add_argument('--image', default='', type=str)
    parser.add_argument('--imageOut', default='result.png', type=str)
    args = parser.parse_args()

    im = np.array(Image.open(args.image))
    img = random_transform(im, rotation_range=10, shear_range=.5, zoom_range=.2, channel_shift_range=10., horizontal_flip=True)
    Image.fromarray(np.uint8(img)).save(args.imageOut)

[(38, 'numpy.random.uniform', 'np.random.uniform', True, 'import numpy as np\n'), (93, 'numpy.random.uniform', 'np.random.uniform', True, 'import numpy as np\n'), (136, 'numpy.rollaxis', 'np.rollaxis', True, 'import numpy as np\n'), (140, 'numpy.stack', 'np.stack', True, 'import numpy as np\n'), (141, 'numpy.rollaxis', 'np.rollaxis', True, 'import numpy as np\n'), (149, 'numpy.random.uniform', 'np.random.uniform', True, 'import numpy as np\n'), (150, 'numpy.random.uniform', 'np.random.uniform', True, 'import numpy as np\n'), (151, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (152, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (153, 'numpy.interp', 'np.interp', True, 'import numpy as np\n'), (159, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (160, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (278, 'numpy.isscalar', 'np.isscalar', True, 'import numpy as np\n'), (318, 'argparse.ArgumentParser', 'argparse.ArgumentParser', False, 'import argparse\n'), (66, 'numpy.random.uniform', 'np.random.uniform', True, 'import numpy as np\n'), (67, 'numpy.random.uniform', 'np.random.uniform', True, 'import numpy as np\n'), (128, 'numpy.random.uniform', 'np.random.uniform', True, 'import numpy as np\n'), (137, 'numpy.min', 'np.min', True, 'import numpy as np\n'), (137, 'numpy.max', 'np.max', True, 'import numpy as np\n'), (161, 'numpy.dot', 'np.dot', True, 'import numpy as np\n'), (194, 'numpy.deg2rad', 'np.deg2rad', True, 'import numpy as np\n'), (201, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (210, 'numpy.deg2rad', 'np.deg2rad', True, 'import numpy as np\n'), (220, 'numpy.array', 'np.array', True, 'import numpy as np\n'), (232, 'numpy.rollaxis', 'np.rollaxis', True, 'import numpy as np\n'), (243, 'numpy.stack', 'np.stack', True, 'import numpy as np\n'), (244, 'numpy.rollaxis', 'np.rollaxis', True, 'import numpy as np\n'), (259, 'numpy.random.uniform', 'np.random.uniform', True, 'import numpy as np\n'), (274, 'numpy.random.uniform', 'np.random.uniform', True, 'import numpy as np\n'), (289, 'numpy.random.uniform', 'np.random.uniform', True, 'import numpy as np\n'), (323, 'PIL.Image.open', 'Image.open', False, 'from PIL import Image\n'), (13, 'numpy.asarray', 'np.asarray', True, 'import numpy as np\n'), (207, 'numpy.dot', 'np.dot', True, 'import numpy as np\n'), (217, 'numpy.dot', 'np.dot', True, 'import numpy as np\n'), (226, 'numpy.dot', 'np.dot', True, 'import numpy as np\n'), (236, 'scipy.ndimage.interpolation.affine_transform', 'ndimage.interpolation.affine_transform', False, 'from scipy import ndimage\n'), (265, 'numpy.random.uniform', 'np.random.uniform', True, 'import numpy as np\n'), (269, 'numpy.random.uniform', 'np.random.uniform', True, 'import numpy as np\n'), (302, 'numpy.random.random', 'np.random.random', True, 'import numpy as np\n'), (306, 'numpy.random.random', 'np.random.random', True, 'import numpy as np\n'), (138, 'numpy.random.uniform', 'np.random.uniform', True, 'import numpy as np\n'), (325, 'numpy.uint8', 'np.uint8', True, 'import numpy as np\n'), (195, 'numpy.cos', 'np.cos', True, 'import numpy as np\n'), (196, 'numpy.sin', 'np.sin', True, 'import numpy as np\n'), (196, 'numpy.cos', 'np.cos', True, 'import numpy as np\n'), (212, 'numpy.cos', 'np.cos', True, 'import numpy as np\n'), (195, 'numpy.sin', 'np.sin', True, 'import numpy as np\n'), (211, 'numpy.sin', 'np.sin', True, 'import numpy as np\n')]