"""
DISCLAIMER : This code has been heavily borrowed from the Keras source code.
Fairly basic set of tools for real-time data augmentation on image data.
Can easily be extended to include new transformations,
new preprocessing methods, etc...
"""
from __future__ import absolute_import
from __future__ import print_function

import numpy as np
from keras.preprocessing.image import apply_affine_transform
from keras_preprocessing.image import flip_axis

try:
    from PIL import Image as pil_image
except ImportError:
    pil_image = None


def random_rotation(x, rg, row_axis=1, col_axis=2, channel_axis=0,
                    fill_mode='nearest', cval=0., theta=None):
    """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'`.
        theta: Value to disable randomness or None.

    # Returns
        Rotated Numpy image tensor.
    """
    theta = np.pi / 180 * np.random.uniform(-rg, rg) if theta is None else theta
    x = apply_affine_transform(x, theta=theta, channel_axis=channel_axis, fill_mode=fill_mode, cval=cval, row_axis=row_axis, col_axis=col_axis)
    return x


def random_shift(x, wrg, hrg, row_axis=1, col_axis=2, channel_axis=0,
                 fill_mode='nearest', cval=0., tx=None, ty=None):
    """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'`.
        tx : Value to disable randomness in X or None.
        ty : Value to disable randomness in Y or None.

    # Returns
        Shifted Numpy image tensor.
    """
    h, w = x.shape[row_axis], x.shape[col_axis]
    tx = np.random.uniform(-hrg, hrg) * h if tx is None else tx
    ty = np.random.uniform(-wrg, wrg) * w if ty is None else ty

    tx *= h
    ty *= w
    x = apply_affine_transform(x, tx=tx, ty=ty, channel_axis=channel_axis, fill_mode=fill_mode, cval=cval)
    return x


def random_shear(x, intensity, row_axis=1, col_axis=2, channel_axis=0,
                 fill_mode='nearest', cval=0., known_intensity=None):
    """Performs a random spatial shear of a Numpy image tensor.

    # Arguments
        x: Input tensor. Must be 3D.
        intensity: Transformation intensity.
        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'`.
        known_intensity: Value to disable randomness or None.

    # Returns
        Sheared Numpy image tensor.
    """
    shear = np.random.uniform(-intensity, intensity) if known_intensity is None else known_intensity
    x = apply_affine_transform(x, shear=shear, channel_axis=channel_axis, fill_mode=fill_mode, cval=cval, row_axis=row_axis, col_axis=col_axis)
    return x


def random_zoom(x, zoom_range, row_axis=1, col_axis=2, channel_axis=0,
                fill_mode='nearest', cval=0., z_known=None):
    """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'`.
        z_known: Value to disable randomness or None.

    # Returns
        Zoomed Numpy image tensor.

    # Raises
        ValueError: if `zoom_range` isn't a tuple.
    """
    if z_known is None:
        if len(zoom_range) != 2:
            raise ValueError('`zoom_range` should be 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)
    else:
        zx, zy = z_known
    x = apply_affine_transform(x, zx=zx, zy=zy, channel_axis=channel_axis, fill_mode=fill_mode, cval=cval, row_axis=row_axis, col_axis=col_axis)
    return x


def random_channel_shift(x, intensity, channel_axis=0, known_intensity=None):
    x = np.rollaxis(x, channel_axis, 0)
    known_intensity = np.random.uniform(-intensity, intensity) if known_intensity is None else known_intensity
    min_x, max_x = np.min(x), np.max(x)
    channel_images = [np.clip(x_channel + known_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


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 flip_horizontal(x, value, col_axis=2):
    if value < 0.5:
        return flip_axis(x, col_axis)
    else:
        return x


def flip_vertical(x, value, row_axis=1):
    if value < 0.5:
        return flip_axis(x, row_axis)
    else:
        return x

"""
DISCLAIMER : This code has been heavily borrowed from the Keras source code.
Fairly basic set of tools for real-time data augmentation on image data.
Can easily be extended to include new transformations,
new preprocessing methods, etc...
"""
from __future__ import absolute_import
from __future__ import print_function

import numpy as np
from keras.preprocessing.image import apply_affine_transform
from keras_preprocessing.image import flip_axis

try:
    from PIL import Image as pil_image
except ImportError:
    pil_image = None


def random_rotation(x, rg, row_axis=1, col_axis=2, channel_axis=0,
                    fill_mode='nearest', cval=0., theta=None):
    """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'`.
        theta: Value to disable randomness or None.

    # Returns
        Rotated Numpy image tensor.
    """
    theta = np.pi / 180 * np.random.uniform(-rg, rg) if theta is None else theta
    x = apply_affine_transform(x, theta=theta, channel_axis=channel_axis, fill_mode=fill_mode, cval=cval, row_axis=row_axis, col_axis=col_axis)
    return x


def random_shift(x, wrg, hrg, row_axis=1, col_axis=2, channel_axis=0,
                 fill_mode='nearest', cval=0., tx=None, ty=None):
    """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'`.
        tx : Value to disable randomness in X or None.
        ty : Value to disable randomness in Y or None.

    # Returns
        Shifted Numpy image tensor.
    """
    h, w = x.shape[row_axis], x.shape[col_axis]
    tx = np.random.uniform(-hrg, hrg) * h if tx is None else tx
    ty = np.random.uniform(-wrg, wrg) * w if ty is None else ty

    tx *= h
    ty *= w
    x = apply_affine_transform(x, tx=tx, ty=ty, channel_axis=channel_axis, fill_mode=fill_mode, cval=cval)
    return x


def random_shear(x, intensity, row_axis=1, col_axis=2, channel_axis=0,
                 fill_mode='nearest', cval=0., known_intensity=None):
    """Performs a random spatial shear of a Numpy image tensor.

    # Arguments
        x: Input tensor. Must be 3D.
        intensity: Transformation intensity.
        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'`.
        known_intensity: Value to disable randomness or None.

    # Returns
        Sheared Numpy image tensor.
    """
    shear = np.random.uniform(-intensity, intensity) if known_intensity is None else known_intensity
    x = apply_affine_transform(x, shear=shear, channel_axis=channel_axis, fill_mode=fill_mode, cval=cval, row_axis=row_axis, col_axis=col_axis)
    return x


def random_zoom(x, zoom_range, row_axis=1, col_axis=2, channel_axis=0,
                fill_mode='nearest', cval=0., z_known=None):
    """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'`.
        z_known: Value to disable randomness or None.

    # Returns
        Zoomed Numpy image tensor.

    # Raises
        ValueError: if `zoom_range` isn't a tuple.
    """
    if z_known is None:
        if len(zoom_range) != 2:
            raise ValueError('`zoom_range` should be 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)
    else:
        zx, zy = z_known
    x = apply_affine_transform(x, zx=zx, zy=zy, channel_axis=channel_axis, fill_mode=fill_mode, cval=cval, row_axis=row_axis, col_axis=col_axis)
    return x


def random_channel_shift(x, intensity, channel_axis=0, known_intensity=None):
    x = np.rollaxis(x, channel_axis, 0)
    known_intensity = np.random.uniform(-intensity, intensity) if known_intensity is None else known_intensity
    min_x, max_x = np.min(x), np.max(x)
    channel_images = [np.clip(x_channel + known_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


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 flip_horizontal(x, value, col_axis=2):
    if value < 0.5:
        return flip_axis(x, col_axis)
    else:
        return x


def flip_vertical(x, value, row_axis=1):
    if value < 0.5:
        return flip_axis(x, row_axis)
    else:
        return x

DISCLAIMER : This code has been heavily borrowed from the Keras source code. Fairly basic set of tools for real-time data augmentation on image data. Can easily be extended to include new transformations, new preprocessing methods, etc... 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'`.         theta: Value to disable randomness or None.      # Returns         Rotated Numpy image tensor. 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'`.         tx : Value to disable randomness in X or None.         ty : Value to disable randomness in Y or None.      # Returns         Shifted Numpy image tensor. Performs a random spatial shear of a Numpy image tensor.      # Arguments         x: Input tensor. Must be 3D.         intensity: Transformation intensity.         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'`.         known_intensity: Value to disable randomness or None.      # Returns         Sheared Numpy image tensor. 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'`.         z_known: Value to disable randomness or None.      # Returns         Zoomed Numpy image tensor.      # Raises         ValueError: if `zoom_range` isn't a tuple.