import os
import math as m
import numpy as np
# np.random.seed(123)
import scipy.io
from sklearn.decomposition import PCA
from sklearn.preprocessing import scale

import random
import keras.backend as K
import warnings
import keras

from keras.preprocessing.sequence import pad_sequences

# import all image transformation utility functions
from .util import random_rotation, random_shift,\
    random_shear, random_zoom, random_channel_shift,\
    transform_matrix_offset_center, apply_transform, flip_axis
from .basegen import Iterator


class DataDirGenerator(object):
    """Generate minibatches of image data with real-time data augmentation.

    Modified from Keras ImageDataGenerator

    # Arguments
        featurewise_center: set input mean to 0 over the dataset.
        samplewise_center: set each sample mean to 0.
        featurewise_std_normalization: divide inputs by std of the dataset.
        samplewise_std_normalization: divide each input by its std.
        rotation_range: degrees (0 to 180).
        width_shift_range: fraction of total width, if < 1, or pixels if >= 1.
        height_shift_range: fraction of total height, if < 1, or pixels if >= 1.
        shear_range: shear intensity (shear angle in degrees).
        zoom_range: amount of zoom. if scalar z, zoom will be randomly picked
            in the range [1-z, 1+z]. A sequence of two can be passed instead
            to select this range.
        channel_shift_range: shift range for each channel.
        fill_mode: points outside the boundaries are filled according to the
            given mode ('constant', 'nearest', 'reflect' or 'wrap'). Default
            is 'nearest'.
            Points outside the boundaries of the input are filled according to the given mode:
                'constant': kkkkkkkk|abcd|kkkkkkkk (cval=k)
                'nearest':  aaaaaaaa|abcd|dddddddd
                'reflect':  abcddcba|abcd|dcbaabcd
                'wrap':  abcdabcd|abcd|abcdabcd
        horizontal_flip: whether to randomly flip images horizontally.
        vertical_flip: whether to randomly flip images vertically.
        rescale: rescaling factor. If None or 0, no rescaling is applied,
            otherwise we multiply the data by the value provided. This is
            applied after the `preprocessing_function` (if any provided)
            but before any other transformation.
        preprocessing_function: function that will be implied on each input.
            The function will run before any other modification on it.
            The function should take one argument:
            one image (Numpy tensor with rank 3),
            and should output a Numpy tensor with the same shape.
        data_format: 'channels_first' or 'channels_last'. In 'channels_first' mode, the channels dimension
            (the depth) is at index 1, in 'channels_last' mode it is at index 3.
            It defaults to the `image_data_format` value found in your
            Keras config file at `~/.keras/keras.json`.
            If you never set it, then it will be "channels_last".
    """

    def __init__(self,
                 featurewise_center=False,
                 samplewise_center=False,
                 featurewise_std_normalization=False,
                 samplewise_std_normalization=False,
                 rotation_range=0.,
                 width_shift_range=0.,
                 height_shift_range=0.,
                 shear_range=0.,
                 zoom_range=0.,
                 channel_shift_range=0.,
                 fill_mode='nearest',
                 cval=0.,
                 horizontal_flip=False,
                 vertical_flip=False,
                 rescale=None,
                 data_format=None,
                 preprocessing_function=None):

        if data_format is None:
            data_format = K.image_data_format()
        self.featurewise_center = featurewise_center
        self.samplewise_center = samplewise_center
        self.featurewise_std_normalization = featurewise_std_normalization
        self.samplewise_std_normalization = samplewise_std_normalization
        self.rotation_range = rotation_range
        self.width_shift_range = width_shift_range
        self.height_shift_range = height_shift_range
        self.shear_range = shear_range
        self.zoom_range = zoom_range
        self.channel_shift_range = channel_shift_range
        self.fill_mode = fill_mode
        self.cval = cval
        self.horizontal_flip = horizontal_flip
        self.vertical_flip = vertical_flip
        self.rescale = rescale

        if data_format not in {'channels_last', 'channels_first'}:
            raise ValueError('`data_format` should be `"channels_last"` (channel after row and '
                             'column) or `"channels_first"` (channel before row and column). '
                             'Received arg: ', data_format)
        self.data_format = data_format
        if data_format == 'channels_first':
            self.channel_axis = 1
            self.row_axis = 2
            self.col_axis = 3
        if data_format == 'channels_last':
            self.channel_axis = 3
            self.row_axis = 1
            self.col_axis = 2

        # img_row_axis = self.row_axis - 1
        # img_col_axis = self.col_axis - 1
        # img_channel_axis = self.channel_axis - 1

        # initialize mean and std and pc for random augmentations
        self.mean = None
        self.std = None
        self.principal_components = None

        if np.isscalar(zoom_range):
            self.zoom_range = [1 - zoom_range, 1 + zoom_range]
        elif len(zoom_range) == 2:
            self.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 featurewise_std_normalization:
            if not featurewise_center:
                self.featurewise_center = True
                warnings.warn('This ImageDataGenerator specifies '
                              '`featurewise_std_normalization`, '
                              'which overrides setting of '
                              '`featurewise_center`.')
        if samplewise_std_normalization:
            if not samplewise_center:
                self.samplewise_center = True
                warnings.warn('This ImageDataGenerator specifies '
                              '`samplewise_std_normalization`, '
                              'which overrides setting of '
                              '`samplewise_center`.')

    def flow_from_directory(self,
                            directory,
                            batch_size=32,
                            num_timesteps=5,
                            num_classes=2,
                            numchannels=3,
                            imsize=32,
                            shuffle=True,
                            seed=None,
                            interpolation='nearest'):
        ''' Returns an iterator through the files we have '''
        return DirectoryIterator(
            directory, self,
            batch_size=batch_size,
            num_timesteps=num_timesteps,
            num_classes=num_classes,
            numchannels=numchannels,
            imsize=imsize,
            shuffle=shuffle,
            seed=seed,
            data_format=self.data_format,
            interpolation=interpolation)

    def random_transform(self, x, seed=None):
        """Randomly augment a single image tensor.

        # Arguments
            x: 3D tensor, single image.
            seed: random seed.

        # Returns
            A randomly transformed version of the input (same shape).
        """
        # x is a single image, so it doesn't have image number at index 0
        img_row_axis = self.row_axis - 1
        img_col_axis = self.col_axis - 1
        img_channel_axis = self.channel_axis - 1

        if seed is not None:
            np.random.seed(seed)

        # use composition of homographies
        # to generate final transform that needs to be applied
        if self.rotation_range:
            theta = np.deg2rad(
                np.random.uniform(-self.rotation_range, self.rotation_range))
        else:
            theta = 0

        if self.height_shift_range:
            tx = np.random.uniform(-self.height_shift_range,
                                   self.height_shift_range)
            if self.height_shift_range < 1:
                tx *= x.shape[img_row_axis]
        else:
            tx = 0

        if self.width_shift_range:
            ty = np.random.uniform(-self.width_shift_range,
                                   self.width_shift_range)
            if self.width_shift_range < 1:
                ty *= x.shape[img_col_axis]
        else:
            ty = 0

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

        if self.zoom_range[0] == 1 and self.zoom_range[1] == 1:
            zx, zy = 1, 1
        else:
            zx, zy = np.random.uniform(
                self.zoom_range[0], self.zoom_range[1], 2)

        transform_matrix = None
        if theta != 0:
            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]])
            transform_matrix = shift_matrix if transform_matrix is None else np.dot(
                transform_matrix, shift_matrix)

        if shear != 0:
            shear_matrix = np.array([[1, -np.sin(shear), 0],
                                     [0, np.cos(shear), 0],
                                     [0, 0, 1]])
            transform_matrix = shear_matrix if transform_matrix is None else 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]])
            transform_matrix = zoom_matrix if transform_matrix is None else np.dot(
                transform_matrix, zoom_matrix)

        if transform_matrix is not None:
            h, w = x.shape[img_row_axis], x.shape[img_col_axis]
            transform_matrix = transform_matrix_offset_center(
                transform_matrix, h, w)
            x = apply_transform(x, transform_matrix, img_channel_axis,
                                fill_mode=self.fill_mode, cval=self.cval)

        if self.channel_shift_range != 0:
            x = random_channel_shift(x,
                                     self.channel_shift_range,
                                     img_channel_axis)
        if self.horizontal_flip:
            if np.random.random() < 0.5:
                x = flip_axis(x, img_col_axis)

        if self.vertical_flip:
            if np.random.random() < 0.5:
                x = flip_axis(x, img_row_axis)
        return x

    def fit(self, x,
            augment=False,
            rounds=1,
            seed=None):
        """Fits internal statistics to some sample data.

        Required for featurewise_center, featurewise_std_normalization
        and zca_whitening.

        # Arguments
            x: Numpy array, the data to fit on. Should have rank 4.
                In case of grayscale data,
                the channels axis should have value 1, and in case
                of RGB data, it should have value 3.
            augment: Whether to fit on randomly augmented samples
            rounds: If `augment`,
                how many augmentation passes to do over the data
            seed: random seed.

        # Raises
            ValueError: in case of invalid input `x`.
        """
        x = np.asarray(x, dtype=K.floatx())
        if x.ndim != 4:
            raise ValueError('Input to `.fit()` should have rank 4. '
                             'Got array with shape: ' + str(x.shape))
        if x.shape[self.channel_axis] not in {1, 3, 4}:
            warnings.warn(
                'Expected input to be images (as Numpy array) '
                '(channels on axis ' + str(self.channel_axis) + '), i.e. expected '
                'either 1, 3 or 4 channels on axis ' +
                str(self.channel_axis) + '. '
                'However, it was passed an array with shape ' + str(x.shape) +
                ' (' + str(x.shape[self.channel_axis]) + ' channels).')

        if seed is not None:
            np.random.seed(seed)

        x = np.copy(x)
        if augment:
            ax = np.zeros(tuple([rounds * x.shape[0]] +
                                list(x.shape)[1:]), dtype=K.floatx())
            for r in range(rounds):
                for i in range(x.shape[0]):
                    ax[i + r * x.shape[0]] = self.random_transform(x[i])
            x = ax

        if self.featurewise_center:
            self.mean = np.mean(x, axis=(0, self.row_axis, self.col_axis))
            broadcast_shape = [1, 1, 1]
            broadcast_shape[self.channel_axis - 1] = x.shape[self.channel_axis]
            self.mean = np.reshape(self.mean, broadcast_shape)
            x -= self.mean

        if self.featurewise_std_normalization:
            self.std = np.std(x, axis=(0, self.row_axis, self.col_axis))
            broadcast_shape = [1, 1, 1]
            broadcast_shape[self.channel_axis - 1] = x.shape[self.channel_axis]
            self.std = np.reshape(self.std, broadcast_shape)
            x /= (self.std + K.epsilon())

        if self.zca_whitening:
            flat_x = np.reshape(
                x, (x.shape[0], x.shape[1] * x.shape[2] * x.shape[3]))
            num_examples = flat_x.shape[0]
            _, s, vt = linalg.svd(flat_x / np.sqrt(num_examples))
            s_expand = np.hstack((s, np.zeros(vt.shape[0] - num_examples,
                                              dtype=flat_x.dtype)))
            self.principal_components = (
                vt.T / np.sqrt(s_expand ** 2 + self.zca_epsilon)).dot(vt)

    def standardize(self, x):
        """Apply the normalization configuration to a batch of inputs.
        # Arguments
            x: batch of inputs to be normalized.
        # Returns
            The inputs, normalized.
        """
        if self.samplewise_center:
            x -= np.mean(x, keepdims=True)
        if self.samplewise_std_normalization:
            x /= (np.std(x, keepdims=True) + K.epsilon())

        if self.featurewise_center:
            if self.mean is not None:
                x -= self.mean
            else:
                warnings.warn('This ImageDataGenerator specifies '
                              '`featurewise_center`, but it hasn\'t '
                              'been fit on any training data. Fit it '
                              'first by calling `.fit(numpy_data)`.')
        if self.featurewise_std_normalization:
            if self.std is not None:
                x /= (self.std + K.epsilon())
            else:
                warnings.warn('This ImageDataGenerator specifies '
                              '`featurewise_std_normalization`, but it hasn\'t '
                              'been fit on any training data. Fit it '
                              'first by calling `.fit(numpy_data)`.')
        return x


class DirectoryIterator(Iterator):
    """Iterator capable of reading images from a directory on disk.

    # Arguments
        directory: Path to the directory to read images from.
            Each subdirectory in this directory will be
            considered to contain images from one class,
            or alternatively you could specify class subdirectories
            via the `classes` argument.
        image_data_generator: Instance of `ImageDataGenerator`
            to use for random transformations and normalization.
        target_size: tuple of integers, dimensions to resize input images to.
        color_mode: One of `"rgb"`, `"grayscale"`. Color mode to read images.
        classes: Optional list of strings, names of subdirectories
            containing images from each class (e.g. `["dogs", "cats"]`).
            It will be computed automatically if not set.
        class_mode: Mode for yielding the targets:
            `"binary"`: binary targets (if there are only two classes),
            `"categorical"`: categorical targets,
            `"sparse"`: integer targets,
            `"input"`: targets are images identical to input images (mainly
                used to work with autoencoders),
            `None`: no targets get yielded (only input images are yielded).
        batch_size: Integer, size of a batch.
        shuffle: Boolean, whether to shuffle the data between epochs.
        seed: Random seed for data shuffling.
        data_format: String, one of `channels_first`, `channels_last`.
        save_to_dir: Optional directory where to save the pictures
            being yielded, in a viewable format. This is useful
            for visualizing the random transformations being
            applied, for debugging purposes.
        save_prefix: String prefix to use for saving sample
            images (if `save_to_dir` is set).
        save_format: Format to use for saving sample images
            (if `save_to_dir` is set).
        interpolation: Interpolation method used to resample the image if the
            target size is different from that of the loaded image.
            Supported methods are "nearest", "bilinear", and "bicubic".
            If PIL version 1.1.3 or newer is installed, "lanczos" is also
            supported. If PIL version 3.4.0 or newer is installed, "box" and
            "hamming" are also supported. By default, "nearest" is used.
    """

    def __init__(self,
                 directory,
                 image_data_generator,
                 batch_size=32,
                 num_timesteps=5,
                 num_classes=2,
                 numchannels=3,
                 imsize=32,
                 shuffle=True,
                 seed=None,
                 data_format=None,
                 interpolation='nearest'):
        if data_format is None:
            data_format = K.image_data_format()
        self.num_timesteps = num_timesteps
        self.data_format = data_format
        self.directory = directory
        self.image_data_generator = image_data_generator
        self.num_classes = num_classes
        self.numchannels = numchannels
        self.imsize = imsize

        # how to apply interpolation
        self.interpolation = interpolation

        # set the file paths of all the sequences we want to iterate over
        self.samples = 0
        self.filepaths = []
        for root, dirs, files in os.walk(self.directory):
            for file in files:
                self.filepaths.append(os.path.join(root, file))
        self.samples = len(self.filepaths)

        print('Found %d image files.' % (self.samples))
        super(DirectoryIterator, self).__init__(
            self.samples, batch_size, shuffle, seed)

    def _get_batches_of_transformed_sequences(self, index_array):
        # build batch of image data
        for i, index in enumerate(index_array):
            fpath = self.filepaths[index]

            # load in a certain file and randomly sample the image
            imgdata = np.load(fpath)
            images = imgdata['image_tensor']
            metadata = imgdata['metadata'].item()
            y = metadata['ylabels']

            # get the random index in this dataset
            numsamples = y.shape[0]

            # reformat data if necessary
            images, y = self._reformatdata(images, y)

            # get it as a sequence
            N = 1  # number of samples to get from this dataset | Todo change into higher number
            assert numsamples == images.shape[0]
            randind = random.sample(range(numsamples), k=N)
            randind = randind[0]
            sampimages, samplabels = self._sample_as_seq(images, y, randind)

            # apply normalization to each image
            for idx in range(0, len(sampimages)):
                sampimage = sampimages[idx, ...]
                # apply transforms and standardize
                sampimage = self.image_data_generator.random_transform(
                    sampimage)
                sampimage = self.image_data_generator.standardize(sampimage)

            if i == 0:
                # init batch of X and Y
                batch_x = np.zeros((len(index_array),) +
                                   sampimages.shape, dtype=K.floatx())
                batch_y = np.zeros((len(index_array),) +
                                   samplabels.shape, dtype=K.floatx())

            batch_x[i] = sampimages
            batch_y[i] = samplabels

        return batch_x, batch_y

    def _sample_as_seq(self, sampdata, samplabels, randind):
        '''
        Get sampled data as a sequence
        '''
        endlen = sampdata.shape[0]

        if randind + self.num_timesteps >= endlen:
            print("padding...")
            sampdata = sampdata[randind:, ...]
            sampdata = self._paddata(sampdata, padding='post')

            samplabels = samplabels[randind:, ...]
            samplabels = self._paddata(samplabels, padding='post')
        else:
            sampdata = sampdata[randind:randind+self.num_timesteps, ...]
            samplabels = samplabels[randind:randind+self.num_timesteps, ...]
        return sampdata, samplabels

    def _paddata(self, sampdata, padding):
        '''
        Pad a sample data on the 
        '''
        samplen = sampdata.shape[0]
        numsamps_topad = self.num_timesteps - samplen
        imshape = sampdata.shape[1:]
        assert imshape[0] == self.imsize
        assert imshape[-1] == self.numchannels

        # pad with zeros
        if samplen < self.num_timesteps:
            padmat = np.zeros(((samplen,)+imshape))
            if padding == 'pre':
                # pad on the pre
                sampdata = np.concatenate((padmat, sampdata), axis=0)
            elif padding == 'post':
                # pad on the pre
                sampdata = np.concatenate((sampdata, padmat), axis=0)
        return sampdata

    def _reformatdata(self, images, y):
        invert_y = 1 - y
        y = np.concatenate((y, invert_y), axis=1)

        # transform the images if necessary
        # reshape
        images = images.reshape(
            (-1, self.numchannels, self.imsize, self.imsize))
        images = images.swapaxes(1, 3)
        images = images.astype("float32")
        return images, y

    def next(self):
        """For python 2.x.

        # Returns
            The next batch.
        """
        with self.lock:
            index_array = next(self.index_generator)
        # The transformation of images is not under thread lock
        # so it can be done in parallel
        return self._get_batches_of_transformed_sequences(index_array)

import os
import math as m
import numpy as np
# np.random.seed(123)
import scipy.io
from sklearn.decomposition import PCA
from sklearn.preprocessing import scale

import random
import keras.backend as K
import warnings
import keras

from keras.preprocessing.sequence import pad_sequences

# import all image transformation utility functions
from .util import random_rotation, random_shift,\
    random_shear, random_zoom, random_channel_shift,\
    transform_matrix_offset_center, apply_transform, flip_axis
from .basegen import Iterator


class DataDirGenerator(object):
    """Generate minibatches of image data with real-time data augmentation.

    Modified from Keras ImageDataGenerator

    # Arguments
        featurewise_center: set input mean to 0 over the dataset.
        samplewise_center: set each sample mean to 0.
        featurewise_std_normalization: divide inputs by std of the dataset.
        samplewise_std_normalization: divide each input by its std.
        rotation_range: degrees (0 to 180).
        width_shift_range: fraction of total width, if < 1, or pixels if >= 1.
        height_shift_range: fraction of total height, if < 1, or pixels if >= 1.
        shear_range: shear intensity (shear angle in degrees).
        zoom_range: amount of zoom. if scalar z, zoom will be randomly picked
            in the range [1-z, 1+z]. A sequence of two can be passed instead
            to select this range.
        channel_shift_range: shift range for each channel.
        fill_mode: points outside the boundaries are filled according to the
            given mode ('constant', 'nearest', 'reflect' or 'wrap'). Default
            is 'nearest'.
            Points outside the boundaries of the input are filled according to the given mode:
                'constant': kkkkkkkk|abcd|kkkkkkkk (cval=k)
                'nearest':  aaaaaaaa|abcd|dddddddd
                'reflect':  abcddcba|abcd|dcbaabcd
                'wrap':  abcdabcd|abcd|abcdabcd
        horizontal_flip: whether to randomly flip images horizontally.
        vertical_flip: whether to randomly flip images vertically.
        rescale: rescaling factor. If None or 0, no rescaling is applied,
            otherwise we multiply the data by the value provided. This is
            applied after the `preprocessing_function` (if any provided)
            but before any other transformation.
        preprocessing_function: function that will be implied on each input.
            The function will run before any other modification on it.
            The function should take one argument:
            one image (Numpy tensor with rank 3),
            and should output a Numpy tensor with the same shape.
        data_format: 'channels_first' or 'channels_last'. In 'channels_first' mode, the channels dimension
            (the depth) is at index 1, in 'channels_last' mode it is at index 3.
            It defaults to the `image_data_format` value found in your
            Keras config file at `~/.keras/keras.json`.
            If you never set it, then it will be "channels_last".
    """

    def __init__(self,
                 featurewise_center=False,
                 samplewise_center=False,
                 featurewise_std_normalization=False,
                 samplewise_std_normalization=False,
                 rotation_range=0.,
                 width_shift_range=0.,
                 height_shift_range=0.,
                 shear_range=0.,
                 zoom_range=0.,
                 channel_shift_range=0.,
                 fill_mode='nearest',
                 cval=0.,
                 horizontal_flip=False,
                 vertical_flip=False,
                 rescale=None,
                 data_format=None,
                 preprocessing_function=None):

        if data_format is None:
            data_format = K.image_data_format()
        self.featurewise_center = featurewise_center
        self.samplewise_center = samplewise_center
        self.featurewise_std_normalization = featurewise_std_normalization
        self.samplewise_std_normalization = samplewise_std_normalization
        self.rotation_range = rotation_range
        self.width_shift_range = width_shift_range
        self.height_shift_range = height_shift_range
        self.shear_range = shear_range
        self.zoom_range = zoom_range
        self.channel_shift_range = channel_shift_range
        self.fill_mode = fill_mode
        self.cval = cval
        self.horizontal_flip = horizontal_flip
        self.vertical_flip = vertical_flip
        self.rescale = rescale

        if data_format not in {'channels_last', 'channels_first'}:
            raise ValueError('`data_format` should be `"channels_last"` (channel after row and '
                             'column) or `"channels_first"` (channel before row and column). '
                             'Received arg: ', data_format)
        self.data_format = data_format
        if data_format == 'channels_first':
            self.channel_axis = 1
            self.row_axis = 2
            self.col_axis = 3
        if data_format == 'channels_last':
            self.channel_axis = 3
            self.row_axis = 1
            self.col_axis = 2

        # img_row_axis = self.row_axis - 1
        # img_col_axis = self.col_axis - 1
        # img_channel_axis = self.channel_axis - 1

        # initialize mean and std and pc for random augmentations
        self.mean = None
        self.std = None
        self.principal_components = None

        if np.isscalar(zoom_range):
            self.zoom_range = [1 - zoom_range, 1 + zoom_range]
        elif len(zoom_range) == 2:
            self.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 featurewise_std_normalization:
            if not featurewise_center:
                self.featurewise_center = True
                warnings.warn('This ImageDataGenerator specifies '
                              '`featurewise_std_normalization`, '
                              'which overrides setting of '
                              '`featurewise_center`.')
        if samplewise_std_normalization:
            if not samplewise_center:
                self.samplewise_center = True
                warnings.warn('This ImageDataGenerator specifies '
                              '`samplewise_std_normalization`, '
                              'which overrides setting of '
                              '`samplewise_center`.')

    def flow_from_directory(self,
                            directory,
                            batch_size=32,
                            num_timesteps=5,
                            num_classes=2,
                            numchannels=3,
                            imsize=32,
                            shuffle=True,
                            seed=None,
                            interpolation='nearest'):
        ''' Returns an iterator through the files we have '''
        return DirectoryIterator(
            directory, self,
            batch_size=batch_size,
            num_timesteps=num_timesteps,
            num_classes=num_classes,
            numchannels=numchannels,
            imsize=imsize,
            shuffle=shuffle,
            seed=seed,
            data_format=self.data_format,
            interpolation=interpolation)

    def random_transform(self, x, seed=None):
        """Randomly augment a single image tensor.

        # Arguments
            x: 3D tensor, single image.
            seed: random seed.

        # Returns
            A randomly transformed version of the input (same shape).
        """
        # x is a single image, so it doesn't have image number at index 0
        img_row_axis = self.row_axis - 1
        img_col_axis = self.col_axis - 1
        img_channel_axis = self.channel_axis - 1

        if seed is not None:
            np.random.seed(seed)

        # use composition of homographies
        # to generate final transform that needs to be applied
        if self.rotation_range:
            theta = np.deg2rad(
                np.random.uniform(-self.rotation_range, self.rotation_range))
        else:
            theta = 0

        if self.height_shift_range:
            tx = np.random.uniform(-self.height_shift_range,
                                   self.height_shift_range)
            if self.height_shift_range < 1:
                tx *= x.shape[img_row_axis]
        else:
            tx = 0

        if self.width_shift_range:
            ty = np.random.uniform(-self.width_shift_range,
                                   self.width_shift_range)
            if self.width_shift_range < 1:
                ty *= x.shape[img_col_axis]
        else:
            ty = 0

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

        if self.zoom_range[0] == 1 and self.zoom_range[1] == 1:
            zx, zy = 1, 1
        else:
            zx, zy = np.random.uniform(
                self.zoom_range[0], self.zoom_range[1], 2)

        transform_matrix = None
        if theta != 0:
            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]])
            transform_matrix = shift_matrix if transform_matrix is None else np.dot(
                transform_matrix, shift_matrix)

        if shear != 0:
            shear_matrix = np.array([[1, -np.sin(shear), 0],
                                     [0, np.cos(shear), 0],
                                     [0, 0, 1]])
            transform_matrix = shear_matrix if transform_matrix is None else 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]])
            transform_matrix = zoom_matrix if transform_matrix is None else np.dot(
                transform_matrix, zoom_matrix)

        if transform_matrix is not None:
            h, w = x.shape[img_row_axis], x.shape[img_col_axis]
            transform_matrix = transform_matrix_offset_center(
                transform_matrix, h, w)
            x = apply_transform(x, transform_matrix, img_channel_axis,
                                fill_mode=self.fill_mode, cval=self.cval)

        if self.channel_shift_range != 0:
            x = random_channel_shift(x,
                                     self.channel_shift_range,
                                     img_channel_axis)
        if self.horizontal_flip:
            if np.random.random() < 0.5:
                x = flip_axis(x, img_col_axis)

        if self.vertical_flip:
            if np.random.random() < 0.5:
                x = flip_axis(x, img_row_axis)
        return x

    def fit(self, x,
            augment=False,
            rounds=1,
            seed=None):
        """Fits internal statistics to some sample data.

        Required for featurewise_center, featurewise_std_normalization
        and zca_whitening.

        # Arguments
            x: Numpy array, the data to fit on. Should have rank 4.
                In case of grayscale data,
                the channels axis should have value 1, and in case
                of RGB data, it should have value 3.
            augment: Whether to fit on randomly augmented samples
            rounds: If `augment`,
                how many augmentation passes to do over the data
            seed: random seed.

        # Raises
            ValueError: in case of invalid input `x`.
        """
        x = np.asarray(x, dtype=K.floatx())
        if x.ndim != 4:
            raise ValueError('Input to `.fit()` should have rank 4. '
                             'Got array with shape: ' + str(x.shape))
        if x.shape[self.channel_axis] not in {1, 3, 4}:
            warnings.warn(
                'Expected input to be images (as Numpy array) '
                '(channels on axis ' + str(self.channel_axis) + '), i.e. expected '
                'either 1, 3 or 4 channels on axis ' +
                str(self.channel_axis) + '. '
                'However, it was passed an array with shape ' + str(x.shape) +
                ' (' + str(x.shape[self.channel_axis]) + ' channels).')

        if seed is not None:
            np.random.seed(seed)

        x = np.copy(x)
        if augment:
            ax = np.zeros(tuple([rounds * x.shape[0]] +
                                list(x.shape)[1:]), dtype=K.floatx())
            for r in range(rounds):
                for i in range(x.shape[0]):
                    ax[i + r * x.shape[0]] = self.random_transform(x[i])
            x = ax

        if self.featurewise_center:
            self.mean = np.mean(x, axis=(0, self.row_axis, self.col_axis))
            broadcast_shape = [1, 1, 1]
            broadcast_shape[self.channel_axis - 1] = x.shape[self.channel_axis]
            self.mean = np.reshape(self.mean, broadcast_shape)
            x -= self.mean

        if self.featurewise_std_normalization:
            self.std = np.std(x, axis=(0, self.row_axis, self.col_axis))
            broadcast_shape = [1, 1, 1]
            broadcast_shape[self.channel_axis - 1] = x.shape[self.channel_axis]
            self.std = np.reshape(self.std, broadcast_shape)
            x /= (self.std + K.epsilon())

        if self.zca_whitening:
            flat_x = np.reshape(
                x, (x.shape[0], x.shape[1] * x.shape[2] * x.shape[3]))
            num_examples = flat_x.shape[0]
            _, s, vt = linalg.svd(flat_x / np.sqrt(num_examples))
            s_expand = np.hstack((s, np.zeros(vt.shape[0] - num_examples,
                                              dtype=flat_x.dtype)))
            self.principal_components = (
                vt.T / np.sqrt(s_expand ** 2 + self.zca_epsilon)).dot(vt)

    def standardize(self, x):
        """Apply the normalization configuration to a batch of inputs.
        # Arguments
            x: batch of inputs to be normalized.
        # Returns
            The inputs, normalized.
        """
        if self.samplewise_center:
            x -= np.mean(x, keepdims=True)
        if self.samplewise_std_normalization:
            x /= (np.std(x, keepdims=True) + K.epsilon())

        if self.featurewise_center:
            if self.mean is not None:
                x -= self.mean
            else:
                warnings.warn('This ImageDataGenerator specifies '
                              '`featurewise_center`, but it hasn\'t '
                              'been fit on any training data. Fit it '
                              'first by calling `.fit(numpy_data)`.')
        if self.featurewise_std_normalization:
            if self.std is not None:
                x /= (self.std + K.epsilon())
            else:
                warnings.warn('This ImageDataGenerator specifies '
                              '`featurewise_std_normalization`, but it hasn\'t '
                              'been fit on any training data. Fit it '
                              'first by calling `.fit(numpy_data)`.')
        return x


class DirectoryIterator(Iterator):
    """Iterator capable of reading images from a directory on disk.

    # Arguments
        directory: Path to the directory to read images from.
            Each subdirectory in this directory will be
            considered to contain images from one class,
            or alternatively you could specify class subdirectories
            via the `classes` argument.
        image_data_generator: Instance of `ImageDataGenerator`
            to use for random transformations and normalization.
        target_size: tuple of integers, dimensions to resize input images to.
        color_mode: One of `"rgb"`, `"grayscale"`. Color mode to read images.
        classes: Optional list of strings, names of subdirectories
            containing images from each class (e.g. `["dogs", "cats"]`).
            It will be computed automatically if not set.
        class_mode: Mode for yielding the targets:
            `"binary"`: binary targets (if there are only two classes),
            `"categorical"`: categorical targets,
            `"sparse"`: integer targets,
            `"input"`: targets are images identical to input images (mainly
                used to work with autoencoders),
            `None`: no targets get yielded (only input images are yielded).
        batch_size: Integer, size of a batch.
        shuffle: Boolean, whether to shuffle the data between epochs.
        seed: Random seed for data shuffling.
        data_format: String, one of `channels_first`, `channels_last`.
        save_to_dir: Optional directory where to save the pictures
            being yielded, in a viewable format. This is useful
            for visualizing the random transformations being
            applied, for debugging purposes.
        save_prefix: String prefix to use for saving sample
            images (if `save_to_dir` is set).
        save_format: Format to use for saving sample images
            (if `save_to_dir` is set).
        interpolation: Interpolation method used to resample the image if the
            target size is different from that of the loaded image.
            Supported methods are "nearest", "bilinear", and "bicubic".
            If PIL version 1.1.3 or newer is installed, "lanczos" is also
            supported. If PIL version 3.4.0 or newer is installed, "box" and
            "hamming" are also supported. By default, "nearest" is used.
    """

    def __init__(self,
                 directory,
                 image_data_generator,
                 batch_size=32,
                 num_timesteps=5,
                 num_classes=2,
                 numchannels=3,
                 imsize=32,
                 shuffle=True,
                 seed=None,
                 data_format=None,
                 interpolation='nearest'):
        if data_format is None:
            data_format = K.image_data_format()
        self.num_timesteps = num_timesteps
        self.data_format = data_format
        self.directory = directory
        self.image_data_generator = image_data_generator
        self.num_classes = num_classes
        self.numchannels = numchannels
        self.imsize = imsize

        # how to apply interpolation
        self.interpolation = interpolation

        # set the file paths of all the sequences we want to iterate over
        self.samples = 0
        self.filepaths = []
        for root, dirs, files in os.walk(self.directory):
            for file in files:
                self.filepaths.append(os.path.join(root, file))
        self.samples = len(self.filepaths)

        print('Found %d image files.' % (self.samples))
        super(DirectoryIterator, self).__init__(
            self.samples, batch_size, shuffle, seed)

    def _get_batches_of_transformed_sequences(self, index_array):
        # build batch of image data
        for i, index in enumerate(index_array):
            fpath = self.filepaths[index]

            # load in a certain file and randomly sample the image
            imgdata = np.load(fpath)
            images = imgdata['image_tensor']
            metadata = imgdata['metadata'].item()
            y = metadata['ylabels']

            # get the random index in this dataset
            numsamples = y.shape[0]

            # reformat data if necessary
            images, y = self._reformatdata(images, y)

            # get it as a sequence
            N = 1  # number of samples to get from this dataset | Todo change into higher number
            assert numsamples == images.shape[0]
            randind = random.sample(range(numsamples), k=N)
            randind = randind[0]
            sampimages, samplabels = self._sample_as_seq(images, y, randind)

            # apply normalization to each image
            for idx in range(0, len(sampimages)):
                sampimage = sampimages[idx, ...]
                # apply transforms and standardize
                sampimage = self.image_data_generator.random_transform(
                    sampimage)
                sampimage = self.image_data_generator.standardize(sampimage)

            if i == 0:
                # init batch of X and Y
                batch_x = np.zeros((len(index_array),) +
                                   sampimages.shape, dtype=K.floatx())
                batch_y = np.zeros((len(index_array),) +
                                   samplabels.shape, dtype=K.floatx())

            batch_x[i] = sampimages
            batch_y[i] = samplabels

        return batch_x, batch_y

    def _sample_as_seq(self, sampdata, samplabels, randind):
        '''
        Get sampled data as a sequence
        '''
        endlen = sampdata.shape[0]

        if randind + self.num_timesteps >= endlen:
            print("padding...")
            sampdata = sampdata[randind:, ...]
            sampdata = self._paddata(sampdata, padding='post')

            samplabels = samplabels[randind:, ...]
            samplabels = self._paddata(samplabels, padding='post')
        else:
            sampdata = sampdata[randind:randind+self.num_timesteps, ...]
            samplabels = samplabels[randind:randind+self.num_timesteps, ...]
        return sampdata, samplabels

    def _paddata(self, sampdata, padding):
        '''
        Pad a sample data on the 
        '''
        samplen = sampdata.shape[0]
        numsamps_topad = self.num_timesteps - samplen
        imshape = sampdata.shape[1:]
        assert imshape[0] == self.imsize
        assert imshape[-1] == self.numchannels

        # pad with zeros
        if samplen < self.num_timesteps:
            padmat = np.zeros(((samplen,)+imshape))
            if padding == 'pre':
                # pad on the pre
                sampdata = np.concatenate((padmat, sampdata), axis=0)
            elif padding == 'post':
                # pad on the pre
                sampdata = np.concatenate((sampdata, padmat), axis=0)
        return sampdata

    def _reformatdata(self, images, y):
        invert_y = 1 - y
        y = np.concatenate((y, invert_y), axis=1)

        # transform the images if necessary
        # reshape
        images = images.reshape(
            (-1, self.numchannels, self.imsize, self.imsize))
        images = images.swapaxes(1, 3)
        images = images.astype("float32")
        return images, y

    def next(self):
        """For python 2.x.

        # Returns
            The next batch.
        """
        with self.lock:
            index_array = next(self.index_generator)
        # The transformation of images is not under thread lock
        # so it can be done in parallel
        return self._get_batches_of_transformed_sequences(index_array)

# np.random.seed(123) # import all image transformation utility functions Generate minibatches of image data with real-time data augmentation.      Modified from Keras ImageDataGenerator      # Arguments         featurewise_center: set input mean to 0 over the dataset.         samplewise_center: set each sample mean to 0.         featurewise_std_normalization: divide inputs by std of the dataset.         samplewise_std_normalization: divide each input by its std.         rotation_range: degrees (0 to 180).         width_shift_range: fraction of total width, if < 1, or pixels if >= 1.         height_shift_range: fraction of total height, if < 1, or pixels if >= 1.         shear_range: shear intensity (shear angle in degrees).         zoom_range: amount of zoom. if scalar z, zoom will be randomly picked             in the range [1-z, 1+z]. A sequence of two can be passed instead             to select this range.         channel_shift_range: shift range for each channel.         fill_mode: points outside the boundaries are filled according to the             given mode ('constant', 'nearest', 'reflect' or 'wrap'). Default             is 'nearest'.             Points outside the boundaries of the input are filled according to the given mode:                 'constant': kkkkkkkk|abcd|kkkkkkkk (cval=k)                 'nearest':  aaaaaaaa|abcd|dddddddd                 'reflect':  abcddcba|abcd|dcbaabcd                 'wrap':  abcdabcd|abcd|abcdabcd         horizontal_flip: whether to randomly flip images horizontally.         vertical_flip: whether to randomly flip images vertically.         rescale: rescaling factor. If None or 0, no rescaling is applied,             otherwise we multiply the data by the value provided. This is             applied after the `preprocessing_function` (if any provided)             but before any other transformation.         preprocessing_function: function that will be implied on each input.             The function will run before any other modification on it.             The function should take one argument:             one image (Numpy tensor with rank 3),             and should output a Numpy tensor with the same shape.         data_format: 'channels_first' or 'channels_last'. In 'channels_first' mode, the channels dimension             (the depth) is at index 1, in 'channels_last' mode it is at index 3.             It defaults to the `image_data_format` value found in your             Keras config file at `~/.keras/keras.json`.             If you never set it, then it will be "channels_last". # img_row_axis = self.row_axis - 1 # img_col_axis = self.col_axis - 1 # img_channel_axis = self.channel_axis - 1 # initialize mean and std and pc for random augmentations Returns an iterator through the files we have Randomly augment a single image tensor.          # Arguments             x: 3D tensor, single image.             seed: random seed.          # Returns             A randomly transformed version of the input (same shape). # x is a single image, so it doesn't have image number at index 0 # use composition of homographies # to generate final transform that needs to be applied Fits internal statistics to some sample data.          Required for featurewise_center, featurewise_std_normalization         and zca_whitening.          # Arguments             x: Numpy array, the data to fit on. Should have rank 4.                 In case of grayscale data,                 the channels axis should have value 1, and in case                 of RGB data, it should have value 3.             augment: Whether to fit on randomly augmented samples             rounds: If `augment`,                 how many augmentation passes to do over the data             seed: random seed.          # Raises             ValueError: in case of invalid input `x`. Apply the normalization configuration to a batch of inputs.         # Arguments             x: batch of inputs to be normalized.         # Returns             The inputs, normalized. Iterator capable of reading images from a directory on disk.      # Arguments         directory: Path to the directory to read images from.             Each subdirectory in this directory will be             considered to contain images from one class,             or alternatively you could specify class subdirectories             via the `classes` argument.         image_data_generator: Instance of `ImageDataGenerator`             to use for random transformations and normalization.         target_size: tuple of integers, dimensions to resize input images to.         color_mode: One of `"rgb"`, `"grayscale"`. Color mode to read images.         classes: Optional list of strings, names of subdirectories             containing images from each class (e.g. `["dogs", "cats"]`).             It will be computed automatically if not set.         class_mode: Mode for yielding the targets:             `"binary"`: binary targets (if there are only two classes),             `"categorical"`: categorical targets,             `"sparse"`: integer targets,             `"input"`: targets are images identical to input images (mainly                 used to work with autoencoders),             `None`: no targets get yielded (only input images are yielded).         batch_size: Integer, size of a batch.         shuffle: Boolean, whether to shuffle the data between epochs.         seed: Random seed for data shuffling.         data_format: String, one of `channels_first`, `channels_last`.         save_to_dir: Optional directory where to save the pictures             being yielded, in a viewable format. This is useful             for visualizing the random transformations being             applied, for debugging purposes.         save_prefix: String prefix to use for saving sample             images (if `save_to_dir` is set).         save_format: Format to use for saving sample images             (if `save_to_dir` is set).         interpolation: Interpolation method used to resample the image if the             target size is different from that of the loaded image.             Supported methods are "nearest", "bilinear", and "bicubic".             If PIL version 1.1.3 or newer is installed, "lanczos" is also             supported. If PIL version 3.4.0 or newer is installed, "box" and             "hamming" are also supported. By default, "nearest" is used. # how to apply interpolation # set the file paths of all the sequences we want to iterate over # build batch of image data # load in a certain file and randomly sample the image # get the random index in this dataset # reformat data if necessary # get it as a sequence # number of samples to get from this dataset | Todo change into higher number # apply normalization to each image # apply transforms and standardize # init batch of X and Y Get sampled data as a sequence Pad a sample data on the # pad with zeros # pad on the pre # pad on the pre # transform the images if necessary # reshape For python 2.x.          # Returns             The next batch. # The transformation of images is not under thread lock # so it can be done in parallel