import tensorflow as tf
from tensorflow.compat.v1 import keras as tfk
from tensorflow.keras.callbacks import Callback, EarlyStopping
from tensorflow.keras.layers import ZeroPadding2D, ZeroPadding3D, LocallyConnected2D, Lambda, Layer
from tensorflow.keras.losses import mean_absolute_error, mean_squared_error
from tensorflow.python.keras.utils import conv_utils
from tensorflow.python.keras.engine.base_layer import InputSpec
from tensorflow.keras import activations, initializers, regularizers, constraints
import numpy as np

try:
    from s2cnn import S2Convolution, SO3Convolution
except ImportError:
    pass

K = tfk.backend

# ==================================================================================================================== #
# Keras utility classes
# ==================================================================================================================== #

class AdamLearningRateTracker(Callback):
    """
    Log the current learning rate used by and Adam optimizer.
    """

    def on_epoch_end(self, epoch, logs=None, beta_1=0.9, beta_2=0.999, ):
        optimizer = self.model.optimizer
        it = K.cast(optimizer.iterations, K.floatx())
        lr = K.cast(optimizer.lr, K.floatx())
        decay = K.cast(optimizer.decay, K.floatx())
        t = K.eval(it + 1.)
        new_lr = K.eval(lr * (1. / (1. + decay * it)))
        lr_t = K.eval(new_lr * (K.sqrt(1. - K.pow(beta_2, t)) / (1. - K.pow(beta_1, t))))
        print(' - LR: {:.6f}'.format(lr_t))

class SGDLearningRateTracker(Callback):
    """
    Log the current learning rate used by an SGD optimizer.
    """

    def on_epoch_end(self, epoch, logs=None):
        optimizer = self.model.optimizer
        it = K.cast(optimizer.iterations, K.floatx())
        lr = K.cast(optimizer.lr, K.floatx())
        decay = K.cast(optimizer.decay, K.floatx())
        new_lr = K.eval(lr * (1. / (1. + decay * it)))
        print(' - LR: {:.6f}'.format(new_lr))

class BatchHistory(Callback):
    """
    Log training metrics for each batch of training data.
    """

    def on_train_begin(self, logs=None):
        self.history = []
        self.epoch = 0

    def on_epoch_begin(self, epoch, logs=None):
        self.history.append({})

    def on_epoch_end(self, epoch, logs=None):
        self.epoch += 1

    def on_batch_end(self, batch, logs=None):
        logs = logs or {}
        for k, v in logs.items():
            self.history[self.epoch].setdefault(k, []).append(v)

class RunHistory(Callback):
    """
    Callback that records events into a `History` object.
    Adapted from keras.callbacks.History to include logging to Azure experiment runs.
    """

    def __init__(self, run):
        super(RunHistory, self).__init__()
        self.epoch = []
        self.history = {}
        self.run = run

    def on_train_begin(self, logs=None):
        self.epoch = []
        self.history = {}

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        self.epoch.append(epoch)
        for k, v in logs.items():
            self.history.setdefault(k, []).append(v)
            self.run.log(k, v)

class RNNResetStates(Callback):
    def on_epoch_begin(self, epoch, logs=None):
        self.model.reset_states()

class EarlyStoppingMin(EarlyStopping):
    """
    Extends the keras.callbacks.EarlyStopping class to provide the option to force training for a minimum number of
    epochs or restore the best weights after the maximum epochs have been reached.
    """

    def __init__(self, min_epochs=0, max_epochs=None, **kwargs):
        """
        :param min_epochs: int: train the network for at least this number of epochs before early stopping
        :param max_epochs: int: train the network for at most this number of epochs before early stopping
        :param kwargs: passed to EarlyStopping.__init__()
        """
        super(EarlyStoppingMin, self).__init__(**kwargs)
        if not isinstance(min_epochs, int) or min_epochs < 0:
            raise ValueError('min_epochs must be an integer >= 0')
        self.min_epochs = int(min_epochs)
        self.max_epochs = int(max_epochs) if max_epochs is not None else None

    def on_epoch_end(self, epoch, logs=None):
        if epoch < self.min_epochs:
            return

        current = self.get_monitor_value(logs)
        if current is None:
            return

        if self.monitor_op(current - self.min_delta, self.best):
            self.best = current
            self.wait = 0
            if self.restore_best_weights:
                self.best_weights = self.model.get_weights()
        else:
            self.wait += 1
            if self.wait >= self.patience:
                self.stopped_epoch = epoch
                self.model.stop_training = True
                if self.restore_best_weights:
                    if self.verbose > 0:
                        print('Restoring model weights from the end of '
                              'the best epoch')
                    self.model.set_weights(self.best_weights)

        if self.max_epochs is not None and epoch >= self.max_epochs:
            self.stopped_epoch = epoch
            self.model.stop_training = True
            if self.restore_best_weights:
                if self.verbose > 0:
                    print('Maximum epochs reached; restoring model weights from the end of '
                          'the best epoch')
                self.model.set_weights(self.best_weights)

        if self.verbose > 1:
            print('EarlyStoppingMin: %d epochs since last minimum' % self.wait)

class SaveWeightsOnEpoch(Callback):
    """
    Saves the model weights to a temporary file at the end of each epoch. This is useful for avoiding complete loss
    of a run that fails for any reason.
    """

    def __init__(self, weights_file, interval=None):
        """
        :param weights_file: str: file name to save weights
        """
        super(SaveWeightsOnEpoch, self).__init__()
        self.weights_file = str(weights_file)
        if interval is not None:
            assert isinstance(interval, int) and interval > 0, "'interval' must be an integer > 0"
        self.interval = interval

    def on_epoch_end(self, epoch, logs=None):
        if self.interval is not None and epoch % self.interval == 0:
            self.model.save_weights('%s.%s' % (self.weights_file, epoch), save_format='h5')
        else:
            try:
                self.model.save_weights(self.weights_file, save_format='h5')
            except OSError:  # If the file is temporarily system locked for any reason, avoid crashing
                pass

class GeneratorEpochEnd(Callback):
    """
    Used in conjunction with a tensorflow.data.Dataset generator to manually execute the on_epoch_end() method of a
    DLWP.model.generators Generator instance after each epoch.
    """

    def __init__(self, generator):
        """
        :param generator: DLWP.model.generator instance
        """
        super(GeneratorEpochEnd, self).__init__()
        self.generator = generator

    def on_epoch_end(self, epoch, logs=None):
        self.generator.on_epoch_end()

# ==================================================================================================================== #
# Keras padding layers
# ==================================================================================================================== #

class PeriodicPadding2D(ZeroPadding2D):
    """Periodic-padding layer for 2D input (e.g. image).

    This layer can add periodic rows and columns at the top, bottom, left and right side of an image tensor.

    Adapted from keras.layers.ZeroPadding2D by @jweyn

    # Arguments
        padding: int, or tuple of 2 ints, or tuple of 2 tuples of 2 ints.
            - If int: the same symmetric padding
                is applied to height and width.
            - If tuple of 2 ints:
                interpreted as two different
                symmetric padding values for height and width:
                `(symmetric_height_pad, symmetric_width_pad)`.
            - If tuple of 2 tuples of 2 ints:
                interpreted as
                `((top_pad, bottom_pad), (left_pad, right_pad))`
        data_format: A string,
            one of `"channels_last"` or `"channels_first"`.
            The ordering of the dimensions in the inputs.
            `"channels_last"` corresponds to inputs with shape
            `(batch, height, width, channels)` while `"channels_first"`
            corresponds to inputs with shape
            `(batch, channels, height, width)`.
            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".

    # Input shape
        4D tensor with shape:
        - If `data_format` is `"channels_last"`:
            `(batch, rows, cols, channels)`
        - If `data_format` is `"channels_first"`:
            `(batch, channels, rows, cols)`

    # Output shape
        4D tensor with shape:
        - If `data_format` is `"channels_last"`:
            `(batch, padded_rows, padded_cols, channels)`
        - If `data_format` is `"channels_first"`:
            `(batch, channels, padded_rows, padded_cols)`
    """

    def __init__(self,
                 padding=(1, 1),
                 data_format=None,
                 **kwargs):
        super(PeriodicPadding2D, self).__init__(padding=padding,
                                                data_format=data_format,
                                                **kwargs)

    def call(self, inputs):
        if K.backend() == 'plaidml.keras.backend':
            shape = inputs.shape.dims
        else:
            shape = inputs.shape
        if self.data_format == 'channels_first':
            top_slice = slice(shape[2] - self.padding[0][0], shape[2])
            bottom_slice = slice(0, self.padding[0][1])
            left_slice = slice(shape[3] - self.padding[1][0], shape[3])
            right_slice = slice(0, self.padding[1][1])
            # Pad the horizontal
            outputs = K.concatenate([inputs[:, :, :, left_slice], inputs, inputs[:, :, :, right_slice]], axis=3)
            # Pad the vertical
            outputs = K.concatenate([outputs[:, :, top_slice], outputs, outputs[:, :, bottom_slice]], axis=2)
        else:
            top_slice = slice(shape[1] - self.padding[0][0], shape[1])
            bottom_slice = slice(0, self.padding[0][1])
            left_slice = slice(shape[2] - self.padding[1][0], shape[2])
            right_slice = slice(0, self.padding[1][1])
            # Pad the horizontal
            outputs = K.concatenate([inputs[:, :, left_slice], inputs, inputs[:, :, right_slice]], axis=2)
            # Pad the vertical
            outputs = K.concatenate([outputs[:, top_slice], outputs, outputs[:, bottom_slice]], axis=1)
        return outputs

class PeriodicPadding3D(ZeroPadding3D):
    """Periodic-padding layer for 3D input (e.g. image).

    This layer can add periodic rows, columns, and depth to an image tensor.

    Adapted from keras.layers.ZeroPadding3D by @jweyn

    # Arguments
        padding: int, or tuple of 3 ints, or tuple of 3 tuples of 2 ints.
            - If int: the same symmetric padding
                is applied to height and width.
            - If tuple of 3 ints:
                interpreted as two different
                symmetric padding values for height and width:
                `(symmetric_dim1_pad, symmetric_dim2_pad, symmetric_dim3_pad)`.
            - If tuple of 3 tuples of 2 ints:
                interpreted as
                `((left_dim1_pad, right_dim1_pad),
                  (left_dim2_pad, right_dim2_pad),
                  (left_dim3_pad, right_dim3_pad))`
        data_format: A string,
            one of `"channels_last"` or `"channels_first"`.
            The ordering of the dimensions in the inputs.
            `"channels_last"` corresponds to inputs with shape
            `(batch, spatial_dim1, spatial_dim2, spatial_dim3, channels)`
            while `"channels_first"` corresponds to inputs with shape
            `(batch, channels, spatial_dim1, spatial_dim2, spatial_dim3)`.
            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".

    # Input shape
        5D tensor with shape:
        - If `data_format` is `"channels_last"`:
            `(batch, first_axis_to_pad, second_axis_to_pad, third_axis_to_pad,
              depth)`
        - If `data_format` is `"channels_first"`:
            `(batch, depth,
              first_axis_to_pad, second_axis_to_pad, third_axis_to_pad)`

    # Output shape
        5D tensor with shape:
        - If `data_format` is `"channels_last"`:
            `(batch, first_padded_axis, second_padded_axis, third_axis_to_pad,
              depth)`
        - If `data_format` is `"channels_first"`:
            `(batch, depth,
              first_padded_axis, second_padded_axis, third_axis_to_pad)`
    """

    def __init__(self,
                 padding=(1, 1, 1),
                 data_format=None,
                 **kwargs):
        super(PeriodicPadding3D, self).__init__(padding=padding,
                                                data_format=data_format,
                                                **kwargs)

    def call(self, inputs):
        if K.backend() == 'plaidml.keras.backend':
            shape = inputs.shape.dims
        else:
            shape = inputs.shape
        if self.data_format == 'channels_first':
            low_slice = slice(shape[2] - self.padding[0][0], shape[2])
            high_slice = slice(0, self.padding[0][1])
            top_slice = slice(shape[3] - self.padding[1][0], shape[3])
            bottom_slice = slice(0, self.padding[1][1])
            left_slice = slice(shape[4] - self.padding[2][0], shape[4])
            right_slice = slice(0, self.padding[2][1])
            # Pad the horizontal
            outputs = K.concatenate([inputs[:, :, :, :, left_slice], inputs, inputs[:, :, :, :, right_slice]], axis=4)
            # Pad the vertical
            outputs = K.concatenate([outputs[:, :, :, top_slice], outputs, outputs[:, :, :, bottom_slice]], axis=3)
            # Pad the depth
            outputs = K.concatenate([outputs[:, :, low_slice], outputs, outputs[:, :, high_slice]], axis=2)
        else:
            low_slice = slice(shape[1] - self.padding[0][0], shape[1])
            high_slice = slice(0, self.padding[0][1])
            top_slice = slice(shape[2] - self.padding[1][0], shape[2])
            bottom_slice = slice(0, self.padding[1][1])
            left_slice = slice(shape[3] - self.padding[2][0], shape[3])
            right_slice = slice(0, self.padding[2][1])
            # Pad the horizontal
            outputs = K.concatenate([inputs[:, :, :, left_slice], inputs, inputs[:, :, :, right_slice]], axis=3)
            # Pad the vertical
            outputs = K.concatenate([outputs[:, :, top_slice], outputs, outputs[:, :, bottom_slice]], axis=2)
            # Pad the depth
            outputs = K.concatenate([outputs[:, low_slice], outputs, outputs[:, high_slice]], axis=1)
        return outputs

class FillPadding2D(ZeroPadding2D):
    """Fill-padding layer for 2D input (e.g. image).

    This layer can add rows or columns that duplicate the edge values.

    Adapted from keras.layers.ZeroPadding2D by @jweyn

    # Arguments
        padding: int, or tuple of 2 ints, or tuple of 2 tuples of 2 ints.
            - If int: the same symmetric padding
                is applied to height and width.
            - If tuple of 2 ints:
                interpreted as two different
                symmetric padding values for height and width:
                `(symmetric_height_pad, symmetric_width_pad)`.
            - If tuple of 2 tuples of 2 ints:
                interpreted as
                `((top_pad, bottom_pad), (left_pad, right_pad))`
        data_format: A string,
            one of `"channels_last"` or `"channels_first"`.
            The ordering of the dimensions in the inputs.
            `"channels_last"` corresponds to inputs with shape
            `(batch, height, width, channels)` while `"channels_first"`
            corresponds to inputs with shape
            `(batch, channels, height, width)`.
            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".

    # Input shape
        4D tensor with shape:
        - If `data_format` is `"channels_last"`:
            `(batch, rows, cols, channels)`
        - If `data_format` is `"channels_first"`:
            `(batch, channels, rows, cols)`

    # Output shape
        4D tensor with shape:
        - If `data_format` is `"channels_last"`:
            `(batch, padded_rows, padded_cols, channels)`
        - If `data_format` is `"channels_first"`:
            `(batch, channels, padded_rows, padded_cols)`
    """

    def __init__(self,
                 padding=(1, 1),
                 data_format=None,
                 **kwargs):
        super(FillPadding2D, self).__init__(padding=padding, data_format=data_format, **kwargs)

    def call(self, inputs):
        if self.data_format == 'channels_first':
            # Pad the vertical
            if self.padding[0][0] > 0:
                top_slice = K.stack([inputs[:, :, 0]] * self.padding[0][0], axis=2)
            else:
                top_slice = inputs[:, :, slice(0, 0)]
            if self.padding[0][1] > 0:
                bottom_slice = K.stack([inputs[:, :, -1]] * self.padding[0][1], axis=2)
            else:
                bottom_slice = inputs[:, :, slice(0, 0)]
            outputs = K.concatenate([top_slice, inputs, bottom_slice], axis=2)
            # Pad the horizontal
            if self.padding[1][0] > 0:
                left_slice = K.stack([outputs[:, :, :, 0]] * self.padding[1][0], axis=3)
            else:
                left_slice = outputs[:, :, :, slice(0, 0)]
            if self.padding[1][1] > 0:
                right_slice = K.stack([outputs[:, :, :, -1]] * self.padding[1][1], axis=3)
            else:
                right_slice = outputs[:, :, :, slice(0, 0)]
            outputs = K.concatenate([left_slice, outputs, right_slice], axis=3)
        else:
            # Pad the vertical
            if self.padding[0][0] > 0:
                top_slice = K.stack([inputs[:, 0]] * self.padding[0][0], axis=1)
            else:
                top_slice = inputs[:, slice(0, 0)]
            if self.padding[0][1] > 0:
                bottom_slice = K.stack([inputs[:, -1]] * self.padding[0][1], axis=1)
            else:
                bottom_slice = inputs[:, slice(0, 0)]
            outputs = K.concatenate([top_slice, inputs, bottom_slice], axis=1)
            # Pad the horizontal
            if self.padding[1][0] > 0:
                left_slice = K.stack([outputs[:, :, 0]] * self.padding[1][0], axis=2)
            else:
                left_slice = outputs[:, :, slice(0, 0)]
            if self.padding[1][1] > 0:
                right_slice = K.stack([outputs[:, :, -1]] * self.padding[1][1], axis=2)
            else:
                right_slice = outputs[:, :, slice(0, 0)]
            outputs = K.concatenate([left_slice, outputs, right_slice], axis=2)
        return outputs

class FillPadding3D(ZeroPadding3D):
    """Fill-padding layer for 3D input (e.g. image).

    This layer can add rows or columns that duplicate the edge values.

    Adapted from keras.layers.ZeroPadding3D by @jweyn

    # Arguments
        padding: int, or tuple of 3 ints, or tuple of 3 tuples of 2 ints.
            - If int: the same symmetric padding
                is applied to height and width.
            - If tuple of 3 ints:
                interpreted as two different
                symmetric padding values for height and width:
                `(symmetric_dim1_pad, symmetric_dim2_pad, symmetric_dim3_pad)`.
            - If tuple of 3 tuples of 2 ints:
                interpreted as
                `((left_dim1_pad, right_dim1_pad),
                  (left_dim2_pad, right_dim2_pad),
                  (left_dim3_pad, right_dim3_pad))`
        data_format: A string,
            one of `"channels_last"` or `"channels_first"`.
            The ordering of the dimensions in the inputs.
            `"channels_last"` corresponds to inputs with shape
            `(batch, spatial_dim1, spatial_dim2, spatial_dim3, channels)`
            while `"channels_first"` corresponds to inputs with shape
            `(batch, channels, spatial_dim1, spatial_dim2, spatial_dim3)`.
            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".

    # Input shape
        5D tensor with shape:
        - If `data_format` is `"channels_last"`:
            `(batch, first_axis_to_pad, second_axis_to_pad, third_axis_to_pad,
              depth)`
        - If `data_format` is `"channels_first"`:
            `(batch, depth,
              first_axis_to_pad, second_axis_to_pad, third_axis_to_pad)`

    # Output shape
        5D tensor with shape:
        - If `data_format` is `"channels_last"`:
            `(batch, first_padded_axis, second_padded_axis, third_axis_to_pad,
              depth)`
        - If `data_format` is `"channels_first"`:
            `(batch, depth,
              first_padded_axis, second_padded_axis, third_axis_to_pad)`
    """

    def __init__(self,
                 padding=(1, 1, 1),
                 data_format=None,
                 **kwargs):
        super(FillPadding3D, self).__init__(padding=padding, data_format=data_format, **kwargs)

    def call(self, inputs):
        if self.data_format == 'channels_first':
            # Pad the depth
            if self.padding[0][0] > 0:
                low_slice = K.stack([inputs[:, :, 0]] * self.padding[0][0], axis=2)
            else:
                low_slice = inputs[:, :, slice(0, 0)]
            if self.padding[0][1] > 0:
                high_slice = K.stack([inputs[:, :, -1]] * self.padding[0][1], axis=2)
            else:
                high_slice = inputs[:, :, slice(0, 0)]
            outputs = K.concatenate([low_slice, inputs, high_slice], axis=2)
            # Pad the vertical
            if self.padding[1][0] > 0:
                top_slice = K.stack([outputs[:, :, :, 0]] * self.padding[1][0], axis=3)
            else:
                top_slice = outputs[:, :, :, slice(0, 0)]
            if self.padding[1][1] > 0:
                bottom_slice = K.stack([outputs[:, :, :, -1]] * self.padding[1][1], axis=3)
            else:
                bottom_slice = outputs[:, :, :, slice(0, 0)]
            outputs = K.concatenate([top_slice, outputs, bottom_slice], axis=3)
            # Pad the horizontal
            if self.padding[2][0] > 0:
                left_slice = K.stack([outputs[:, :, :, :, 0]] * self.padding[2][0], axis=4)
            else:
                left_slice = outputs[:, :, :, :, slice(0, 0)]
            if self.padding[2][1] > 0:
                right_slice = K.stack([outputs[:, :, :, :, -1]] * self.padding[2][1], axis=4)
            else:
                right_slice = outputs[:, :, :, :, slice(0, 0)]
            outputs = K.concatenate([left_slice, outputs, right_slice], axis=4)
        else:
            # Pad the depth
            if self.padding[0][0] > 0:
                low_slice = K.stack([inputs[:, 0]] * self.padding[0][0], axis=1)
            else:
                low_slice = inputs[:, slice(0, 0)]
            if self.padding[0][1] > 0:
                high_slice = K.stack([inputs[:, -1]] * self.padding[0][1], axis=1)
            else:
                high_slice = inputs[:, slice(0, 0)]
            outputs = K.concatenate([low_slice, inputs, high_slice], axis=1)
            # Pad the vertical
            if self.padding[1][0] > 0:
                top_slice = K.stack([outputs[:, :, 0]] * self.padding[1][0], axis=2)
            else:
                top_slice = outputs[:, :, slice(0, 0)]
            if self.padding[1][1] > 0:
                bottom_slice = K.stack([outputs[:, :, -1]] * self.padding[1][1], axis=2)
            else:
                bottom_slice = outputs[:, :, slice(0, 0)]
            outputs = K.concatenate([top_slice, outputs, bottom_slice], axis=2)
            # Pad the horizontal
            if self.padding[2][0] > 0:
                left_slice = K.stack([outputs[:, :, :, 0]] * self.padding[2][0], axis=3)
            else:
                left_slice = outputs[:, :, :, slice(0, 0)]
            if self.padding[2][1] > 0:
                right_slice = K.stack([outputs[:, :, :, -1]] * self.padding[2][1], axis=3)
            else:
                right_slice = outputs[:, :, :, slice(0, 0)]
            outputs = K.concatenate([left_slice, outputs, right_slice], axis=3)
        return outputs

class TFPadding2D(ZeroPadding2D):
    """Padding layer for 2D input (e.g. image) using TensorFlow's padding function.

    Adapted from keras.layers.ZeroPadding2D by @jweyn

    # Arguments
        padding: int, or tuple of 2 ints, or tuple of 2 tuples of 2 ints.
            - If int: the same symmetric padding
                is applied to height and width.
            - If tuple of 2 ints:
                interpreted as two different
                symmetric padding values for height and width:
                `(symmetric_height_pad, symmetric_width_pad)`.
            - If tuple of 2 tuples of 2 ints:
                interpreted as
                `((top_pad, bottom_pad), (left_pad, right_pad))`
        data_format: A string,
            one of `"channels_last"` or `"channels_first"`.
            The ordering of the dimensions in the inputs.
            `"channels_last"` corresponds to inputs with shape
            `(batch, height, width, channels)` while `"channels_first"`
            corresponds to inputs with shape
            `(batch, channels, height, width)`.
            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".
        mode: A string,
            one of `"CONSTANT"`, `"SYMMETRIC"`, or `"REFLECT"`.
        constant_values: A float. The value to pad if mode=='CONSTANT'.

    # Input shape
        5D tensor with shape:
        - If `data_format` is `"channels_last"`:
            `(batch, first_axis_to_pad, second_axis_to_pad, third_axis_to_pad,
              depth)`
        - If `data_format` is `"channels_first"`:
            `(batch, depth,
              first_axis_to_pad, second_axis_to_pad, third_axis_to_pad)`

    # Output shape
        5D tensor with shape:
        - If `data_format` is `"channels_last"`:
            `(batch, first_padded_axis, second_padded_axis, third_axis_to_pad,
              depth)`
        - If `data_format` is `"channels_first"`:
            `(batch, depth,
              first_padded_axis, second_padded_axis, third_axis_to_pad)`
    """

    def __init__(self,
                 padding=(1, 1),
                 data_format=None,
                 mode='CONSTANT',
                 constant_values=0.,
                 **kwargs):
        super(TFPadding2D, self).__init__(padding=padding, data_format=data_format, **kwargs)
        self.mode = mode
        self.constant_values = constant_values

    def call(self, inputs):
        if self.data_format == 'channels_first':
            padding = ((0, 0), (0, 0)) + self.padding
        else:
            padding = ((0, 0),) + self.padding + ((0, 0),)
        return tf.pad(inputs, padding, mode=self.mode, constant_values=self.constant_values)

    def get_config(self):
        config = {'padding': self.padding,
                  'data_format': self.data_format,
                  'mode': self.mode,
                  'constant_values': self.constant_values}
        base_config = super(TFPadding2D, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class TFPadding3D(ZeroPadding3D):
    """Padding layer for 3D input (e.g. image) using TensorFlow's padding function.

    Adapted from keras.layers.ZeroPadding3D by @jweyn

    # Arguments
        padding: int, or tuple of 3 ints, or tuple of 3 tuples of 2 ints.
            - If int: the same symmetric padding
                is applied to height and width.
            - If tuple of 3 ints:
                interpreted as two different
                symmetric padding values for height and width:
                `(symmetric_dim1_pad, symmetric_dim2_pad, symmetric_dim3_pad)`.
            - If tuple of 3 tuples of 2 ints:
                interpreted as
                `((left_dim1_pad, right_dim1_pad),
                  (left_dim2_pad, right_dim2_pad),
                  (left_dim3_pad, right_dim3_pad))`
        data_format: A string,
            one of `"channels_last"` or `"channels_first"`.
            The ordering of the dimensions in the inputs.
            `"channels_last"` corresponds to inputs with shape
            `(batch, spatial_dim1, spatial_dim2, spatial_dim3, channels)`
            while `"channels_first"` corresponds to inputs with shape
            `(batch, channels, spatial_dim1, spatial_dim2, spatial_dim3)`.
            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".
        mode: A string,
            one of `"CONSTANT"`, `"SYMMETRIC"`, or `"REFLECT"`.
        constant_values: A float. The value to pad if mode=='CONSTANT'.

    # Input shape
        4D tensor with shape:
        - If `data_format` is `"channels_last"`:
            `(batch, rows, cols, channels)`
        - If `data_format` is `"channels_first"`:
            `(batch, channels, rows, cols)`

    # Output shape
        4D tensor with shape:
        - If `data_format` is `"channels_last"`:
            `(batch, padded_rows, padded_cols, channels)`
        - If `data_format` is `"channels_first"`:
            `(batch, channels, padded_rows, padded_cols)`
    """

    def __init__(self,
                 padding=(1, 1, 1),
                 data_format=None,
                 mode='CONSTANT',
                 constant_values=0.,
                 **kwargs):
        super(TFPadding3D, self).__init__(padding=padding, data_format=data_format, **kwargs)
        self.mode = mode
        self.constant_values = constant_values

    def call(self, inputs):
        if self.data_format == 'channels_first':
            padding = ((0, 0), (0, 0)) + self.padding
        else:
            padding = ((0, 0),) + self.padding + ((0, 0),)
        return tf.pad(inputs, padding, mode=self.mode, constant_values=self.constant_values)

    def get_config(self):
        config = {'padding': self.padding,
                  'data_format': self.data_format,
                  'mode': self.mode,
                  'constant_values': self.constant_values}
        base_config = super(TFPadding3D, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

# ==================================================================================================================== #
# Keras classes for cube sphere convolution
# ==================================================================================================================== #

class CubeSphereConv2D(Layer):
    """
    2D convolutional layer for data that is assumed on a cubed sphere. The requirements for using this layer are as
    follows:
    - The input data is 5-dimensional:
        (batch, channels, 6, height, width) for channels_first
        (batch, 6, height, width, channels) for channels_last
    - The first spatial dimension must have a length of 6 for the 6 faces of the cubed sphere
    - The last two faces (indices 4 and 5) are the polar faces

    This layer learns two separate convolutional kernels and biases, one for the equatorial faces of the cube, and one
    for the polar faces. Optionally, it can learn separate kernels and biases for each polar face.

    Note that this layer should be preceded by CubeSpherePadding2D otherwise there is no connection between faces of
    the cube.

    Adapted from keras.layers._Conv by @jweyn

    # Arguments
        filters: Integer, the dimensionality of the output space
            (i.e. the number of output filters in the convolution).
        kernel_size: An integer or tuple/list of n integers, specifying the
            dimensions of the convolution window.
        strides: An integer or tuple/list of n integers,
            specifying the strides of the convolution.
            Specifying any stride value != 1 is incompatible with specifying
            any `dilation_rate` value != 1.
        padding: One of `"valid"` or `"same"` (case-insensitive).
        data_format: A string,
            one of `"channels_last"` or `"channels_first"`.
            The ordering of the dimensions in the inputs.
            `"channels_last"` corresponds to inputs with shape
            `(batch, ..., channels)` while `"channels_first"` corresponds to
            inputs with shape `(batch, channels, ...)`.
            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".
        dilation_rate: An integer or tuple/list of n integers, specifying
            the dilation rate to use for dilated convolution.
            Currently, specifying any `dilation_rate` value != 1 is
            incompatible with specifying any `strides` value != 1.
        activation: Activation function to use
            (see [activations](../activations.md)).
            If you don't specify anything, no activation is applied
            (ie. "linear" activation: `a(x) = x`).
        use_bias: Boolean, whether the layer uses a bias vector.
        flip_north_pole: Boolean, whether to reverse the direction of the
            north pole should that be necessary to match the rotation
            direction of the south pole in the data
        independent_north_pole: Boolean, if true, learn separate filters
            for the north and south poles (3 total sets)
        kernel_initializer: Initializer for the `kernel` weights matrix
            (see [initializers](../initializers.md)).
        bias_initializer: Initializer for the bias vector
            (see [initializers](../initializers.md)).
        kernel_regularizer: Regularizer function applied to
            the `kernel` weights matrix
            (see [regularizer](../regularizers.md)).
        bias_regularizer: Regularizer function applied to the bias vector
            (see [regularizer](../regularizers.md)).
        activity_regularizer: Regularizer function applied to
            the output of the layer (its "activation").
            (see [regularizer](../regularizers.md)).
        kernel_constraint: Constraint function applied to the kernel matrix
            (see [constraints](../constraints.md)).
        bias_constraint: Constraint function applied to the bias vector
            (see [constraints](../constraints.md)).
    """

    def __init__(self,
                 filters,
                 kernel_size,
                 strides=1,
                 padding='valid',
                 data_format='channels_first',
                 dilation_rate=1,
                 activation=None,
                 use_bias=True,
                 flip_north_pole=True,
                 independent_north_pole=False,
                 kernel_initializer='glorot_uniform',
                 bias_initializer='zeros',
                 kernel_regularizer=None,
                 bias_regularizer=None,
                 activity_regularizer=None,
                 kernel_constraint=None,
                 bias_constraint=None,
                 **kwargs):
        super(CubeSphereConv2D, self).__init__(**kwargs)
        self.filters = filters
        self.kernel_size = conv_utils.normalize_tuple(kernel_size, 2, 'kernel_size')
        self.strides = conv_utils.normalize_tuple(strides, 2, 'strides')
        self.padding = conv_utils.normalize_padding(padding)
        self.data_format = conv_utils.normalize_data_format(data_format)
        self.dilation_rate = conv_utils.normalize_tuple(dilation_rate, 2, 'dilation_rate')
        self.activation = activations.get(activation)
        self.use_bias = use_bias
        self.flip_north_pole = flip_north_pole
        self.independent_north_pole = independent_north_pole
        self.kernel_initializer = initializers.get(kernel_initializer)
        self.bias_initializer = initializers.get(bias_initializer)
        self.kernel_regularizer = regularizers.get(kernel_regularizer)
        self.bias_regularizer = regularizers.get(bias_regularizer)
        self.activity_regularizer = regularizers.get(activity_regularizer)
        self.kernel_constraint = constraints.get(kernel_constraint)
        self.bias_constraint = constraints.get(bias_constraint)
        self.rank = 3
        self.input_spec = InputSpec(ndim=self.rank + 2)

        self.equatorial_kernel = None
        self.equatorial_bias = None
        self.polar_kernel = None
        self.polar_bias = None
        self.north_pole_kernel = None
        self.north_pole_bias = None

    def build(self, input_shape):
        if self.data_format == 'channels_first':
            channel_axis = 1
        else:
            channel_axis = -1
        if input_shape[channel_axis] is None:
            raise ValueError('The channel dimension of the inputs '
                             'should be defined. Found `None`.')
        input_dim = input_shape[channel_axis]
        kernel_shape = self.kernel_size + (input_dim, self.filters)

        self.equatorial_kernel = self.add_weight(shape=kernel_shape,
                                                 initializer=self.kernel_initializer,
                                                 name='equatorial_kernel',
                                                 regularizer=self.kernel_regularizer,
                                                 constraint=self.kernel_constraint)
        self.polar_kernel = self.add_weight(shape=kernel_shape,
                                            initializer=self.kernel_initializer,
                                            name='polar_kernel',
                                            regularizer=self.kernel_regularizer,
                                            constraint=self.kernel_constraint)
        if self.independent_north_pole:
            self.north_pole_kernel = self.add_weight(shape=kernel_shape,
                                                     initializer=self.kernel_initializer,
                                                     name='north_pole_kernel',
                                                     regularizer=self.kernel_regularizer,
                                                     constraint=self.kernel_constraint)
        if self.use_bias:
            self.equatorial_bias = self.add_weight(shape=(self.filters,),
                                                   initializer=self.bias_initializer,
                                                   name='equatorial_bias',
                                                   regularizer=self.bias_regularizer,
                                                   constraint=self.bias_constraint)
            self.polar_bias = self.add_weight(shape=(self.filters,),
                                              initializer=self.bias_initializer,
                                              name='polar_bias',
                                              regularizer=self.bias_regularizer,
                                              constraint=self.bias_constraint)
            if self.independent_north_pole:
                self.north_pole_bias = self.add_weight(shape=(self.filters,),
                                                       initializer=self.bias_initializer,
                                                       name='north_pole_bias',
                                                       regularizer=self.bias_regularizer,
                                                       constraint=self.bias_constraint)

        # Set input spec.
        self.input_spec = InputSpec(ndim=self.rank + 2,
                                    axes={channel_axis: input_dim})
        self.built = True

    def call(self, inputs, **kwargs):
        outputs = []
        channels_first = self.data_format == 'channels_first'

        # Equatorial faces
        for f in range(4):
            outputs.append(
                K.conv2d(
                    inputs[:, :, f, :, :] if channels_first else inputs[:, f],
                    self.equatorial_kernel,
                    strides=self.strides,
                    padding=self.padding,
                    data_format=self.data_format,
                    dilation_rate=self.dilation_rate
                )
            )
            if self.use_bias:
                outputs[f] = K.bias_add(
                    outputs[f],
                    self.equatorial_bias,
                    data_format=self.data_format
                )
            outputs[f] = K.expand_dims(outputs[f], 2 if channels_first else 1)

        # South pole face
        outputs.append(
            K.conv2d(
                inputs[:, :, 4, :, :] if channels_first else inputs[:, 4],
                self.polar_kernel,
                strides=self.strides,
                padding=self.padding,
                data_format=self.data_format,
                dilation_rate=self.dilation_rate
            )
        )
        if self.use_bias:
            outputs[4] = K.bias_add(
                outputs[4],
                self.polar_bias,
                data_format=self.data_format
            )
        outputs[4] = K.expand_dims(outputs[4], 2 if channels_first else 1)

        # North pole face
        if self.flip_north_pole:
            # Reverse the height dimension, 2 if channels_first otherwise 1
            outputs.append(
                K.conv2d(
                    K.reverse(inputs[:, :, 5, :, :], 2) if channels_first else K.reverse(inputs[:, 5], 1),
                    self.north_pole_kernel if self.independent_north_pole else self.polar_kernel,
                    strides=self.strides,
                    padding=self.padding,
                    data_format=self.data_format,
                    dilation_rate=self.dilation_rate
                )
            )
        else:
            outputs.append(
                K.conv2d(
                    inputs[:, :, 5, :, :] if channels_first else inputs[:, 5],
                    self.north_pole_kernel if self.independent_north_pole else self.polar_kernel,
                    strides=self.strides,
                    padding=self.padding,
                    data_format=self.data_format,
                    dilation_rate=self.dilation_rate
                )
            )
        if self.use_bias:
            outputs[5] = K.bias_add(
                outputs[5],
                self.north_pole_bias if self.independent_north_pole else self.polar_bias,
                data_format=self.data_format
            )
        if self.flip_north_pole:
            outputs[5] = K.reverse(outputs[5], 2 if channels_first else 1)
        outputs[5] = K.expand_dims(outputs[5], 2 if channels_first else 1)

        outputs = K.concatenate(outputs, axis=2 if channels_first else 1)

        if self.activation is not None:
            return self.activation(outputs)
        return outputs

    def compute_output_shape(self, input_shape):
        if self.data_format == 'channels_last':
            # batch, face, height, width, ...
            space = input_shape[2:4]
            new_space = []
            for i in range(len(space)):
                new_dim = conv_utils.conv_output_length(
                    space[i],
                    self.kernel_size[i],
                    padding=self.padding,
                    stride=self.strides[i],
                    dilation=self.dilation_rate[i])
                new_space.append(new_dim)
            return (input_shape[0], 6) + tuple(new_space) + (self.filters,)
        if self.data_format == 'channels_first':
            # batch, channels, ..., face, height, width
            space = input_shape[-2:]
            new_space = []
            for i in range(len(space)):
                new_dim = conv_utils.conv_output_length(
                    space[i],
                    self.kernel_size[i],
                    padding=self.padding,
                    stride=self.strides[i],
                    dilation=self.dilation_rate[i])
                new_space.append(new_dim)
            return (input_shape[0], self.filters, 6) + tuple(new_space)

    def get_config(self):
        config = {
            'filters': self.filters,
            'kernel_size': self.kernel_size,
            'strides': self.strides,
            'padding': self.padding,
            'data_format': self.data_format,
            'dilation_rate': self.dilation_rate,
            'activation': activations.serialize(self.activation),
            'use_bias': self.use_bias,
            'flip_north_pole': self.flip_north_pole,
            'independent_north_pole': self.independent_north_pole,
            'kernel_initializer': initializers.serialize(self.kernel_initializer),
            'bias_initializer': initializers.serialize(self.bias_initializer),
            'kernel_regularizer': regularizers.serialize(self.kernel_regularizer),
            'bias_regularizer': regularizers.serialize(self.bias_regularizer),
            'activity_regularizer':
                regularizers.serialize(self.activity_regularizer),
            'kernel_constraint': constraints.serialize(self.kernel_constraint),
            'bias_constraint': constraints.serialize(self.bias_constraint)
        }
        base_config = super(CubeSphereConv2D, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

class CubeSpherePadding2D(ZeroPadding3D):
    """
    Padding layer for 2D data on a cubed sphere. The requirements for using this layer are as follows:
    - The input data is 5-dimensional (batch, channels, height, width, 6)
    - Must follow "channels_first" order
    - The last dimension must have a length of 6 for the 6 faces of the cubed sphere
    - The last two faces (indices 4 and 5) are the polar faces

    Adapted from keras.layers.ZeroPadding3D by @jweyn

    # Arguments
        padding: int
        data_format: 'channels_first' or 'channels_last'
    """

    def __init__(self,
                 padding=(1, 1),
                 data_format='channels_first',
                 **kwargs):
        data_format = conv_utils.normalize_data_format(data_format)
        super(CubeSpherePadding2D, self).__init__(padding=padding,
                                                  data_format=data_format,
                                                  **kwargs)
        self.padding = ((0, 0),) + self.padding[1:]

    def call(self, inputs):
        p = self.padding[1][0]

        if self.data_format == 'channels_first':
            tr = (0, 1, 3, 2)

            # Pad the equatorial upper/lower boundaries and the polar upper/lower boundaries
            out = list()
            # Face 0
            out.append(K.expand_dims(
                K.concatenate([
                    inputs[:, :, 4, -p:, :],
                    inputs[:, :, 0],
                    inputs[:, :, 5, :p, :]
                ], axis=2), 2
            ))
            # Face 1
            out.append(K.expand_dims(
                K.concatenate([
                    tf.transpose(inputs[:, :, 4, ::-1, -p:], tr),
                    inputs[:, :, 1],
                    tf.transpose(K.reverse(inputs[:, :, 5, :, -p:], 3), tr)
                ], axis=2), 2
            ))
            # Face 2
            out.append(K.expand_dims(
                K.concatenate([
                    K.reverse(inputs[:, :, 4, :p, ::-1], 2),
                    inputs[:, :, 2],
                    K.reverse(inputs[:, :, 5, -p:, ::-1], 2)
                ], axis=2), 2
            ))
            # Face 3
            out.append(K.expand_dims(
                K.concatenate([
                    tf.transpose(K.reverse(inputs[:, :, 4, :, :p], 3), tr),
                    inputs[:, :, 3],
                    tf.transpose(inputs[:, :, 5, ::-1, :p], tr)
                ], axis=2), 2
            ))
            # Face 4 (south pole)
            out.append(K.expand_dims(
                K.concatenate([
                    K.reverse(inputs[:, :, 2, :p, ::-1], 2),
                    inputs[:, :, 4],
                    inputs[:, :, 0, :p, :]
                ], axis=2), 2
            ))
            # Face 5 (north pole)
            out.append(K.expand_dims(
                K.concatenate([
                    inputs[:, :, 0, -p:, :],
                    inputs[:, :, 5],
                    K.reverse(inputs[:, :, 2, -p:, ::-1], 2)
                ], axis=2), 2
            ))

            out1 = K.concatenate(out, axis=2)
            del out

            # Pad the equatorial periodic lateral boundaries and the polar left/right boundaries
            out = list()
            # Face 0
            out.append(K.expand_dims(
                K.concatenate([
                    out1[:, :, 3, :, -p:],
                    out1[:, :, 0],
                    out1[:, :, 1, :, :p]
                ], axis=3), 2
            ))
            # Face 1
            out.append(K.expand_dims(
                K.concatenate([
                    out1[:, :, 0, :, -p:],
                    out1[:, :, 1],
                    out1[:, :, 2, :, :p]
                ], axis=3), 2
            ))
            # Face 2
            out.append(K.expand_dims(
                K.concatenate([
                    out1[:, :, 1, :, -p:],
                    out1[:, :, 2],
                    out1[:, :, 3, :, :p]
                ], axis=3), 2
            ))
            # Face 3
            out.append(K.expand_dims(
                K.concatenate([
                    out1[:, :, 2, :, -p:],
                    out1[:, :, 3],
                    out1[:, :, 0, :, :p]
                ], axis=3), 2
            ))
            # Face 4
            out.append(K.expand_dims(
                K.concatenate([
                    tf.transpose(K.reverse(out[3][:, :, 0, p:2 * p, :], 2), tr),
                    out1[:, :, 4],
                    tf.transpose(out[1][:, :, 0, p:2 * p, ::-1], tr)
                ], axis=3), 2
            ))
            # Face 5
            out.append(K.expand_dims(
                K.concatenate([
                    tf.transpose(out[3][:, :, 0, -2 * p:-p, ::-1], tr),
                    out1[:, :, 5],
                    tf.transpose(K.reverse(out[1][:, :, 0, -2 * p:-p, :], 2), tr)
                ], axis=3), 2
            ))

            del out1
            outputs = K.concatenate(out, axis=2)
            del out
            return outputs

        else:  # channels_last
            tr = (0, 2, 1, 3)
            # Pad the equatorial upper/lower boundaries and the polar upper/lower boundaries
            out = list()
            # Face 0
            out.append(K.expand_dims(
                K.concatenate([
                    inputs[:, 4, -p:, :],
                    inputs[:, 0],
                    inputs[:, 5, :p, :]
                ], axis=1), 1
            ))
            # Face 1
            out.append(K.expand_dims(
                K.concatenate([
                    tf.transpose(inputs[:, 4, ::-1, -p:], tr),
                    inputs[:, 1],
                    tf.transpose(K.reverse(inputs[:, 5, :, -p:], 2), tr)
                ], axis=1), 1
            ))
            # Face 2
            out.append(K.expand_dims(
                K.concatenate([
                    K.reverse(inputs[:, 4, :p, ::-1], 1),
                    inputs[:, 2],
                    K.reverse(inputs[:, 5, -p:, ::-1], 1)
                ], axis=1), 1
            ))
            # Face 3
            out.append(K.expand_dims(
                K.concatenate([
                    tf.transpose(K.reverse(inputs[:, 4, :, :p], 2), tr),
                    inputs[:, 3],
                    tf.transpose(inputs[:, 5, ::-1, :p], tr)
                ], axis=1), 1
            ))
            # Face 4 (south pole)
            out.append(K.expand_dims(
                K.concatenate([
                    K.reverse(inputs[:, 2, :p, ::-1], 1),
                    inputs[:, 4],
                    inputs[:, 0, :p, :]
                ], axis=1), 1
            ))
            # Face 5 (north pole)
            out.append(K.expand_dims(
                K.concatenate([
                    inputs[:, 0, -p:, :],
                    inputs[:, 5],
                    K.reverse(inputs[:, 2, -p:, ::-1], 1)
                ], axis=1), 1
            ))

            out1 = K.concatenate(out, axis=1)
            del out

            # Pad the equatorial periodic lateral boundaries and the polar left/right boundaries
            out = list()
            # Face 0
            out.append(K.expand_dims(
                K.concatenate([
                    out1[:, 3, :, -p:],
                    out1[:, 0],
                    out1[:, 1, :, :p]
                ], axis=2), 1
            ))
            # Face 1
            out.append(K.expand_dims(
                K.concatenate([
                    out1[:, 0, :, -p:],
                    out1[:, 1],
                    out1[:, 2, :, :p]
                ], axis=2), 1
            ))
            # Face 2
            out.append(K.expand_dims(
                K.concatenate([
                    out1[:, 1, :, -p:],
                    out1[:, 2],
                    out1[:, 3, :, :p]
                ], axis=2), 1
            ))
            # Face 3
            out.append(K.expand_dims(
                K.concatenate([
                    out1[:, 2, :, -p:],
                    out1[:, 3],
                    out1[:, 0, :, :p]
                ], axis=2), 1
            ))
            # Face 4
            out.append(K.expand_dims(
                K.concatenate([
                    tf.transpose(K.reverse(out[3][:, 0, p:2 * p, :], 1), tr),
                    out1[:, 4],
                    tf.transpose(out[1][:, 0, p:2 * p, ::-1], tr)
                ], axis=2), 1
            ))
            # Face 5
            out.append(K.expand_dims(
                K.concatenate([
                    tf.transpose(out[3][:, 0, -2 * p:-p, ::-1], tr),
                    out1[:, 5],
                    tf.transpose(K.reverse(out[1][:, 0, -2 * p:-p, :], 1), tr)
                ], axis=2), 1
            ))

            del out1
            outputs = K.concatenate(out, axis=1)
            del out
            return outputs

# ==================================================================================================================== #
# Keras other layer classes and methods
# ==================================================================================================================== #

def slice_layer(start, end, step=None, axis=1):
    """
    Return a Lambda layer that performs slicing on a tensor.

    :param start: int: start index
    :param end: int: end index
    :param step: int: stepping parameter
    :param axis: int: axis along which to slice
    """
    if axis < 1:
        raise ValueError("'slice_layer' can only work on a specified axis > 0")

    def slice_func(x):
        slices = [slice(None)] * axis
        slices.append(slice(start, end, step))
        return x[tuple(slices)]

    return Lambda(slice_func)

class RowConnected2D(LocallyConnected2D):
    """Row-connected layer for 2D inputs.

    The `RowConnected2D` layer works similarly
    to the `Conv2D` layer, except that weights are shared only along rows,
    that is, a different set of filters is applied at each
    different row of the input.

    Adapted from keras.layers.local.LocallyConnected2D by @jweyn

    # Examples
    ```python
        # apply a 3x3 unshared weights convolution with 64 output filters
        # on a 32x32 image with `data_format="channels_last"`:
        model = Sequential()
        model.add(LocallyConnected2D(64, (3, 3), input_shape=(32, 32, 3)))
        # now model.output_shape == (None, 30, 30, 64)
        # notice that this layer will consume (30*30)*(3*3*3*64)
        # + (30*30)*64 parameters

        # add a 3x3 unshared weights convolution on top, with 32 output filters:
        model.add(LocallyConnected2D(32, (3, 3)))
        # now model.output_shape == (None, 28, 28, 32)
    ```

    # Arguments
        filters: Integer, the dimensionality of the output space
            (i.e. the number of output filters in the convolution).
        kernel_size: An integer or tuple/list of 2 integers, specifying the
            width and height of the 2D convolution window.
            Can be a single integer to specify the same value for
            all spatial dimensions.
        strides: An integer or tuple/list of 2 integers,
            specifying the strides of the convolution along the width and height.
            Can be a single integer to specify the same value for
            all spatial dimensions.
        padding: Currently only support `"valid"` (case-insensitive).
            `"same"` will be supported in future.
        data_format: A string,
            one of `channels_last` (default) or `channels_first`.
            The ordering of the dimensions in the inputs.
            `channels_last` corresponds to inputs with shape
            `(batch, height, width, channels)` while `channels_first`
            corresponds to inputs with shape
            `(batch, channels, height, width)`.
            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".
        activation: Activation function to use
            (see [activations](../activations.md)).
            If you don't specify anything, no activation is applied
            (ie. "linear" activation: `a(x) = x`).
        use_bias: Boolean, whether the layer uses a bias vector.
        kernel_initializer: Initializer for the `kernel` weights matrix
            (see [initializers](../initializers.md)).
        bias_initializer: Initializer for the bias vector
            (see [initializers](../initializers.md)).
        kernel_regularizer: Regularizer function applied to
            the `kernel` weights matrix
            (see [regularizer](../regularizers.md)).
        bias_regularizer: Regularizer function applied to the bias vector
            (see [regularizer](../regularizers.md)).
        activity_regularizer: Regularizer function applied to
            the output of the layer (its "activation").
            (see [regularizer](../regularizers.md)).
        kernel_constraint: Constraint function applied to the kernel matrix
            (see [constraints](../constraints.md)).
        bias_constraint: Constraint function applied to the bias vector
            (see [constraints](../constraints.md)).

    # Input shape
        4D tensor with shape:
        `(samples, channels, rows, cols)` if data_format='channels_first'
        or 4D tensor with shape:
        `(samples, rows, cols, channels)` if data_format='channels_last'.

    # Output shape
        4D tensor with shape:
        `(samples, filters, new_rows, new_cols)` if data_format='channels_first'
        or 4D tensor with shape:
        `(samples, new_rows, new_cols, filters)` if data_format='channels_last'.
        `rows` and `cols` values might have changed due to padding.
    """

    def __init__(self, *args, **kwargs):
        super(RowConnected2D, self).__init__(*args, **kwargs)

    def build(self, input_shape):
        if self.data_format == 'channels_last':
            input_row, input_col = input_shape[1:-1]
            input_filter = input_shape[3]
        else:
            input_row, input_col = input_shape[2:]
            input_filter = input_shape[1]
        if input_row is None or input_col is None:
            raise ValueError('The spatial dimensions of the inputs to '
                             ' a LocallyConnected2D layer '
                             'should be fully-defined, but layer received '
                             'the inputs shape ' + str(input_shape))
        output_row = conv_utils.conv_output_length(input_row, self.kernel_size[0],
                                                   self.padding, self.strides[0])
        output_col = conv_utils.conv_output_length(input_col, self.kernel_size[1],
                                                   self.padding, self.strides[1])
        self.output_row = output_row
        self.output_col = output_col
        self.kernel_shape = (
            output_row,
            self.kernel_size[0],
            self.kernel_size[1],
            input_filter,
            self.filters)
        self.kernel = self.add_weight(shape=self.kernel_shape,
                                      initializer=self.kernel_initializer,
                                      name='kernel',
                                      regularizer=self.kernel_regularizer,
                                      constraint=self.kernel_constraint)
        if self.use_bias:
            self.bias = self.add_weight(shape=(output_row, 1, self.filters),
                                        initializer=self.bias_initializer,
                                        name='bias',
                                        regularizer=self.bias_regularizer,
                                        constraint=self.bias_constraint)
        else:
            self.bias = None
        if self.data_format == 'channels_first':
            self.input_spec = InputSpec(ndim=4, axes={1: input_filter})
        else:
            self.input_spec = InputSpec(ndim=4, axes={-1: input_filter})
        self.built = True

    def call(self, inputs):
        output = row_conv2d(inputs,
                            self.kernel,
                            self.kernel_size,
                            self.strides,
                            (self.output_row, self.output_col),
                            self.data_format)

        if self.use_bias:
            output = K.bias_add(output, self.bias, data_format=self.data_format)

        output = self.activation(output)
        return output

def row_conv2d(inputs, kernel, kernel_size, strides, output_shape, data_format=None):
    """Apply 2D conv with weights shared only along rows.

    Adapted from K.local_conv2d by @jweyn

    # Arguments
        inputs: 4D tensor with shape:
                (batch_size, filters, new_rows, new_cols)
                if data_format='channels_first'
                or 4D tensor with shape:
                (batch_size, new_rows, new_cols, filters)
                if data_format='channels_last'.
        kernel: the row-shared weights for convolution,
                with shape (output_rows, kernel_size, input_channels, filters)
        kernel_size: a tuple of 2 integers, specifying the
                     width and height of the 2D convolution window.
        strides: a tuple of 2 integers, specifying the strides
                 of the convolution along the width and height.
        output_shape: a tuple with (output_row, output_col)
        data_format: the data format, channels_first or channels_last

    # Returns
        A 4d tensor with shape:
        (batch_size, filters, new_rows, new_cols)
        if data_format='channels_first'
        or 4D tensor with shape:
        (batch_size, new_rows, new_cols, filters)
        if data_format='channels_last'.

    # Raises
        ValueError: if `data_format` is neither
                    `channels_last` or `channels_first`.
    """
    data_format = K.normalize_data_format(data_format)

    stride_row, stride_col = strides
    output_row, output_col = output_shape

    out = []
    for i in range(output_row):
        # Slice the rows with the neighbors they need
        slice_row = slice(i * stride_row, i * stride_col + kernel_size[0])
        if data_format == 'channels_first':
            x = inputs[:, :, slice_row, :]  # batch, 16, 5, 144
        else:
            x = inputs[:, slice_row, :, :]  # batch, 5, 144, 16
        # Convolve, resulting in an array with only one row: batch, 1, 140, 6 or batch, 6, 1, 140
        x = K.conv2d(x, kernel[i], strides=strides, padding='valid', data_format=data_format)
        out.append(x)

    if data_format == 'channels_first':
        output = K.concatenate(out, axis=2)
    else:
        output = K.concatenate(out, axis=1)
    del x
    del out
    return output

# ==================================================================================================================== #
# Keras loss functions
# ==================================================================================================================== #

def latitude_weighted_loss(loss_function=mean_squared_error, lats=None, output_shape=(), axis=-2, weighting='cosine'):
    """
    Create a loss function that weights inputs by a function of latitude before calculating the loss.

    :param loss_function: method: Keras loss function to apply after the weighting
    :param lats: ndarray: 1-dimensional array of latitude coordinates
    :param output_shape: tuple: shape of expected model output
    :param axis: int: latitude axis in model output shape
    :param weighting: str: type of weighting to apply. Options are:
            cosine: weight by the cosine of the latitude (default)
            midlatitude: weight by the cosine of the latitude but also apply a 25% reduction to the equator and boost
                to the mid-latitudes
    :return: callable loss function
    """
    if weighting not in ['cosine', 'midlatitude']:
        raise ValueError("'weighting' must be one of 'cosine' or 'midlatitude'")
    if lats is not None:
        lat_tensor = K.zeros(lats.shape)
        lat_tensor.assign(K.cast_to_floatx(lats[:]))

        weights = K.cos(lat_tensor * np.pi / 180.)
        if weighting == 'midlatitude':
            weights = weights + 0.5 * K.pow(K.sin(lat_tensor * 2 * np.pi / 180.), 2.)

        weight_shape = output_shape[axis:]
        for d in weight_shape[1:]:
            weights = K.expand_dims(weights, axis=-1)
            weights = K.repeat_elements(weights, d, axis=-1)

    else:
        weights = K.ones(output_shape)

    def lat_loss(y_true, y_pred):
        return loss_function(y_true * weights, y_pred * weights)

    return lat_loss

def anomaly_correlation(y_true, y_pred, mean=0., regularize_mean='mse', reverse=True):
    """
    Calculate the anomaly correlation. FOR NOW, ASSUMES THAT THE CLIMATOLOGICAL MEAN IS 0, AND THEREFORE REQUIRES DATA
    TO BE SCALED TO REMOVE SPATIALLY-DEPENDENT MEAN.

    :param y_true: Tensor: target values
    :param y_pred: Tensor: model-predicted values
    :param mean: float: subtract this global mean from all predicted and target array values. IGNORED FOR NOW.
    :param regularize_mean: str or None: if not None, also penalizes a form of mean squared error:
        global: penalize differences in the global mean
        spatial: penalize differences in spatially-averaged mean (last two dimensions)
        mse: penalize the mean squared error
        mae: penalize the mean absolute error
    :param reverse: bool: if True, inverts the loss so that -1 is the target score
    :return: float: anomaly correlation loss
    """
    if regularize_mean is not None:
        assert regularize_mean in ['global', 'spatial', 'mse', 'mae']
    a = (K.mean(y_pred * y_true)
         / K.sqrt(K.mean(K.square(y_pred)) * K.mean(K.square(y_true))))
    if regularize_mean is not None:
        if regularize_mean == 'global':
            m = K.abs((K.mean(y_true) - K.mean(y_pred)) / K.mean(y_true))
        elif regularize_mean == 'spatial':
            m = K.mean(K.abs((K.mean(y_true, axis=[-2, -1]) - K.mean(y_pred, axis=[-2, -1]))
                             / K.mean(y_true, axis=[-2, -1])))
        elif regularize_mean == 'mse':
            m = mean_squared_error(y_true, y_pred)
        elif regularize_mean == 'mae':
            m = mean_absolute_error(y_true, y_pred)
    if reverse:
        if regularize_mean is not None:
            return m - a
        else:
            return -a
    else:
        if regularize_mean:
            return a - m
        else:
            return a

def anomaly_correlation_loss(mean=None, regularize_mean='mse', reverse=True):
    """
    Create a Keras loss function for anomaly correlation.

    :param mean: ndarray or None: if not None, must be an array with the same shape as the expected prediction, except
        that the first (batch) axis should have a dimension of 1.
    :param regularize_mean: str or None: if not None, also penalizes a form of mean squared error:
        global: penalize differences in the global mean
        spatial: penalize differences in spatially-averaged mean (last two dimensions)
        mse: penalize the mean squared error
        mae: penalize the mean absolute error
    :param reverse: bool: if True, inverts the loss so that -1 is the (minimized) target score. Must be True if
        regularize_mean is not None.
    :return: method: anomaly correlation loss function
    """
    if mean is not None:
        assert len(mean.shape) > 1
        assert mean.shape[0] == 1
        mean_tensor = K.variable(mean, name='anomaly_correlation_mean')

    if regularize_mean is not None:
        assert regularize_mean in ['global', 'spatial', 'mse', 'mae']
        reverse = True

    def acc_loss(y_true, y_pred):
        if mean is not None:
            a = (K.mean((y_pred - mean_tensor) * (y_true - mean_tensor))
                 / K.sqrt(K.mean(K.square((y_pred - mean_tensor))) * K.mean(K.square((y_true - mean_tensor)))))
        else:
            a = (K.mean(y_pred * y_true)
                 / K.sqrt(K.mean(K.square(y_pred)) * K.mean(K.square(y_true))))
        if regularize_mean is not None:
            if regularize_mean == 'global':
                m = K.abs((K.mean(y_true) - K.mean(y_pred)) / K.mean(y_true))
            elif regularize_mean == 'spatial':
                m = K.mean(K.abs((K.mean(y_true, axis=[-2, -1]) - K.mean(y_pred, axis=[-2, -1]))
                                 / K.mean(y_true, axis=[-2, -1])))
            elif regularize_mean == 'mse':
                m = mean_squared_error(y_true, y_pred)
            elif regularize_mean == 'mae':
                m = mean_absolute_error(y_true, y_pred)
        if reverse:
            if regularize_mean is not None:
                return m - a
            else:
                return -a
        else:
            if regularize_mean:
                return a - m
            else:
                return a

    return acc_loss

# Compatibility names
# lat_loss = latitude_weighted_loss()
# acc_loss = anomaly_correlation_loss()

class AsymptoteReLU(Layer):
    """
    Rectified Linear Unit activation function.

    Returns:
      `f(x) = x / (beta * x + 1)` for `x >= 0`,
      `f(x) = alpha * x` otherwise.

      Input shape:
        Arbitrary. Use the keyword argument `input_shape`
        (tuple of integers, does not include the samples axis)
        when using this layer as the first layer in a model.

      Output shape:
        Same shape as the input.

      Arguments:
        alpha: Float >= 0. Negative slope coefficient.
        beta: Float >= 0. Positive asymptote coefficient.
    """

    def __init__(self, alpha=0., beta=1., **kwargs):
        super(AsymptoteReLU, self).__init__(**kwargs)
        if alpha < 0.:
            raise ValueError('alpha of AsymptoteReLU layer '
                             'cannot be negative value: ' + str(alpha))
        if beta <= 0.:
            raise ValueError('beta of AsymptoteReLU layer '
                             'cannot be zero or negative value: ' + str(beta))

        self.support_masking = True
        self.alpha = K.cast_to_floatx(alpha)
        self.beta = K.cast_to_floatx(beta)
        self.one = K.cast_to_floatx(1.)

    def call(self, inputs, **kwargs):
        return K.relu(inputs / (K.abs(self.beta * inputs) + self.one), alpha=self.alpha)

    def get_config(self):
        config = {
            'alpha': self.alpha,
            'beta': self.beta
        }
        base_config = super(AsymptoteReLU, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

    def compute_output_shape(self, input_shape):
        return input_shape

# ==================================================================================================================== #
# PyTorch classes
# ==================================================================================================================== #

class TorchReshape(object):
    def __init__(self, shape):
        if not isinstance(shape, tuple):
            raise ValueError("'shape' must be a tuple of integers")
        self.shape = shape

    def __call__(self, x):
        return x.view(*self.shape)


# -*- coding: utf-8 -*-
from tensorflow.python.keras.models import Model
from math import sqrt, log, exp
from time import time
import numpy as np
import random

class Node:
    def __init__(self, playerJustMoved, move=None, parent=None, proba=1, untriedMoves=[]):
        self.move = move
        self.parentNode = parent
        self.childNodes = []
        self.wins = 0
        self.visits = 0
        self.quotient = 0
        self.proba = proba
        self.untriedMoves = untriedMoves
        self.playerJustMoved = playerJustMoved
        self.isevaluate = False

    def __repr__(self):
        return "[M:" + str(self.move) + " W/V:" + str(round(self.wins,1)) + "/" + str(self.visits) + " P:" + str(round(self.proba*100,1))\
            + " Q:" + str(round(self.quotient*100,1)) + " U:" + str(self.untriedMoves) + "]"
    def TreeToString(self, indent):
        s = self.IndentString(indent) + self.__repr__()
        if self.childNodes != []:
            for c in self.childNodes:
                if c.visits > 0:
                    s += c.TreeToString(indent+1)
        return s
    def IndentString(self,indent):
        s = "\n"
        for i in range (1,indent+1):
            s += "| "
        return s

class AI:
    count = 1
    def __init__(self, itermax=300, alpha=1, name="", model=None, formula="", use_subtree=False, stochastic=False, temperature=0.5, add_data=True):
        self.turn = 0
        self.add_data = add_data
        self.itermax = itermax
        self.alpha = alpha
        self.use_subtree = use_subtree
        self.name = "a" + str(AI.count) if name == "" else name
        if type(model) is Model:
            self.method, self.model, self.alpha = neural_network_tree_search, model, model.output_shape[0][1]*alpha
        else:
            self.method, self.model, self.alpha = monte_carlo_tree_search, None, alpha
        if formula == "ucb1":
            self.formula = _ucb1
        else:
            self.formula = _ucb2
        self.selection, self.temperature = (_stochastic, temperature) if stochastic else (_competitive, 0)
        AI.count += 1

    def __repr__(self):
        res = "I:" + str(self.itermax) + ", A:" + str(self.alpha) + ", F:" + self.formula.__name__ + ", M:" + self.method.__name__ + ", N:" + self.name
        if self.method is neural_network_tree_search: res += "_" + self.model.name
        res += ", S:" + self.selection.__name__
        if self.selection is _stochastic: res += ", T:" + str(round(self.temperature,2))
        return res + ", U:" + str(self.use_subtree) + ", D:" + str(self.add_data)

    def getnode(self, rootenv, rootnode):
        return self.selection(self.method(rootenv, rootnode, self.itermax, self.alpha, self.formula, self.model), self.temperature)

    def updatenode(self, node, move):
        if (type(node) is Node) and self.use_subtree:
            for c in node.childNodes:
                if c.move == move: return c
        return None

# Upper Confident Bounce Formula
def _ucb1(node, Visits, alpha): return node.quotient + alpha*node.proba*sqrt(Visits+1)/(node.visits+1)
def _ucb2(node, Visits, alpha): return node.quotient + alpha*node.proba*sqrt(2*log(Visits+1)/(node.visits+1))

# Move Selection Method
def _competitive(node, temperature):
    return sorted(node.childNodes, key = lambda c: c.visits)[-1]
def _stochastic(node, temperature):
    cursor, step, inverse = 0, 0, 1/temperature
    for c in node.childNodes: cursor += c.visits**inverse
    cursor = cursor*random.random()
    for c in node.childNodes:
        step += c.visits**inverse
        if cursor <= step: return c
    return None

# Node Builder Method
def monte_carlo_tree_search(rootenv, rootnode, itermax, alpha, formula, rnn):
    if type(rootnode) is not Node:
        rootnode = Node(playerJustMoved=rootenv.playerJustMoved, untriedMoves=rootenv.getmoves())

    for i in range(itermax):
        node = rootnode
        env = rootenv.clone()

        #select
        while node.untriedMoves == [] and node.childNodes != []:
            node = sorted(node.childNodes, key = lambda c: formula(c, node.visits, alpha))[-1]
            env.domove(node.move)

        #expand
        if node.untriedMoves != []:
            m = random.choice(node.untriedMoves)
            env.domove(m)
            n = Node(playerJustMoved=env.playerJustMoved, move=m, parent=node, untriedMoves=env.getmoves())
            node.untriedMoves.remove(m)
            node.childNodes.append(n)
            node = n

        #rollout - this can often be made orders of magnitude quicker using a env.getrandommove() function
        moves = env.getmoves()
        while moves != []:
            env.domove(random.choice(moves))
            moves = env.getmoves()

        #backpropagate
        value, pjm = env.value(), env.playerJustMoved
        while node != None:
            node.wins += 0.5 + (value - 0.5)*pjm*node.playerJustMoved
            node.visits += 1
            node.quotient = node.wins/node.visits
            node = node.parentNode

    return rootnode
def neural_network_tree_search(rootenv, rootnode, itermax, alpha, formula, rnn):
    if type(rootnode) is not Node:
        rootnode = Node(playerJustMoved=rootenv.playerJustMoved)

    if rootnode.childNodes == []:
        moves = rootenv.getmoves()
        if moves != []:
            predicts = rnn.predict(np.array([rootenv.x()]))
            tab = predicts[0][0,:]
            for m in moves:
                rootnode.childNodes.append(Node(playerJustMoved=-1*rootnode.playerJustMoved, move=m, parent=rootnode, proba=tab[m]))
    i = 0

    while i < itermax:
        node = rootnode
        env = rootenv.clone()

        #select
        while node.isevaluate:
            node = sorted(node.childNodes, key = lambda c: formula(c, node.visits, alpha))[-1]
            env.domove(node.move)

        #expand & evaluate
        childnodes = node.childNodes
        k, value, pjm = 1, env.value(), env.playerJustMoved
        if childnodes != []:
            xlist, mlist, vlist = [], [], []
            for c in childnodes:
                _env = env.clone()
                _env.domove(c.move)
                xlist.append(_env.x())
                vlist.append(_env.value())
                mlist.append(_env.getmoves())
            predicts = rnn.predict(np.array(xlist))

            for j in range(len(childnodes)):
                c = childnodes[j]
                moves = mlist[j]
                if moves != []:
                    tab = predicts[0][j,:]
                    vlist[j] = predicts[1][j,0]
                    for m in moves:
                        c.childNodes.append(Node(playerJustMoved=-1*c.playerJustMoved, move=m, parent=c, proba=tab[m]))
                c.quotient = c.wins = vlist[j]
                c.visits = 1
            k, value, pjm = len(vlist), sum(vlist), -1*pjm
            node.isevaluate = True

        #backpropagate
        while node != None:
            node.wins += 0.5*k + (value - 0.5*k)*pjm*node.playerJustMoved
            node.visits += k
            node.quotient = node.wins/node.visits
            node = node.parentNode
        i += k

    return rootnode


# -*- coding: utf-8 -*-
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

# Semilla
np.random.seed(50)

# Datos
dir_trab = r'C:\Users\coyol\Desktop\bayesian_python\py_programs'
data_all = pd.read_csv(dir_trab  + r'\datos\icecore.csv', sep=',' ,index_col=0)

# Se realizan los graficos
tmp_std = (data_all.tmp - np.mean(data_all.tmp)) / np.std(data_all.tmp)
co2_std = (data_all.co2 - np.mean(data_all.co2)) / np.std(data_all.co2)

## grafico 10.7
# Grafico 1
plt.plot(data_all.year, tmp_std, label='temp')
plt.plot(data_all.year, co2_std, label='CO2')
plt.legend()
plt.show()

# Grafico 2
plt.plot(data_all.co2, data_all.tmp, 'o')


import numpy as np
import matplotlib.pyplot as plt
import math
import random
from HRNN2 import HRNN2
np.random.seed(10000)
np.set_printoptions(threshold=np.nan)

def missmatch(pattern,activation):
    diff = np.abs(pattern-activation)
    return np.sum(np.where(np.sum(diff,axis = 1)>0,0,1))

def add_noise(pattern, amt_of_noise):
    # Calculate the percentages
    amt_of_noise = int(amt_of_noise * 100)
    # Pick out a row vector of indices
    index = np.random.choice(100, amt_of_noise, replace=False)
    # Add some noise to the patterns by flipping the randomly selected units
    pattern[index] = pattern[index] + 1
    return np.copy(pattern)

def main():
    random_data = np.random.randint(1,size=(100,100))
    for i, pattern in enumerate(random_data):
        random_data[i, :] = add_noise(pattern, 0.10)
    correct_classed_patterns = [None] * random_data.shape[0]
    biases = [0, .001, .01, .1, 1, 10]
    for bias in biases:
        for p in range(2, 30):
            selected_patterns = random_data[0:p, :]
            print(selected_patterns)
            Network = HRNN2(selected_patterns)
            output = Network.recall(selected_patterns, 1, True, bias)
            correct_classed_patterns[p-2] = missmatch(Network.activations, np.copy(selected_patterns))
            print(correct_classed_patterns[p-2])
        plt.plot(correct_classed_patterns)
        plt.show()

    # random_data = np.where((np.random.rand(300,100) * 10) - 2.5 >=0,1,-1)
    # print(random_data)
    # random_data = np.where(npv.random.rand(2(300,100))==0,-1,1)
    # print(random_data)
#     for p in range(2,random_data.shape[0]):
#         selected_patterns = random_data[0:p, :]
#         # Initialize the Hopfield RNN with three selected patterns from the pict dataset
#         Network = HRNN2(selected_patterns)
#         new_selected_patterns = add_noise(np.copy(selected_patterns), 0.0)
#         epochs = 1 #int(math.log2(len(selected_patterns)))
#         # Initialize a subplot object
#         Network.recall(np.copy(new_selected_patterns[0:p, :]), epochs, True)
#         correct_classed_patterns[p-2] = missmatch(Network.activations, np.copy(selected_patterns))
#
main()


"""
Problem Description
The program takes a list and swaps the first and last value of the list.

Problem Solution
1. Take the number of elements in the list and store it in a variable.
2. Accept the values into the list using a for loop and insert them into the list.
3. Using a temporary variable, switch the first and last element in the list.
4. Print the newly formed list.
5. Exit.
"""
print("---------------------------METHOD 1----------------------")
a=[]
n=int(input("Enter the number of elements in the list: "))
for i in range(n):
	data=int(input(f"Enter elements of the array at position {i}: "))
	a.append(data)
temp=a[0]
a[0]=a[n-1]
a[n-1]=temp
print(f"New list: {a}")

print("---------------------METHOD 2-----------------------------")
"""The last element of the list can be referred as list[-1]. Therefore, we can simply swap list[0] with list[-1]"""
#Swap function
def swaplist(newList):
	newList[0], newList[-1]=newList[-1],newList[0]
	return newList
newList=[]
n=int(input("Enter the number of elements in the list: "))
for  i in range(n):
	data=int(input(f"Enter elements of the array at position {i}:  "))
	newList.append(data)
print(swaplist(newList))

print("---------------METHOD 3-------------------------")
"""Swap the first and last element is using tuple variable.
Store the first and last element as a pair in a tuple variable,
say get, and unpack those elements with first and last element in that list.
Now, the First and last values in that list are swapped."""
#Swap function
def swaplist(list):
	#Storing the first and last element
	#as a pair in a tuple variable get
    get=list[-1],list[0]
    #unpacking those elements
    list[0],list[-1]=get
    return list
newList=[]
n=int(input("Enter the number of elements in the list: "))
for  i in range(n):
	data=int(input(f"Enter elements of the array at position {i}:  "))
	newList.append(data)
print(swaplist(newList))

print("-------------METHOD 4--------------------------")
"""Using * operand.
This operand proposes a change to iterable unpacking syntax,
allowing to specify a “catch-all” name which will be assigned
a list of all items not assigned to a “regular” name."""
#usage of *operand
def swaplist(list):
	start, *mid, end = list
	list=[end, *mid, start]
	return list
newList=[]
n=int(input("Enter the number of elements in the list: "))
for  i in range(n):
	data=int(input(f"Enter elements of the array at position {i}:  "))
	newList.append(data)
print(swaplist(newList))

print("----------------METHOD 5-----------------------")
"""
Swap the first and last elements is to use inbuilt function list.pop().
Pop the first element and store it in a variable.
Similarily pop the last element and store it in another variable.
Now insert the two popped element at each other’s original position.
"""
def swaplist(list):
	first=list.pop(0)
	last=list.pop(-1)
	list.insert(0,last)
	list.append(first)
	return list

newList=[]
n=int(input("Enter the number of elements in the list: "))
for  i in range(n):
	data=int(input(f"Enter elements of the array at position {i}:  "))
	newList.append(data)
print(swaplist(newList))

list/Swap the First and Last Value of a List.py


#!/usr/bin/env python

import csv
import logging
import math
import os
import subprocess
import urllib
from pathlib import Path

import pandas as pd
from sox import Transformer

import swifter
from deepspeech_training.util.importers import get_importers_parser, get_validate_label

__version__ = "0.1.0"
_logger = logging.getLogger(__name__)

MAX_SECS = 10
BITDEPTH = 16
N_CHANNELS = 1
SAMPLE_RATE = 16000

DEV_PERCENTAGE = 0.10
TRAIN_PERCENTAGE = 0.80

def parse_args(args):
    """Parse command line parameters
    Args:
      args ([str]): Command line parameters as list of strings
    Returns:
      :obj:`argparse.Namespace`: command line parameters namespace
    """
    parser = get_importers_parser(description="Imports GramVaani data for Deep Speech")
    parser.add_argument(
        "--version",
        action="version",
        version="GramVaaniImporter {ver}".format(ver=__version__),
    )
    parser.add_argument(
        "-v",
        "--verbose",
        action="store_const",
        required=False,
        help="set loglevel to INFO",
        dest="loglevel",
        const=logging.INFO,
    )
    parser.add_argument(
        "-vv",
        "--very-verbose",
        action="store_const",
        required=False,
        help="set loglevel to DEBUG",
        dest="loglevel",
        const=logging.DEBUG,
    )
    parser.add_argument(
        "-c",
        "--csv_filename",
        required=True,
        help="Path to the GramVaani csv",
        dest="csv_filename",
    )
    parser.add_argument(
        "-t",
        "--target_dir",
        required=True,
        help="Directory in which to save the importer GramVaani data",
        dest="target_dir",
    )
    return parser.parse_args(args)

def setup_logging(level):
    """Setup basic logging
    Args:
      level (int): minimum log level for emitting messages
    """
    format = "[%(asctime)s] %(levelname)s:%(name)s:%(message)s"
    logging.basicConfig(
        level=level, stream=sys.stdout, format=format, datefmt="%Y-%m-%d %H:%M:%S"
    )

class GramVaaniCSV:
    """GramVaaniCSV representing a GramVaani dataset.
    Args:
      csv_filename (str): Path to the GramVaani csv
    Attributes:
        data (:class:`pandas.DataFrame`): `pandas.DataFrame` Containing the GramVaani csv data
    """

    def __init__(self, csv_filename):
        self.data = self._parse_csv(csv_filename)

    def _parse_csv(self, csv_filename):
        _logger.info("Parsing csv file...%s", os.path.abspath(csv_filename))
        data = pd.read_csv(
            os.path.abspath(csv_filename),
            names=[
                "piece_id",
                "audio_url",
                "transcript_labelled",
                "transcript",
                "labels",
                "content_filename",
                "audio_length",
                "user_id",
            ],
            usecols=["audio_url", "transcript", "audio_length"],
            skiprows=[0],
            engine="python",
            encoding="utf-8",
            quotechar='"',
            quoting=csv.QUOTE_ALL,
        )
        data.dropna(inplace=True)
        _logger.info("Parsed %d lines csv file." % len(data))
        return data

class GramVaaniDownloader:
    """GramVaaniDownloader downloads a GramVaani dataset.
    Args:
      gram_vaani_csv (GramVaaniCSV): A GramVaaniCSV representing the data to download
      target_dir (str): The path to download the data to
    Attributes:
        data (:class:`pandas.DataFrame`): `pandas.DataFrame` Containing the GramVaani csv data
    """

    def __init__(self, gram_vaani_csv, target_dir):
        self.target_dir = target_dir
        self.data = gram_vaani_csv.data

    def download(self):
        """Downloads the data associated with this instance
        Return:
          mp3_directory (os.path): The directory into which the associated mp3's were downloaded
        """
        mp3_directory = self._pre_download()
        self.data.swifter.apply(
            func=lambda arg: self._download(*arg, mp3_directory), axis=1, raw=True
        )
        return mp3_directory

    def _pre_download(self):
        mp3_directory = os.path.join(self.target_dir, "mp3")
        if not os.path.exists(self.target_dir):
            _logger.info("Creating directory...%s", self.target_dir)
            os.mkdir(self.target_dir)
        if not os.path.exists(mp3_directory):
            _logger.info("Creating directory...%s", mp3_directory)
            os.mkdir(mp3_directory)
        return mp3_directory

    def _download(self, audio_url, transcript, audio_length, mp3_directory):
        if audio_url == "audio_url":
            return
        mp3_filename = os.path.join(mp3_directory, os.path.basename(audio_url))
        if not os.path.exists(mp3_filename):
            _logger.debug("Downloading mp3 file...%s", audio_url)
            urllib.request.urlretrieve(audio_url, mp3_filename)
        else:
            _logger.debug("Already downloaded mp3 file...%s", audio_url)

class GramVaaniConverter:
    """GramVaaniConverter converts the mp3's to wav's for a GramVaani dataset.
    Args:
      target_dir (str): The path to download the data from
      mp3_directory (os.path): The path containing the GramVaani mp3's
    Attributes:
        target_dir (str): The target directory passed as a command line argument
        mp3_directory (os.path): The path containing the GramVaani mp3's
    """

    def __init__(self, target_dir, mp3_directory):
        self.target_dir = target_dir
        self.mp3_directory = Path(mp3_directory)

    def convert(self):
        """Converts the mp3's associated with this instance to wav's
        Return:
          wav_directory (os.path): The directory into which the associated wav's were downloaded
        """
        wav_directory = self._pre_convert()
        for mp3_filename in self.mp3_directory.glob("**/*.mp3"):
            wav_filename = os.path.join(
                wav_directory,
                os.path.splitext(os.path.basename(mp3_filename))[0] + ".wav",
            )
            if not os.path.exists(wav_filename):
                _logger.debug(
                    "Converting mp3 file %s to wav file %s"
                    % (mp3_filename, wav_filename)
                )
                transformer = Transformer()
                transformer.convert(
                    samplerate=SAMPLE_RATE, n_channels=N_CHANNELS, bitdepth=BITDEPTH
                )
                transformer.build(str(mp3_filename), str(wav_filename))
            else:
                _logger.debug(
                    "Already converted mp3 file %s to wav file %s"
                    % (mp3_filename, wav_filename)
                )
        return wav_directory

    def _pre_convert(self):
        wav_directory = os.path.join(self.target_dir, "wav")
        if not os.path.exists(self.target_dir):
            _logger.info("Creating directory...%s", self.target_dir)
            os.mkdir(self.target_dir)
        if not os.path.exists(wav_directory):
            _logger.info("Creating directory...%s", wav_directory)
            os.mkdir(wav_directory)
        return wav_directory

class GramVaaniDataSets:
    def __init__(self, target_dir, wav_directory, gram_vaani_csv):
        self.target_dir = target_dir
        self.wav_directory = wav_directory
        self.csv_data = gram_vaani_csv.data
        self.raw = pd.DataFrame(columns=["wav_filename", "wav_filesize", "transcript"])
        self.valid = pd.DataFrame(
            columns=["wav_filename", "wav_filesize", "transcript"]
        )
        self.train = pd.DataFrame(
            columns=["wav_filename", "wav_filesize", "transcript"]
        )
        self.dev = pd.DataFrame(columns=["wav_filename", "wav_filesize", "transcript"])
        self.test = pd.DataFrame(columns=["wav_filename", "wav_filesize", "transcript"])

    def create(self):
        self._convert_csv_data_to_raw_data()
        self.raw.index = range(len(self.raw.index))
        self.valid = self.raw[self._is_valid_raw_rows()]
        self.valid = self.valid.sample(frac=1).reset_index(drop=True)
        train_size, dev_size, test_size = self._calculate_data_set_sizes()
        self.train = self.valid.loc[0:train_size]
        self.dev = self.valid.loc[train_size : train_size + dev_size]
        self.test = self.valid.loc[
            train_size + dev_size : train_size + dev_size + test_size
        ]

    def _convert_csv_data_to_raw_data(self):
        self.raw[["wav_filename", "wav_filesize", "transcript"]] = self.csv_data[
            ["audio_url", "transcript", "audio_length"]
        ].swifter.apply(
            func=lambda arg: self._convert_csv_data_to_raw_data_impl(*arg),
            axis=1,
            raw=True,
        )
        self.raw.reset_index()

    def _convert_csv_data_to_raw_data_impl(self, audio_url, transcript, audio_length):
        if audio_url == "audio_url":
            return pd.Series(["wav_filename", "wav_filesize", "transcript"])
        mp3_filename = os.path.basename(audio_url)
        wav_relative_filename = os.path.join(
            "wav", os.path.splitext(os.path.basename(mp3_filename))[0] + ".wav"
        )
        wav_filesize = os.path.getsize(
            os.path.join(self.target_dir, wav_relative_filename)
        )
        transcript = validate_label(transcript)
        if None == transcript:
            transcript = ""
        return pd.Series([wav_relative_filename, wav_filesize, transcript])

    def _is_valid_raw_rows(self):
        is_valid_raw_transcripts = self._is_valid_raw_transcripts()
        is_valid_raw_wav_frames = self._is_valid_raw_wav_frames()
        is_valid_raw_row = [
            (is_valid_raw_transcript & is_valid_raw_wav_frame)
            for is_valid_raw_transcript, is_valid_raw_wav_frame in zip(
                is_valid_raw_transcripts, is_valid_raw_wav_frames
            )
        ]
        series = pd.Series(is_valid_raw_row)
        return series

    def _is_valid_raw_transcripts(self):
        return pd.Series([bool(transcript) for transcript in self.raw.transcript])

    def _is_valid_raw_wav_frames(self):
        transcripts = [str(transcript) for transcript in self.raw.transcript]
        wav_filepaths = [
            os.path.join(self.target_dir, str(wav_filename))
            for wav_filename in self.raw.wav_filename
        ]
        wav_frames = [
            int(
                subprocess.check_output(
                    ["soxi", "-s", wav_filepath], stderr=subprocess.STDOUT
                )
            )
            for wav_filepath in wav_filepaths
        ]
        is_valid_raw_wav_frames = [
            self._is_wav_frame_valid(wav_frame, transcript)
            for wav_frame, transcript in zip(wav_frames, transcripts)
        ]
        return pd.Series(is_valid_raw_wav_frames)

    def _is_wav_frame_valid(self, wav_frame, transcript):
        is_wav_frame_valid = True
        if int(wav_frame / SAMPLE_RATE * 1000 / 10 / 2) < len(str(transcript)):
            is_wav_frame_valid = False
        elif wav_frame / SAMPLE_RATE > MAX_SECS:
            is_wav_frame_valid = False
        return is_wav_frame_valid

    def _calculate_data_set_sizes(self):
        total_size = len(self.valid)
        dev_size = math.floor(total_size * DEV_PERCENTAGE)
        train_size = math.floor(total_size * TRAIN_PERCENTAGE)
        test_size = total_size - (train_size + dev_size)
        return (train_size, dev_size, test_size)

    def save(self):
        datasets = ["train", "dev", "test"]
        for dataset in datasets:
            self._save(dataset)

    def _save(self, dataset):
        dataset_path = os.path.join(self.target_dir, dataset + ".csv")
        dataframe = getattr(self, dataset)
        dataframe.to_csv(
            dataset_path,
            index=False,
            encoding="utf-8",
            escapechar="\\",
            quoting=csv.QUOTE_MINIMAL,
        )

def main(args):
    """Main entry point allowing external calls
    Args:
      args ([str]): command line parameter list
    """
    args = parse_args(args)
    validate_label = get_validate_label(args)
    setup_logging(args.loglevel)
    _logger.info("Starting GramVaani importer...")
    _logger.info("Starting loading GramVaani csv...")
    csv = GramVaaniCSV(args.csv_filename)
    _logger.info("Starting downloading GramVaani mp3's...")
    downloader = GramVaaniDownloader(csv, args.target_dir)
    mp3_directory = downloader.download()
    _logger.info("Starting converting GramVaani mp3's to wav's...")
    converter = GramVaaniConverter(args.target_dir, mp3_directory)
    wav_directory = converter.convert()
    datasets = GramVaaniDataSets(args.target_dir, wav_directory, csv)
    datasets.create()
    datasets.save()
    _logger.info("Finished GramVaani importer...")

main(sys.argv[1:])