from collections import deque
from collections import OrderedDict
from difflib import get_close_matches
import functools
import inspect
from itertools import chain
import math
from warnings import warn

import numpy as np
import tensorflow as tf
from tensorflow.python.training import moving_averages


class G:
    pass


G._n_layers = 0


def create_param(spec, shape, name, trainable=True, regularizable=True):
    if not hasattr(spec, '__call__'):
        assert isinstance(spec, (tf.Tensor, tf.Variable))
        return spec
    assert hasattr(spec, '__call__')
    if regularizable:
        # use the default regularizer
        regularizer = None
    else:
        # do not regularize this variable
        def f(_):
            return tf.constant(0.)

        regularizer = f
    return tf.compat.v1.get_variable(
        name=name,
        shape=shape,
        initializer=spec,
        trainable=trainable,
        regularizer=regularizer,
        dtype=tf.float32)


def as_tuple(x, n, t=None):
    try:
        x = tuple(x)
    except TypeError:
        x = (x, ) * n

    if (t is not None) and not all(isinstance(v, t) for v in x):
        raise TypeError('expected a single value or an iterable '
                        'of {0}, got {1} instead'.format(t.__name__, x))

    if len(x) != n:
        raise ValueError('expected a single value or an iterable '
                         'with length {0}, got {1} instead'.format(n, x))

    return x


def conv_output_length(input_length, filter_size, stride, pad=0):
    """Helper function to compute the output size of a convolution operation
    This function computes the length along a single axis, which corresponds
    to a 1D convolution. It can also be used for convolutions with higher
    dimensionalities by using it individually for each axis.
    Parameters
    ----------
    input_length : int or None
        The size of the input.
    filter_size : int
        The size of the filter.
    stride : int
        The stride of the convolution operation.
    pad : int, 'full' or 'same' (default: 0)
        By default, the convolution is only computed where the input and the
        filter fully overlap (a valid convolution). When ``stride=1``, this
        yields an output that is smaller than the input by ``filter_size - 1``.
        The `pad` argument allows you to implicitly pad the input with zeros,
        extending the output size.
        A single integer results in symmetric zero-padding of the given size on
        both borders.
        ``'full'`` pads with one less than the filter size on both sides. This
        is equivalent to computing the convolution wherever the input and the
        filter overlap by at least one position.
        ``'same'`` pads with half the filter size on both sides (one less on
        the second side for an even filter size). When ``stride=1``, this
        results in an output size equal to the input size.
    Returns
    -------
    int or None
        The output size corresponding to the given convolution parameters, or
        ``None`` if `input_size` is ``None``.
    Raises
    ------
    ValueError
        When an invalid padding is specified, a `ValueError` is raised.
    """
    if input_length is None:
        return None
    if pad == 'valid':
        output_length = input_length - filter_size + 1
    elif pad == 'full':
        output_length = input_length + filter_size - 1
    elif pad == 'same':
        output_length = input_length
    elif isinstance(pad, int):
        output_length = input_length + 2 * pad - filter_size + 1
    else:
        raise ValueError('Invalid pad: {0}'.format(pad))

    # This is the integer arithmetic equivalent to
    # np.ceil(output_length / stride)
    output_length = (output_length + stride - 1) // stride

    return output_length


class Layer:
    def __init__(self,
                 incoming,
                 name=None,
                 variable_reuse=None,
                 weight_normalization=False,
                 **kwargs):
        if isinstance(incoming, tuple):
            self.input_shape = incoming
            self.input_layer = None
        else:
            self.input_shape = incoming.output_shape
            self.input_layer = incoming
        self.params = OrderedDict()
        self.weight_normalization = weight_normalization

        if name is None:
            name = '%s_%d' % (type(self).__name__, G._n_layers)
            G._n_layers += 1

        self.name = name
        self.variable_reuse = variable_reuse
        self.get_output_kwargs = []
        self.variable_scope = tf.compat.v1.variable_scope(
            name, reuse=self.variable_reuse)

        if any(d is not None and d <= 0 for d in self.input_shape):
            raise ValueError(
                ('Cannot create Layer with a non-positive input_shape '
                 'dimension. input_shape=%r, self.name=%r') %
                (self.input_shape, self.name))

    @property
    def output_shape(self):
        shape = self.get_output_shape_for(self.input_shape)
        if any(isinstance(s, (tf.Variable, tf.Tensor)) for s in shape):
            raise ValueError('%s returned a symbolic output shape from its '
                             'get_output_shape_for() method: %r. This is not '
                             'allowed; shapes must be tuples of integers for '
                             'fixed-size dimensions and Nones for variable '
                             'dimensions.' % (self.__class__.__name__, shape))
        return shape

    def get_output_shape_for(self, input_shape):
        raise NotImplementedError

    def get_output_for(self, input, **kwargs):
        raise NotImplementedError

    def add_param_plain(self, spec, shape, name, **tags):
        with self.variable_scope:
            tags['trainable'] = tags.get('trainable', True)
            tags['regularizable'] = tags.get('regularizable', True)
            param = create_param(spec, shape, name, **tags)
            self.params[param] = set(
                tag for tag, value in list(tags.items()) if value)
            return param

    def add_param(self, spec, shape, name, **kwargs):
        param = self.add_param_plain(spec, shape, name, **kwargs)
        if name is not None and name.startswith(
                'W') and self.weight_normalization:
            # Hacky: check if the parameter is a weight matrix. If so, apply
            # weight normalization
            if len(param.get_shape()) == 2:
                v = param
                g = self.add_param_plain(
                    tf.ones_initializer(), (shape[1], ), name=name + '_wn/g')
                param = v * (tf.reshape(g, (1, -1)) / tf.sqrt(
                    tf.reduce_sum(tf.square(v), 0, keep_dims=True)))
            elif len(param.get_shape()) == 4:
                v = param
                g = self.add_param_plain(
                    tf.ones_initializer(), (shape[3], ), name=name + '_wn/g')
                param = v * (tf.reshape(g, (1, 1, 1, -1)) / tf.sqrt(
                    tf.reduce_sum(tf.square(v), [0, 1, 2], keep_dims=True)))
            else:
                raise NotImplementedError
        return param

    def get_params(self, **tags):
        result = list(self.params.keys())

        only = set(tag for tag, value in list(tags.items()) if value)
        if only:
            # retain all parameters that have all of the tags in `only`
            result = [
                param for param in result if not (only - self.params[param])
            ]

        exclude = set(tag for tag, value in list(tags.items()) if not value)
        if exclude:
            # retain all parameters that have none of the tags in `exclude`
            result = [
                param for param in result if not (self.params[param] & exclude)
            ]

        return result


class InputLayer(Layer):
    def __init__(self, shape, input_var=None, **kwargs):
        super(InputLayer, self).__init__(shape, **kwargs)
        self.shape = shape
        if input_var is None:
            with tf.compat.v1.variable_scope(self.name):
                input_var = tf.compat.v1.placeholder(
                    tf.float32, shape=shape, name='input')
        self.input_var = input_var

    @Layer.output_shape.getter
    def output_shape(self):
        return self.shape


class MergeLayer(Layer):
    def __init__(self, incomings, name=None, **kwargs):
        self.input_shapes = [
            incoming if isinstance(incoming, tuple) else incoming.output_shape
            for incoming in incomings
        ]
        self.input_layers = [
            None if isinstance(incoming, tuple) else incoming
            for incoming in incomings
        ]
        if name is None:
            name = '%s_%d' % (type(self).__name__, G._n_layers)
            G._n_layers += 1
        self.name = name
        self.params = OrderedDict()
        self.get_output_kwargs = []

    @Layer.output_shape.getter
    def output_shape(self):
        shape = self.get_output_shape_for(self.input_shapes)
        if any(isinstance(s, (tf.Variable, tf.Tensor)) for s in shape):
            raise ValueError('%s returned a symbolic output shape from its '
                             'get_output_shape_for() method: %r. This is not '
                             'allowed; shapes must be tuples of integers for '
                             'fixed-size dimensions and Nones for variable '
                             'dimensions.' % (self.__class__.__name__, shape))
        return shape

    def get_output_shape_for(self, input_shapes):
        raise NotImplementedError

    def get_output_for(self, inputs, **kwargs):
        raise NotImplementedError


class ConcatLayer(MergeLayer):
    """
    Concatenates multiple inputs along the specified axis. Inputs should have
    the same shape except for the dimension specified in axis, which can have
    different sizes.
    Parameters
    -----------
    incomings : a list of :class:`Layer` instances or tuples
        The layers feeding into this layer, or expected input shapes
    axis : int
        Axis which inputs are joined over
    """

    def __init__(self, incomings, axis=1, **kwargs):
        super(ConcatLayer, self).__init__(incomings, **kwargs)
        self.axis = axis

    def get_output_shape_for(self, input_shapes):
        # Infer the output shape by grabbing, for each axis, the first
        # input size that is not `None` (if there is any)
        output_shape = [
            next((s for s in sizes if s is not None), None)
            for sizes in zip(*input_shapes)
        ]

        def match(shape1, shape2):
            return (len(shape1) == len(shape2) and all(
                i == self.axis or s1 is None or s2 is None or s1 == s2
                for i, (s1, s2) in enumerate(zip(shape1, shape2))))

        # Check for compatibility with inferred output shape
        if not all(match(shape, output_shape) for shape in input_shapes):
            raise ValueError('Mismatch: input shapes must be the same except '
                             'in the concatenation axis')
        # Infer output shape on concatenation axis and return
        sizes = [input_shape[self.axis] for input_shape in input_shapes]
        concat_size = None if any(s is None for s in sizes) else sum(sizes)
        output_shape[self.axis] = concat_size
        return tuple(output_shape)

    def get_output_for(self, inputs, **kwargs):
        with tf.name_scope(self.name, values=[inputs]):
            dtypes = [x.dtype.as_numpy_dtype for x in inputs]
            if len(set(dtypes)) > 1:
                # need to convert to common data type
                common_dtype = np.core.numerictypes.find_common_type([],
                                                                     dtypes)
                inputs = [tf.cast(x, common_dtype) for x in inputs]
            return tf.concat(axis=self.axis, values=inputs)


concat = ConcatLayer  # shortcut


class XavierUniformInitializer:
    def __call__(self, shape, dtype=tf.float32, name=None, *args, **kwargs):
        if len(shape) == 2:
            n_inputs, n_outputs = shape
        else:
            receptive_field_size = np.prod(shape[:2])
            n_inputs = shape[-2] * receptive_field_size
            n_outputs = shape[-1] * receptive_field_size
        init_range = math.sqrt(6.0 / (n_inputs + n_outputs))
        with tf.name_scope(name, 'XavierUniformInitializer'):
            return tf.random_uniform_initializer(-init_range, init_range)(
                shape, dtype=dtype)


class HeUniformInitializer:
    def __call__(self, shape, dtype=tf.float32, name=None, *args, **kwargs):
        if len(shape) == 2:
            n_inputs, _ = shape
        else:
            receptive_field_size = np.prod(shape[:2])
            n_inputs = shape[-2] * receptive_field_size
        init_range = math.sqrt(1.0 / n_inputs)
        with tf.name_scope(name, 'HeUniformInitializer'):
            return tf.random_uniform_initializer(
                -init_range, init_range, dtype=dtype)(shape)


def py_ortho_init(scale):
    def _init(shape):
        u, s, v = np.linalg.svd(np.random.uniform(size=shape))
        return np.cast['float32'](u * scale)

    return _init


class OrthogonalInitializer:
    def __init__(self, scale=1.1):
        self.scale = scale

    def __call__(self, shape, dtype=tf.float32, *args, **kwargs):
        result, = tf.py_func(py_ortho_init(self.scale), [shape], [tf.float32])
        result.set_shape(shape)
        return result


class ParamLayer(Layer):
    def __init__(self,
                 incoming,
                 num_units,
                 param=tf.zeros_initializer(),
                 trainable=True,
                 **kwargs):
        super(ParamLayer, self).__init__(incoming, **kwargs)
        self.num_units = num_units
        self.param = self.add_param(
            param, (num_units, ), name='param', trainable=trainable)

    def get_output_shape_for(self, input_shape):
        return input_shape[:-1] + (self.num_units, )

    def get_output_for(self, input, **kwargs):
        with tf.name_scope(self.name, values=[input]):
            ndim = input.get_shape().ndims
            reshaped_param = tf.reshape(
                self.param, (1, ) * (ndim - 1) + (self.num_units, ))
            tile_arg = tf.concat(
                axis=0, values=[tf.shape(input)[:ndim - 1], [1]])
            tiled = tf.tile(reshaped_param, tile_arg)
            return tiled


class OpLayer(MergeLayer):
    def __init__(self,
                 incoming,
                 op,
                 shape_op=lambda x: x,
                 extras=None,
                 **kwargs):
        if extras is None:
            extras = []
        incomings = [incoming] + extras
        super(OpLayer, self).__init__(incomings, **kwargs)
        self.op = op
        self.shape_op = shape_op
        self.incomings = incomings

    def get_output_shape_for(self, input_shapes):
        return self.shape_op(*input_shapes)

    def get_output_for(self, inputs, **kwargs):
        with tf.name_scope(self.name, values=[inputs]):
            return self.op(*inputs)


class DenseLayer(Layer):
    def __init__(self,
                 incoming,
                 num_units,
                 nonlinearity=None,
                 w=XavierUniformInitializer(),
                 b=tf.zeros_initializer(),
                 **kwargs):
        super(DenseLayer, self).__init__(incoming, **kwargs)
        if nonlinearity is None:
            self.nonlinearity = tf.identity
        else:
            self.nonlinearity = nonlinearity

        self.num_units = num_units

        num_inputs = int(np.prod(self.input_shape[1:]))

        self.w = self.add_param(w, (num_inputs, num_units), name='W')
        if b is None:
            self.b = None
        else:
            self.b = self.add_param(
                b, (num_units, ), name='b', regularizable=False)

    def get_output_shape_for(self, input_shape):
        return (input_shape[0], self.num_units)

    def get_output_for(self, input, **kwargs):
        with tf.name_scope(self.name, values=[input]):
            if input.get_shape().ndims > 2:
                # if the input has more than two dimensions, flatten it into a
                # batch of feature vectors.
                input = tf.reshape(input, tf.stack([tf.shape(input)[0], -1]))
            activation = tf.matmul(input, self.w)
            if self.b is not None:
                activation = activation + tf.expand_dims(self.b, 0)
            return self.nonlinearity(activation)


class BaseConvLayer(Layer):
    def __init__(self,
                 incoming,
                 num_filters,
                 filter_size,
                 stride=1,
                 pad='VALID',
                 untie_biases=False,
                 w=XavierUniformInitializer(),
                 b=tf.zeros_initializer(),
                 nonlinearity=tf.nn.relu,
                 n=None,
                 **kwargs):
        """
        Input is assumed to be of shape batch*height*width*channels
        """
        super(BaseConvLayer, self).__init__(incoming, **kwargs)
        if nonlinearity is None:
            self.nonlinearity = tf.identity
        else:
            self.nonlinearity = nonlinearity

        if n is None:
            n = len(self.input_shape) - 2
        elif n != len(self.input_shape) - 2:
            raise ValueError('Tried to create a %dD convolution layer with '
                             'input shape %r. Expected %d input dimensions '
                             '(batchsize, channels, %d spatial dimensions).' %
                             (n, self.input_shape, n + 2, n))
        self.n = n
        self.num_filters = num_filters
        self.filter_size = as_tuple(filter_size, n, int)
        self.stride = as_tuple(stride, n, int)
        self.untie_biases = untie_biases

        self.pad = pad

        if pad == 'SAME':
            if any(s % 2 == 0 for s in self.filter_size):
                raise NotImplementedError(
                    '`same` padding requires odd filter size.')

        self.w = self.add_param(w, self.get_w_shape(), name='W')
        if b is None:
            self.b = None
        else:
            if self.untie_biases:
                biases_shape = self.output_shape[1:3] + (
                    num_filters, )  # + self.output_shape[2:]
            else:
                biases_shape = (num_filters, )
            self.b = self.add_param(
                b, biases_shape, name='b', regularizable=False)

    def get_w_shape(self):
        """Get the shape of the weight matrix `W`.
        Returns
        -------
        tuple of int
            The shape of the weight matrix.
        """
        num_input_channels = self.input_shape[-1]
        return self.filter_size + (num_input_channels, self.num_filters)

    def get_output_shape_for(self, input_shape):
        if self.pad == 'SAME':
            pad = ('same', ) * self.n
        elif self.pad == 'VALID':
            pad = (0, ) * self.n
        else:
            raise NotImplementedError

        # if isinstance(self.pad, tuple):
        #     pad = self.pad
        # else:
        #     pad = (self.pad,) * self.n
        batchsize = input_shape[0]
        return ((batchsize, ) + tuple(
            conv_output_length(input, filter, stride, p)
            for input, filter, stride, p in zip(
                input_shape[1:3], self.filter_size, self.stride, pad))) + (
                    self.num_filters, )

    def get_output_for(self, input, **kwargs):
        with tf.name_scope(self.name, values=[input]):
            conved = self.convolve(input, **kwargs)

            if self.b is None:
                activation = conved
            elif self.untie_biases:
                # raise NotImplementedError
                activation = conved + tf.expand_dims(self.b, 0)
            else:
                activation = conved + tf.reshape(self.b,
                                                 (1, 1, 1, self.num_filters))

            return self.nonlinearity(activation)

    def convolve(self, input, name=None, **kwargs):
        """
        Symbolically convolves `input` with ``self.w``, producing an output of
        shape ``self.output_shape``. To be implemented by subclasses.
        Parameters
        ----------
        input : TensorFlow tensor
            The input minibatch to convolve
        **kwargs
            Any additional keyword arguments from :meth:`get_output_for`
        Returns
        -------
        TensorFlow tensor
            `input` convolved according to the configuration of this layer,
            without any bias or nonlinearity applied.
        """
        raise NotImplementedError('BaseConvLayer does not implement the '
                                  'convolve() method. You will want to '
                                  'use a subclass such as Conv2DLayer.')


class Conv2DLayer(BaseConvLayer):
    def __init__(self,
                 incoming,
                 num_filters,
                 filter_size,
                 stride=(1, 1),
                 pad='VALID',
                 untie_biases=False,
                 w=XavierUniformInitializer(),
                 b=tf.zeros_initializer(),
                 nonlinearity=tf.nn.relu,
                 convolution=tf.nn.conv2d,
                 **kwargs):
        super(Conv2DLayer, self).__init__(
            incoming=incoming,
            num_filters=num_filters,
            filter_size=filter_size,
            stride=stride,
            pad=pad,
            untie_biases=untie_biases,
            w=w,
            b=b,
            nonlinearity=nonlinearity,
            n=2,
            **kwargs)
        self.convolution = convolution

    def convolve(self, input, name=None, **kwargs):
        with tf.name_scope(name, 'convolve', values=[input]):
            conved = self.convolution(
                input,
                self.w,
                strides=(1, ) + self.stride + (1, ),
                padding=self.pad)
            return conved


def pool_output_length(input_length, pool_size, stride, pad):
    if input_length is None or pool_size is None:
        return None

    if pad == 'SAME':
        return int(np.ceil(float(input_length) / float(stride)))

    return int(np.ceil(float(input_length - pool_size + 1) / float(stride)))


class Pool2DLayer(Layer):
    def __init__(self,
                 incoming,
                 pool_size,
                 stride=None,
                 pad='VALID',
                 mode='max',
                 **kwargs):
        super(Pool2DLayer, self).__init__(incoming, **kwargs)

        self.pool_size = as_tuple(pool_size, 2)

        if len(self.input_shape) != 4:
            raise ValueError('Tried to create a 2D pooling layer with '
                             'input shape %r. Expected 4 input dimensions '
                             '(batchsize, 2 spatial dimensions, channels).' %
                             (self.input_shape, ))

        if stride is None:
            self.stride = self.pool_size
        else:
            self.stride = as_tuple(stride, 2)

        self.pad = pad

        self.mode = mode

    def get_output_shape_for(self, input_shape):
        output_shape = list(input_shape)  # copy / convert to mutable list

        output_shape[1] = pool_output_length(
            input_shape[1],
            pool_size=self.pool_size[0],
            stride=self.stride[0],
            pad=self.pad,
        )

        output_shape[2] = pool_output_length(
            input_shape[2],
            pool_size=self.pool_size[1],
            stride=self.stride[1],
            pad=self.pad,
        )

        return tuple(output_shape)

    def get_output_for(self, input, **kwargs):
        with tf.name_scope(self.name, values=[input]):
            assert self.mode == 'max'
            pooled = tf.nn.max_pool(
                input,
                ksize=(1, ) + self.pool_size + (1, ),
                strides=(1, ) + self.stride + (1, ),
                padding=self.pad,
            )
            return pooled


def spatial_expected_softmax(x, temp=1):
    assert len(x.get_shape()) == 4
    vals = []
    for dim in [0, 1]:
        dim_val = x.get_shape()[dim + 1].value
        lin = tf.linspace(-1.0, 1.0, dim_val)
        lin = tf.expand_dims(lin, 1 - dim)
        lin = tf.expand_dims(lin, 0)
        lin = tf.expand_dims(lin, 3)
        m = tf.reduce_max(x, [1, 2], keep_dims=True)
        e = tf.exp((x - m) / temp) + 1e-5
        val = tf.reduce_sum(e * lin, [1, 2]) / (tf.reduce_sum(e, [1, 2]))
        vals.append(tf.expand_dims(val, 2))

    return tf.reshape(
        tf.concat(axis=2, values=vals), [-1, x.get_shape()[-1].value * 2])


class SpatialExpectedSoftmaxLayer(Layer):
    """
    Computes the softmax across a spatial region, separately for each channel,
    followed by an expectation operation.
    """

    def __init__(self, incoming, **kwargs):
        super().__init__(incoming, **kwargs)
        # self.temp = self.add_param(
        #     tf.ones_initializer, shape=(), name="temperature")

    def get_output_shape_for(self, input_shape):
        return (input_shape[0], input_shape[-1] * 2)

    def get_output_for(self, input, **kwargs):
        with tf.name_scope(self.name, values=[input]):
            return spatial_expected_softmax(input)  # , self.temp)
        # max_ = tf.reduce_max(input, reduction_indices=[1, 2], keep_dims=True)
        # exp = tf.exp(input - max_) + 1e-5

        # vals = []
        #
        # for dim in [0, 1]:
        #     dim_val = input.get_shape()[dim + 1].value
        #     lin = tf.linspace(-1.0, 1.0, dim_val)
        #     lin = tf.expand_dims(lin, 1 - dim)
        #     lin = tf.expand_dims(lin, 0)
        #     lin = tf.expand_dims(lin, 3)
        #     m = tf.reduce_max(input, [1, 2], keep_dims=True)
        #     e = tf.exp(input - m) + 1e-5
        #     val = tf.reduce_sum(e * lin, [1, 2]) / (tf.reduce_sum(e, [1, 2]))
        #     vals.append(tf.expand_dims(val, 2))
        #
        # return tf.reshape(
        #     tf.concat(2, vals), [-1, input.get_shape()[-1].value * 2])

        # import ipdb; ipdb.set_trace()

        # input.get_shape()
        # exp / tf.reduce_sum(exp, reduction_indices=[1, 2], keep_dims=True)
        # import ipdb;
        # ipdb.set_trace()
        # spatial softmax?

        # for dim in range(2):
        #     val = obs.get_shape()[dim + 1].value
        #     lin = tf.linspace(-1.0, 1.0, val)
        #     lin = tf.expand_dims(lin, 1 - dim)
        #     lin = tf.expand_dims(lin, 0)
        #     lin = tf.expand_dims(lin, 3)
        #     m = tf.reduce_max(e, [1, 2], keep_dims=True)
        #     e = tf.exp(e - m) + 1e-3
        #     val = tf.reduce_sum(e * lin, [1, 2]) / (tf.reduce_sum(e, [1, 2]))


class DropoutLayer(Layer):
    def __init__(self, incoming, p=0.5, rescale=True, **kwargs):
        super(DropoutLayer, self).__init__(incoming, **kwargs)
        self.p = p
        self.rescale = rescale

    def get_output_for(self, input, deterministic=False, **kwargs):
        """
        Parameters
        ----------
        input : tensor
            output from the previous layer
        deterministic : bool
            If true dropout and scaling is disabled, see notes
        """
        with tf.name_scope(self.name, values=[input]):
            if deterministic or self.p == 0:
                return input
            else:
                retain_prob = 1. - self.p
                if self.rescale:
                    input /= retain_prob

                # use nonsymbolic shape for dropout mask if possible
                return tf.nn.dropout(input, keep_prob=retain_prob)

    def get_output_shape_for(self, input_shape):
        return input_shape


# TODO: add Conv3DLayer


class FlattenLayer(Layer):
    """
    A layer that flattens its input. The leading ``outdim-1`` dimensions of
    the output will have the same shape as the input. The remaining dimensions
    are collapsed into the last dimension.
    Parameters
    ----------
    incoming : a :class:`Layer` instance or a tuple
        The layer feeding into this layer, or the expected input shape.
    outdim : int
        The number of dimensions in the output.
    See Also
    --------
    flatten  : Shortcut
    """

    def __init__(self, incoming, outdim=2, **kwargs):
        super(FlattenLayer, self).__init__(incoming, **kwargs)
        self.outdim = outdim

        if outdim < 1:
            raise ValueError('Dim must be >0, was %i', outdim)

    def get_output_shape_for(self, input_shape):
        to_flatten = input_shape[self.outdim - 1:]

        if any(s is None for s in to_flatten):
            flattened = None
        else:
            flattened = int(np.prod(to_flatten))

        return input_shape[:self.outdim - 1] + (flattened, )

    def get_output_for(self, input, **kwargs):
        with tf.name_scope(self.name, values=[input]):
            # total_entries = tf.reduce_prod(tf.shape(input))
            pre_shape = tf.shape(input)[:self.outdim - 1]
            to_flatten = tf.reduce_prod(tf.shape(input)[self.outdim - 1:])
            return tf.reshape(
                input,
                tf.concat(axis=0, values=[pre_shape,
                                          tf.stack([to_flatten])]))


flatten = FlattenLayer  # shortcut


class ReshapeLayer(Layer):
    def __init__(self, incoming, shape, **kwargs):
        super(ReshapeLayer, self).__init__(incoming, **kwargs)
        shape = tuple(shape)
        for s in shape:
            if isinstance(s, int):
                if s == 0 or s < -1:
                    raise ValueError('`shape` integers must be positive or -1')
            elif isinstance(s, list):
                if len(s) != 1 or not isinstance(s[0], int) or s[0] < 0:
                    raise ValueError('`shape` input references must be '
                                     'single-element lists of int >= 0')
            elif isinstance(s, (tf.Tensor, tf.Variable)):
                raise NotImplementedError
            else:
                raise ValueError('`shape` must be a tuple of int and/or [int]')
        if sum(s == -1 for s in shape) > 1:
            raise ValueError('`shape` cannot contain multiple -1')
        self.shape = shape
        # try computing the output shape once as a sanity check
        self.get_output_shape_for(self.input_shape)

    def get_output_shape_for(self, input_shape, **kwargs):
        # Initialize output shape from shape specification
        output_shape = list(self.shape)
        # First, replace all `[i]` with the corresponding input dimension, and
        # mask parts of the shapes thus becoming irrelevant for -1 inference
        masked_input_shape = list(input_shape)
        masked_output_shape = list(output_shape)
        for dim, o in enumerate(output_shape):
            if isinstance(o, list):
                if o[0] >= len(input_shape):
                    raise ValueError('specification contains [%d], but input '
                                     'shape has %d dimensions only' %
                                     (o[0], len(input_shape)))
                output_shape[dim] = input_shape[o[0]]
                masked_output_shape[dim] = input_shape[o[0]]
                if (input_shape[o[0]] is None) \
                        and (masked_input_shape[o[0]] is None):
                    # first time we copied this unknown input size: mask
                    # it, we have a 1:1 correspondence between out[dim] and
                    # in[o[0]] and can ignore it for -1 inference even if
                    # it is unknown.
                    masked_input_shape[o[0]] = 1
                    masked_output_shape[dim] = 1
        # Secondly, replace all symbolic shapes with `None`, as we cannot
        # infer their size here.
        for dim, o in enumerate(output_shape):
            if isinstance(o, (tf.Tensor, tf.Variable)):
                raise NotImplementedError
                # output_shape[dim] = None
                # masked_output_shape[dim] = None
        # From the shapes, compute the sizes of the input and output tensor
        input_size = (None if any(
            x is None
            for x in masked_input_shape) else np.prod(masked_input_shape))
        output_size = (None if any(
            x is None
            for x in masked_output_shape) else np.prod(masked_output_shape))
        del masked_input_shape, masked_output_shape
        # Finally, infer value for -1 if needed
        if -1 in output_shape:
            dim = output_shape.index(-1)
            if (input_size is None) or (output_size is None):
                output_shape[dim] = None
                output_size = None
            else:
                output_size *= -1
                output_shape[dim] = input_size // output_size
                output_size *= output_shape[dim]
        # Sanity check
        if (input_size is not None) and (output_size is not None) \
                and (input_size != output_size):
            raise ValueError(
                '%s cannot be reshaped to specification %s. '
                'The total size mismatches.' % (input_shape, self.shape))
        return tuple(output_shape)

    def get_output_for(self, input, **kwargs):
        with tf.name_scope(self.name, values=[input]):
            # Replace all `[i]` with the corresponding input dimension
            output_shape = list(self.shape)
            for dim, o in enumerate(output_shape):
                if isinstance(o, list):
                    output_shape[dim] = tf.shape(input)[o[0]]
            # Everything else is handled by TensorFlow
            return tf.reshape(input, tf.stack(output_shape))


reshape = ReshapeLayer  # shortcut


class SliceLayer(Layer):
    def __init__(self, incoming, indices, axis=-1, **kwargs):
        super(SliceLayer, self).__init__(incoming, **kwargs)
        self.slice = indices
        self.axis = axis

    def get_output_shape_for(self, input_shape):
        output_shape = list(input_shape)
        if isinstance(self.slice, int):
            del output_shape[self.axis]
        elif input_shape[self.axis] is not None:
            output_shape[self.axis] = len(
                list(range(*self.slice.indices(input_shape[self.axis]))))
        else:
            output_shape[self.axis] = None
        return tuple(output_shape)

    def get_output_for(self, input, **kwargs):
        with tf.name_scope(self.name, values=[input]):
            axis = self.axis
            ndims = input.get_shape().ndims
            if axis < 0:
                axis += ndims
            if isinstance(self.slice, int) and self.slice < 0:
                return tf.reverse(
                    input, [self.axis + 1
                            ])[(slice(None), ) * axis + (-1 - self.slice, ) +
                               (slice(None), ) * (ndims - axis - 1)]
            # import ipdb; ipdb.set_trace()
            return input[(slice(None), ) * axis + (self.slice, ) +
                         (slice(None), ) * (ndims - axis - 1)]


class DimshuffleLayer(Layer):
    def __init__(self, incoming, pattern, **kwargs):
        super(DimshuffleLayer, self).__init__(incoming, **kwargs)

        # Sanity check the pattern
        used_dims = set()
        for p in pattern:
            if isinstance(p, int):
                # Dimension p
                if p in used_dims:
                    raise ValueError('pattern contains dimension {0} more '
                                     'than once'.format(p))
                used_dims.add(p)
            elif p == 'x':
                # Broadcast
                pass
            else:
                raise ValueError('pattern should only contain dimension'
                                 "indices or 'x', not {0}".format(p))

        self.pattern = pattern

        # try computing the output shape once as a sanity check
        self.get_output_shape_for(self.input_shape)

    def get_output_shape_for(self, input_shape):
        # Build output shape while keeping track of the dimensions that we are
        # attempting to collapse, so we can ensure that they are broadcastable
        output_shape = []
        dims_used = [False] * len(input_shape)
        for p in self.pattern:
            if isinstance(p, int):
                if p < 0 or p >= len(input_shape):
                    raise ValueError('pattern contains {0}, but input shape '
                                     'has {1} dimensions '
                                     'only'.format(p, len(input_shape)))
                # Dimension p
                o = input_shape[p]
                dims_used[p] = True
            elif p == 'x':
                # Broadcast; will be of size 1
                o = 1
            output_shape.append(o)

        for i, (dim_size, used) in enumerate(zip(input_shape, dims_used)):
            if not used and dim_size != 1 and dim_size is not None:
                raise ValueError(
                    'pattern attempted to collapse dimension '
                    '{0} of size {1}; dimensions with size != 1/None are not'
                    'broadcastable and cannot be '
                    'collapsed'.format(i, dim_size))

        return tuple(output_shape)

    def get_output_for(self, input, **kwargs):
        with tf.name_scope(self.name, values=[input]):
            return tf.transpose(input, self.pattern)


dimshuffle = DimshuffleLayer  # shortcut


def apply_ln(layer):
    with tf.name_scope('apply_ln'):

        def _normalize(x, prefix):
            eps = 1e-5
            dim = x.get_shape()[-1].value

            bias_name = prefix + '_ln/bias'
            scale_name = prefix + '_ln/scale'

            if bias_name not in layer.norm_params:
                layer.norm_params[bias_name] = layer.add_param(
                    tf.zeros_initializer(), (dim, ),
                    name=bias_name,
                    regularizable=False)
            if scale_name not in layer.norm_params:
                layer.norm_params[scale_name] = layer.add_param(
                    tf.ones_initializer(), (dim, ), name=scale_name)

            bias = layer.norm_params[bias_name]
            scale = layer.norm_params[scale_name]
            mean, var = tf.nn.moments(x, axes=[1], keep_dims=True)
            x_normed = (x - mean) / tf.sqrt(var + eps)
            return x_normed * scale + bias

        return _normalize


class GRULayer(Layer):
    """
    A gated recurrent unit implements the following update mechanism:
    Reset gate:        r(t) = f_r(x(t) @ W_xr + h(t-1) @ W_hr + b_r)
    Update gate:       u(t) = f_u(x(t) @ W_xu + h(t-1) @ W_hu + b_u)
    Cell gate:         c(t) = f_c(x(t) @ W_xc + r(t) * (h(t-1) @ W_hc) + b_c)
    New hidden state:  h(t) = (1 - u(t)) * c(t) + u_t * h(t-1)
    Note that the reset, update, and cell vectors must have the same dimension
    as the hidden state
    """

    def __init__(self,
                 incoming,
                 num_units,
                 hidden_nonlinearity,
                 gate_nonlinearity=tf.nn.sigmoid,
                 w_x_init=XavierUniformInitializer(),
                 w_h_init=OrthogonalInitializer(),
                 b_init=tf.zeros_initializer(),
                 hidden_init=tf.zeros_initializer(),
                 hidden_init_trainable=False,
                 layer_normalization=False,
                 **kwargs):

        if hidden_nonlinearity is None:
            hidden_nonlinearity = tf.identity

        if gate_nonlinearity is None:
            gate_nonlinearity = tf.identity

        super(GRULayer, self).__init__(incoming, **kwargs)

        with tf.compat.v1.variable_scope('gru_layer_step'):
            input_shape = self.input_shape[2:]

            input_dim = np.prod(input_shape)

            self.layer_normalization = layer_normalization

            # Weights for the initial hidden state
            self.h0 = self.add_param(
                hidden_init, (num_units, ),
                name='h0',
                trainable=hidden_init_trainable,
                regularizable=False)
            # Weights for the reset gate
            self.W_xr = self.add_param(
                w_x_init, (input_dim, num_units), name='W_xr')
            self.W_hr = self.add_param(
                w_h_init, (num_units, num_units), name='W_hr')
            self.b_r = self.add_param(
                b_init, (num_units, ), name='b_r', regularizable=False)
            # Weights for the update gate
            self.W_xu = self.add_param(
                w_x_init, (input_dim, num_units), name='W_xu')
            self.W_hu = self.add_param(
                w_h_init, (num_units, num_units), name='W_hu')
            self.b_u = self.add_param(
                b_init, (num_units, ), name='b_u', regularizable=False)
            # Weights for the cell gate
            self.W_xc = self.add_param(
                w_x_init, (input_dim, num_units), name='W_xc')
            self.W_hc = self.add_param(
                w_h_init, (num_units, num_units), name='W_hc')
            self.b_c = self.add_param(
                b_init, (num_units, ), name='b_c', regularizable=False)

            self.W_x_ruc = tf.concat(
                axis=1,
                values=[self.W_xr, self.W_xu, self.W_xc],
                name='W_x_ruc')
            self.W_h_ruc = tf.concat(
                axis=1,
                values=[self.W_hr, self.W_hu, self.W_hc],
                name='W_h_ruc')
            self.W_x_ru = tf.concat(
                axis=1, values=[self.W_xr, self.W_xu], name='W_x_ru')
            self.W_h_ru = tf.concat(
                axis=1, values=[self.W_hr, self.W_hu], name='W_h_ru')
            self.b_ruc = tf.concat(
                axis=0, values=[self.b_r, self.b_u, self.b_c], name='b_ruc')

            self.gate_nonlinearity = gate_nonlinearity
            self.num_units = num_units
            self.nonlinearity = hidden_nonlinearity
            self.norm_params = dict()

            # pre-run the step method to initialize the normalization
            # parameters
            h_dummy = tf.compat.v1.placeholder(
                dtype=tf.float32, shape=(None, num_units), name='h_dummy')
            x_dummy = tf.compat.v1.placeholder(
                dtype=tf.float32, shape=(None, input_dim), name='x_dummy')
            self.step(h_dummy, x_dummy)

    def step(self, hprev, x, name=None):
        with tf.name_scope(name, 'step', values=[hprev, x]):
            if self.layer_normalization:
                ln = apply_ln(self)
                x_ru = ln(tf.matmul(x, self.W_x_ru), 'x_ru')
                h_ru = ln(tf.matmul(hprev, self.W_h_ru), 'h_ru')
                x_r, x_u = tf.split(axis=1, num_or_size_splits=2, value=x_ru)
                h_r, h_u = tf.split(axis=1, num_or_size_splits=2, value=h_ru)
                x_c = ln(tf.matmul(x, self.W_xc), 'x_c')
                h_c = ln(tf.matmul(hprev, self.W_hc), 'h_c')
                r = self.gate_nonlinearity(x_r + h_r)
                u = self.gate_nonlinearity(x_u + h_u)
                c = self.nonlinearity(x_c + r * h_c)
                h = (1 - u) * c + u * hprev
                return h
            else:
                xb_ruc = tf.matmul(x, self.W_x_ruc) + tf.reshape(
                    self.b_ruc, (1, -1))
                h_ruc = tf.matmul(hprev, self.W_h_ruc)
                xb_r, xb_u, xb_c = tf.split(
                    axis=1, num_or_size_splits=3, value=xb_ruc)
                h_r, h_u, h_c = tf.split(
                    axis=1, num_or_size_splits=3, value=h_ruc)
                r = self.gate_nonlinearity(xb_r + h_r)
                u = self.gate_nonlinearity(xb_u + h_u)
                c = self.nonlinearity(xb_c + r * h_c)
                h = (1 - u) * c + u * hprev
                return h

    def get_step_layer(self, l_in, l_prev_hidden, name=None):
        return GRUStepLayer(
            incomings=[l_in, l_prev_hidden], recurrent_layer=self, name=name)

    def get_output_shape_for(self, input_shape):
        n_batch, n_steps = input_shape[:2]
        return n_batch, n_steps, self.num_units

    def get_output_for(self, input, **kwargs):
        with tf.name_scope(self.name, values=[input]):
            input_shape = tf.shape(input)
            n_batches = input_shape[0]
            n_steps = input_shape[1]
            input = tf.reshape(input, tf.stack([n_batches, n_steps, -1]))
            if ('recurrent_state' in kwargs
                    and self in kwargs['recurrent_state']):
                h0s = kwargs['recurrent_state'][self]
            else:
                h0s = tf.tile(
                    tf.reshape(self.h0, (1, self.num_units)), (n_batches, 1))
            # flatten extra dimensions
            shuffled_input = tf.transpose(input, (1, 0, 2))
            hs = tf.scan(self.step, elems=shuffled_input, initializer=h0s)
            shuffled_hs = tf.transpose(hs, (1, 0, 2))
            if 'recurrent_state_output' in kwargs:
                kwargs['recurrent_state_output'][self] = shuffled_hs
            return shuffled_hs


class GRUStepLayer(MergeLayer):
    def __init__(self, incomings, recurrent_layer, **kwargs):
        super(GRUStepLayer, self).__init__(incomings, **kwargs)
        self._gru_layer = recurrent_layer

    def get_params(self, **tags):
        return self._gru_layer.get_params(**tags)

    def get_output_shape_for(self, input_shapes):
        n_batch = input_shapes[0][0]
        return n_batch, self._gru_layer.num_units

    def get_output_for(self, inputs, **kwargs):
        with tf.name_scope(self.name, values=[inputs]):
            x, hprev = inputs
            n_batch = tf.shape(x)[0]
            x = tf.reshape(x, tf.stack([n_batch, -1]))
            x.set_shape((None, self.input_shapes[0][1]))
            return self._gru_layer.step(hprev, x)


class TfGRULayer(Layer):
    """
    Use TensorFlow's built-in GRU implementation
    """

    def __init__(self,
                 incoming,
                 num_units,
                 hidden_nonlinearity,
                 horizon=None,
                 hidden_init_trainable=False,
                 **kwargs):
        assert len(incoming.output_shape) == 3
        input_dim = incoming.shape[2]
        gru = tf.nn.rnn_cell.GRUCell(
            num_units=num_units, activation=hidden_nonlinearity)
        self.num_units = num_units
        self.horizon = horizon
        self.gru = gru
        self.hidden_nonlinearity = hidden_nonlinearity
        Layer.__init__(self, incoming=incoming, **kwargs)
        # dummy input variable

        with tf.compat.v1.variable_scope(self.name) as vs:
            input_dummy = tf.compat.v1.placeholder(
                tf.float32, (None, input_dim), 'input_dummy')
            hidden_dummy = tf.compat.v1.placeholder(
                tf.float32, (None, num_units), 'hidden_dummy')
            gru(input_dummy, hidden_dummy, scope=vs)
            vs.reuse_variables()
            self.scope = vs
            all_vars = [
                v for v in tf.compat.v1.global_variables()
                if v.name.startswith(vs.name)
            ]
            trainable_vars = [
                v for v in tf.compat.v1.trainable_variables()
                if v.name.startswith(vs.name)
            ]

        for var in trainable_vars:
            self.add_param(spec=var, shape=None, name=None, trainable=True)
        for var in set(all_vars) - set(trainable_vars):
            self.add_param(spec=var, shape=None, name=None, trainable=False)
        self.h0 = self.add_param(
            tf.zeros_initializer(), (num_units, ),
            name='h0',
            trainable=hidden_init_trainable,
            regularizable=False)

    def step(self, hprev, x, name=None):
        with tf.name_scope(name, 'step', values=[hprev, x]):
            return self.gru(x, hprev, scope=self.scope)[1]

    def get_output_for(self, input, **kwargs):
        with tf.name_scope(self.name, values=[input]):
            input_shape = tf.shape(input)
            n_batches = input_shape[0]
            state = tf.tile(
                tf.reshape(self.h0, (1, self.num_units)), (n_batches, 1))
            state.set_shape((None, self.num_units))
            if self.horizon is not None:
                outputs = []
                for idx in range(self.horizon):
                    output, state = self.gru(
                        input[:, idx, :], state,
                        scope=self.scope)  # self.name)
                    outputs.append(tf.expand_dims(output, 1))
                outputs = tf.concat(axis=1, values=outputs)
                return outputs
            else:
                n_steps = input_shape[1]
                input = tf.reshape(input, tf.stack([n_batches, n_steps, -1]))
                # flatten extra dimensions
                shuffled_input = tf.transpose(input, (1, 0, 2))
                shuffled_input.set_shape((None, None, self.input_shape[-1]))
                hs = tf.scan(
                    self.step, elems=shuffled_input, initializer=state)
                shuffled_hs = tf.transpose(hs, (1, 0, 2))
                return shuffled_hs

    def get_output_shape_for(self, input_shape):
        n_batch, n_steps = input_shape[:2]
        return n_batch, n_steps, self.num_units

    def get_step_layer(self, l_in, l_prev_hidden, name=None):
        return GRUStepLayer(
            incomings=[l_in, l_prev_hidden], recurrent_layer=self, name=name)


class PseudoLSTMLayer(Layer):
    """
    A Pseudo LSTM unit implements the following update mechanism:

    Incoming gate:     i(t) = σ(W_hi @ h(t-1)) + W_xi @ x(t) + b_i)
    Forget gate:       f(t) = σ(W_hf @ h(t-1)) + W_xf @ x(t) + b_f)
    Out gate:          o(t) = σ(W_ho @ h(t-1)) + W_xo @ x(t) + b_o)
    New cell gate:     c_new(t) = ϕ(W_hc @ (o(t) * h(t-1)) + W_xc @ x(t) + b_c)
    Cell gate:         c(t) = f(t) * c(t-1) + i(t) * c_new(t)
    Hidden state:      h(t) = ϕ(c(t))
    Output:            out  = h(t)

    If gate_squash_inputs is set to True, we have the following updates
    instead:

    Out gate:          o(t) = σ(W_ho @ h(t-1)) + W_xo @ x(t) + b_o)
    Incoming gate:     i(t) = σ(W_hi @ (o(t) * h(t-1)) + W_xi @ x(t) + b_i)
    Forget gate:       f(t) = σ(W_hf @ (o(t) * h(t-1)) + W_xf @ x(t) + b_f)
    New cell gate:     c_new(t) = ϕ(W_hc @ (o(t) * h(t-1)) + W_xc @ x(t) + b_c)
    Cell state:        c(t) = f(t) * c(t-1) + i(t) * c_new(t)
    Hidden state:      h(t) = ϕ(c(t))
    Output:            out  = h(t)

    Note that the incoming, forget, cell, and out vectors must have the same
    dimension as the hidden state

    The notation is slightly different from
    http://r2rt.com/written-memories-understanding-deriving-and-extending-the-lstm.html.
    Here we introduce the cell gate and swap its role with the hidden state, so
    that the output is the same as the hidden state (and we can use this as a
    drop-in replacement for LSTMLayer).
    """

    def __init__(self,
                 incoming,
                 num_units,
                 hidden_nonlinearity=tf.tanh,
                 gate_nonlinearity=tf.nn.sigmoid,
                 w_x_init=XavierUniformInitializer(),
                 w_h_init=OrthogonalInitializer(),
                 forget_bias=1.0,
                 b_init=tf.zeros_initializer(),
                 hidden_init=tf.zeros_initializer(),
                 hidden_init_trainable=False,
                 cell_init=tf.zeros_initializer(),
                 cell_init_trainable=False,
                 gate_squash_inputs=False,
                 layer_normalization=False,
                 **kwargs):

        if hidden_nonlinearity is None:
            hidden_nonlinearity = tf.identity

        if gate_nonlinearity is None:
            gate_nonlinearity = tf.identity

        super(PseudoLSTMLayer, self).__init__(incoming, **kwargs)

        with tf.compat.v1.variable_scope('pseudo_lstm_layer'):
            self.layer_normalization = layer_normalization

            input_shape = self.input_shape[2:]

            input_dim = np.prod(input_shape)
            # Weights for the initial hidden state (this is actually not used,
            # since the initial hidden state is determined by the initial cell
            # state via h0 = self.nonlinearity(c0)). It is here merely for
            # interface convenience
            self.h0 = self.add_param(
                hidden_init, (num_units, ),
                name='h0',
                trainable=hidden_init_trainable,
                regularizable=False)
            # Weights for the initial cell state
            self.c0 = self.add_param(
                cell_init, (num_units, ),
                name='c0',
                trainable=cell_init_trainable,
                regularizable=False)
            # Weights for the incoming gate
            self.W_xi = self.add_param(
                w_x_init, (input_dim, num_units), name='W_xi')
            self.W_hi = self.add_param(
                w_h_init, (num_units, num_units), name='W_hi')
            self.b_i = self.add_param(
                b_init, (num_units, ), name='b_i', regularizable=False)
            # Weights for the forget gate
            self.W_xf = self.add_param(
                w_x_init, (input_dim, num_units), name='W_xf')
            self.W_hf = self.add_param(
                w_h_init, (num_units, num_units), name='W_hf')
            self.b_f = self.add_param(
                b_init, (num_units, ), name='b_f', regularizable=False)
            # Weights for the out gate
            self.W_xo = self.add_param(
                w_x_init, (input_dim, num_units), name='W_xo')
            self.W_ho = self.add_param(
                w_h_init, (num_units, num_units), name='W_ho')
            self.b_o = self.add_param(
                b_init, (num_units, ), name='b_o', regularizable=False)
            # Weights for the cell gate
            self.W_xc = self.add_param(
                w_x_init, (input_dim, num_units), name='W_xc')
            self.W_hc = self.add_param(
                w_h_init, (num_units, num_units), name='W_hc')
            self.b_c = self.add_param(
                b_init, (num_units, ), name='b_c', regularizable=False)

            self.gate_nonlinearity = gate_nonlinearity
            self.num_units = num_units
            self.nonlinearity = hidden_nonlinearity
            self.forget_bias = forget_bias
            self.gate_squash_inputs = gate_squash_inputs

            self.W_x_ifo = tf.concat(
                axis=1, values=[self.W_xi, self.W_xf, self.W_xo])
            self.W_h_ifo = tf.concat(
                axis=1, values=[self.W_hi, self.W_hf, self.W_ho])

            self.W_x_if = tf.concat(axis=1, values=[self.W_xi, self.W_xf])
            self.W_h_if = tf.concat(axis=1, values=[self.W_hi, self.W_hf])

            self.norm_params = dict()

    def step(self, hcprev, x, name=None):
        with tf.name_scope(name, 'step', values=[hcprev, x]):
            hprev = hcprev[:, :self.num_units]
            cprev = hcprev[:, self.num_units:]

            if self.layer_normalization:
                ln = apply_ln(self)
            else:

                def f(x, *args):
                    return x

                ln = f

            if self.gate_squash_inputs:
                """
                Out gate:          o(t) = σ(W_ho @ h(t-1)) + W_xo @ x(t) + b_o)
                Incoming gate:     i(t) = σ(W_hi @ (o(t) * h(t-1)) +
                                            W_xi @ x(t) + b_i)
                Forget gate:       f(t) = σ(W_hf @ (o(t) * h(t-1)) +
                                            W_xf @ x(t) + b_f)
                New cell gate:     c_new(t) = ϕ(W_hc @ (o(t) * h(t-1)) +
                                                W_xc @ x(t) + b_c)
                Cell state:        c(t) = f(t) * c(t-1) + i(t) * c_new(t)
                Hidden state:      h(t) = ϕ(c(t))
                Output:            out  = h(t)
                """

                o = self.nonlinearity(
                    ln(tf.matmul(hprev, self.W_ho), 'h_o') +
                    ln(tf.matmul(x, self.W_xo), 'x_o') + self.b_o)

                x_if = ln(tf.matmul(x, self.W_x_if), 'x_if')
                h_if = ln(tf.matmul(o * hprev, self.W_h_if), 'h_if')

                x_i, x_f = tf.split(axis=1, num_or_size_splits=2, value=x_if)
                h_i, h_f = tf.split(axis=1, num_or_size_splits=2, value=h_if)

                i = self.gate_nonlinearity(x_i + h_i + self.b_i)
                f = self.gate_nonlinearity(x_f + h_f + self.b_f +
                                           self.forget_bias)
                c_new = self.nonlinearity(
                    ln(tf.matmul(o * hprev, self.W_hc), 'h_c') +
                    ln(tf.matmul(x, self.W_xc), 'x_c') + self.b_c)
                c = f * cprev + i * c_new
                h = self.nonlinearity(ln(c, 'c'))
                return tf.concat(axis=1, values=[h, c])
            else:
                """
                Incoming gate:     i(t) = σ(W_hi @ h(t-1)) + W_xi @ x(t) + b_i)
                Forget gate:       f(t) = σ(W_hf @ h(t-1)) + W_xf @ x(t) + b_f)
                Out gate:          o(t) = σ(W_ho @ h(t-1)) + W_xo @ x(t) + b_o)
                New cell gate:     c_new(t) = ϕ(W_hc @ (o(t) * h(t-1)) +
                                                W_xc @ x(t) + b_c)
                Cell gate:         c(t) = f(t) * c(t-1) + i(t) * c_new(t)
                Hidden state:      h(t) = ϕ(c(t))
                Output:            out  = h(t)
                """

                x_ifo = ln(tf.matmul(x, self.W_x_ifo), 'x_ifo')
                h_ifo = ln(tf.matmul(hprev, self.W_h_ifo), 'h_ifo')

                x_i, x_f, x_o = tf.split(
                    axis=1, num_or_size_splits=3, value=x_ifo)
                h_i, h_f, h_o = tf.split(
                    axis=1, num_or_size_splits=3, value=h_ifo)

                i = self.gate_nonlinearity(x_i + h_i + self.b_i)
                f = self.gate_nonlinearity(x_f + h_f + self.b_f +
                                           self.forget_bias)
                o = self.gate_nonlinearity(x_o + h_o + self.b_o)
                c_new = self.nonlinearity(
                    ln(tf.matmul(o * hprev, self.W_hc), 'h_c') +
                    ln(tf.matmul(x, self.W_xc), 'x_c') + self.b_c)
                c = f * cprev + i * c_new
                h = self.nonlinearity(ln(c, 'c'))
                return tf.concat(axis=1, values=[h, c])

    def get_step_layer(self, l_in, l_prev_state, name=None):
        return LSTMStepLayer(
            incomings=[l_in, l_prev_state], recurrent_layer=self, name=name)

    def get_output_shape_for(self, input_shape):
        n_batch, n_steps = input_shape[:2]
        return n_batch, n_steps, self.num_units

    def get_output_for(self, input, **kwargs):
        with tf.name_scope(self.name, values=[input]):
            input_shape = tf.shape(input)
            n_batches = input_shape[0]
            n_steps = input_shape[1]
            input = tf.reshape(input, tf.stack([n_batches, n_steps, -1]))
            c0s = tf.tile(
                tf.reshape(self.c0, (1, self.num_units)), (n_batches, 1))
            h0s = self.nonlinearity(c0s)
            # flatten extra dimensions
            shuffled_input = tf.transpose(input, (1, 0, 2))
            hcs = tf.scan(
                self.step,
                elems=shuffled_input,
                initializer=tf.concat(axis=1, values=[h0s, c0s]))
            shuffled_hcs = tf.transpose(hcs, (1, 0, 2))
            shuffled_hs = shuffled_hcs[:, :, :self.num_units]
            return shuffled_hs


class LSTMLayer(Layer):
    """
    A LSTM unit implements the following update mechanism:

    Incoming gate:    i(t) = f_i(x(t) @ W_xi + h(t-1) @ W_hi +
                                 w_ci * c(t-1) + b_i)
    Forget gate:      f(t) = f_f(x(t) @ W_xf + h(t-1) @ W_hf +
                                 w_cf * c(t-1) + b_f)
    Cell gate:        c(t) = f(t) * c(t - 1) + i(t) * f_c(x(t) @ W_xc +
                             h(t-1) @ W_hc + b_c)
    Out gate:         o(t) = f_o(x(t) @ W_xo + h(t-1) W_ho + w_co * c(t) + b_o)
    New hidden state: h(t) = o(t) * f_h(c(t))

    Note that the incoming, forget, cell, and out vectors must have the same
    dimension as the hidden state
    """

    def __init__(self,
                 incoming,
                 num_units,
                 hidden_nonlinearity=tf.tanh,
                 gate_nonlinearity=tf.nn.sigmoid,
                 w_x_init=XavierUniformInitializer(),
                 w_h_init=OrthogonalInitializer(),
                 forget_bias=1.0,
                 use_peepholes=False,
                 w_init=tf.random_normal_initializer(stddev=0.1),
                 b_init=tf.zeros_initializer(),
                 hidden_init=tf.zeros_initializer(),
                 hidden_init_trainable=False,
                 cell_init=tf.zeros_initializer(),
                 cell_init_trainable=False,
                 layer_normalization=False,
                 **kwargs):

        if hidden_nonlinearity is None:
            hidden_nonlinearity = tf.identity

        if gate_nonlinearity is None:
            gate_nonlinearity = tf.identity

        super(LSTMLayer, self).__init__(incoming, **kwargs)

        with tf.compat.v1.variable_scope('lstm_layer'):
            self.layer_normalization = layer_normalization

            input_shape = self.input_shape[2:]

            input_dim = np.prod(input_shape)
            # Weights for the initial hidden state
            self.h0 = self.add_param(
                hidden_init, (num_units, ),
                name='h0',
                trainable=hidden_init_trainable,
                regularizable=False)
            # Weights for the initial cell state
            self.c0 = self.add_param(
                cell_init, (num_units, ),
                name='c0',
                trainable=cell_init_trainable,
                regularizable=False)
            # Weights for the incoming gate
            self.W_xi = self.add_param(
                w_x_init, (input_dim, num_units), name='W_xi')
            self.W_hi = self.add_param(
                w_h_init, (num_units, num_units), name='W_hi')
            if use_peepholes:
                self.w_ci = self.add_param(w_init, (num_units, ), name='w_ci')
            self.b_i = self.add_param(
                b_init, (num_units, ), name='b_i', regularizable=False)
            # Weights for the forget gate
            self.W_xf = self.add_param(
                w_x_init, (input_dim, num_units), name='W_xf')
            self.W_hf = self.add_param(
                w_h_init, (num_units, num_units), name='W_hf')
            if use_peepholes:
                self.w_cf = self.add_param(w_init, (num_units, ), name='w_cf')
            self.b_f = self.add_param(
                b_init, (num_units, ), name='b_f', regularizable=False)
            # Weights for the cell gate
            self.W_xc = self.add_param(
                w_x_init, (input_dim, num_units), name='W_xc')
            self.W_hc = self.add_param(
                w_h_init, (num_units, num_units), name='W_hc')
            self.b_c = self.add_param(
                b_init, (num_units, ), name='b_c', regularizable=False)
            # Weights for the reset gate
            self.W_xr = self.add_param(
                w_x_init, (input_dim, num_units), name='W_xr')
            self.W_hr = self.add_param(
                w_h_init, (num_units, num_units), name='W_hr')
            self.b_r = self.add_param(
                b_init, (num_units, ), name='b_r', regularizable=False)
            # Weights for the out gate
            self.W_xo = self.add_param(
                w_x_init, (input_dim, num_units), name='W_xo')
            self.W_ho = self.add_param(
                w_h_init, (num_units, num_units), name='W_ho')
            if use_peepholes:
                self.w_co = self.add_param(w_init, (num_units, ), name='w_co')
            self.b_o = self.add_param(
                b_init, (num_units, ), name='b_o', regularizable=False)
            self.gate_nonlinearity = gate_nonlinearity
            self.num_units = num_units
            self.nonlinearity = hidden_nonlinearity
            self.forget_bias = forget_bias
            self.use_peepholes = use_peepholes

            self.W_x_ifco = tf.concat(
                axis=1, values=[self.W_xi, self.W_xf, self.W_xc, self.W_xo])
            self.W_h_ifco = tf.concat(
                axis=1, values=[self.W_hi, self.W_hf, self.W_hc, self.W_ho])

            if use_peepholes:
                self.w_c_ifo = tf.concat(
                    axis=0, values=[self.w_ci, self.w_cf, self.w_co])

            self.norm_params = dict()

    def step(self, hcprev, x, name=None):
        """
        Incoming gate:     i(t) = f_i(x(t) @ W_xi + h(t-1) @ W_hi +
                                      w_ci * c(t-1) + b_i)
        Forget gate:       f(t) = f_f(x(t) @ W_xf + h(t-1) @ W_hf +
                                      w_cf * c(t-1) + b_f)
        Cell gate:         c(t) = f(t) * c(t - 1) + i(t) * f_c(x(t) @ W_xc +
                                  h(t-1) @ W_hc + b_c)
        Out gate:          o(t) = f_o(x(t) @ W_xo + h(t-1) W_ho +
                                      w_co * c(t) + b_o)
        New hidden state:  h(t) = o(t) * f_h(c(t))
        """
        with tf.name_scope(name, 'step', values=[hcprev, x]):
            hprev = hcprev[:, :self.num_units]
            cprev = hcprev[:, self.num_units:]

            if self.layer_normalization:
                ln = apply_ln(self)
            else:

                def f(x, *args):
                    return x

                ln = f

            x_ifco = ln(tf.matmul(x, self.W_x_ifco), 'x_ifco')
            h_ifco = ln(tf.matmul(hprev, self.W_h_ifco), 'h_ifco')
            x_i, x_f, x_c, x_o = tf.split(
                axis=1, num_or_size_splits=4, value=x_ifco)
            h_i, h_f, h_c, h_o = tf.split(
                axis=1, num_or_size_splits=4, value=h_ifco)

            if self.use_peepholes:
                i = self.gate_nonlinearity(x_i + h_i + self.w_ci * cprev +
                                           self.b_i)
                f = self.gate_nonlinearity(x_f + h_f + self.w_cf * cprev +
                                           self.b_f + self.forget_bias)

                o = self.gate_nonlinearity(x_o + h_o + self.w_co * cprev +
                                           self.b_o)
            else:
                i = self.gate_nonlinearity(x_i + h_i + self.b_i)
                f = self.gate_nonlinearity(x_f + h_f + self.b_f +
                                           self.forget_bias)
                o = self.gate_nonlinearity(x_o + h_o + self.b_o)

            c = f * cprev + i * self.nonlinearity(x_c + h_c + self.b_c)
            h = o * self.nonlinearity(ln(c, 'c'))

            return tf.concat(axis=1, values=[h, c])

    def get_step_layer(self, l_in, l_prev_state, name=None):
        return LSTMStepLayer(
            incomings=[l_in, l_prev_state], recurrent_layer=self, name=name)

    def get_output_shape_for(self, input_shape):
        n_batch, n_steps = input_shape[:2]
        return n_batch, n_steps, self.num_units

    def get_output_for(self, input, **kwargs):
        with tf.name_scope(self.name, values=[input]):
            input_shape = tf.shape(input)
            n_batches = input_shape[0]
            n_steps = input_shape[1]
            input = tf.reshape(input, tf.stack([n_batches, n_steps, -1]))
            h0s = tf.tile(
                tf.reshape(self.h0, (1, self.num_units)), (n_batches, 1))
            c0s = tf.tile(
                tf.reshape(self.c0, (1, self.num_units)), (n_batches, 1))
            # flatten extra dimensions
            shuffled_input = tf.transpose(input, (1, 0, 2))
            hcs = tf.scan(
                self.step,
                elems=shuffled_input,
                initializer=tf.concat(axis=1, values=[h0s, c0s]))
            shuffled_hcs = tf.transpose(hcs, (1, 0, 2))
            shuffled_hs = shuffled_hcs[:, :, :self.num_units]
            if 'recurrent_state_output' in kwargs:
                kwargs['recurrent_state_output'][self] = shuffled_hcs
            return shuffled_hs


class LSTMStepLayer(MergeLayer):
    def __init__(self, incomings, recurrent_layer, **kwargs):
        super(LSTMStepLayer, self).__init__(incomings, **kwargs)
        self._recurrent_layer = recurrent_layer

    def get_params(self, **tags):
        return self._recurrent_layer.get_params(**tags)

    def get_output_shape_for(self, input_shapes):
        n_batch = input_shapes[0][0]
        return n_batch, 2 * self._recurrent_layer.num_units

    def get_output_for(self, inputs, **kwargs):
        with tf.name_scope(self.name, values=[inputs]):
            x, hcprev = inputs
            n_batch = tf.shape(x)[0]
            x = tf.reshape(x, tf.stack([n_batch, -1]))
            hc = self._recurrent_layer.step(hcprev, x)
            return hc


class TfBasicLSTMLayer(Layer):
    """
    Use TensorFlow's built-in (basic) LSTM implementation
    """

    def __init__(self,
                 incoming,
                 num_units,
                 hidden_nonlinearity,
                 horizon=None,
                 hidden_init_trainable=False,
                 forget_bias=1.0,
                 use_peepholes=False,
                 **kwargs):
        assert not use_peepholes, 'Basic LSTM does not support peepholes!'
        assert len(incoming.output_shape) == 3
        input_dim = incoming.shape[2]
        lstm = tf.contrib.rnn.BasicLSTMCell(
            num_units=num_units,
            activation=hidden_nonlinearity,
            state_is_tuple=True,
            forget_bias=forget_bias)
        self.num_units = num_units
        self.horizon = horizon
        self.lstm = lstm
        self.hidden_nonlinearity = hidden_nonlinearity
        Layer.__init__(self, incoming=incoming, **kwargs)

        with tf.compat.v1.variable_scope(self.name):
            # dummy input variable
            input_dummy = tf.compat.v1.placeholder(
                tf.float32, (None, input_dim), 'input_dummy')
            hidden_dummy = tf.compat.v1.placeholder(
                tf.float32, (None, num_units), 'hidden_dummy')
            cell_dummy = tf.compat.v1.placeholder(
                tf.float32, (None, num_units), 'cell_dummy')

            with tf.compat.v1.variable_scope('lstm') as vs:
                lstm(input_dummy, (cell_dummy, hidden_dummy), scope=vs)
                vs.reuse_variables()
                self.scope = vs
                all_vars = [
                    v for v in tf.compat.v1.global_variables()
                    if v.name.startswith(vs.name)
                ]
                trainable_vars = [
                    v for v in tf.compat.v1.trainable_variables()
                    if v.name.startswith(vs.name)
                ]

        for var in trainable_vars:
            self.add_param(spec=var, shape=None, name=None, trainable=True)
        for var in set(all_vars) - set(trainable_vars):
            self.add_param(spec=var, shape=None, name=None, trainable=False)

        self.h0 = self.add_param(
            tf.zeros_initializer(), (num_units, ),
            name='h0',
            trainable=hidden_init_trainable,
            regularizable=False)
        self.c0 = self.add_param(
            tf.zeros_initializer(), (num_units, ),
            name='c0',
            trainable=hidden_init_trainable,
            regularizable=False)

    def step(self, hcprev, x, name=None):
        with tf.name_scope(name, 'step', values=[hcprev, x]):
            hprev = hcprev[:, :self.num_units]
            cprev = hcprev[:, self.num_units:]
            x.set_shape((None, self.input_shape[-1]))
            c, h = self.lstm(x, (cprev, hprev), scope=self.scope)[1]
            return tf.concat(axis=1, values=[h, c])

    def get_output_for(self, input, **kwargs):
        with tf.name_scope(self.name, values=[input]):
            input_shape = tf.shape(input)
            n_batches = input_shape[0]
            h0s = tf.tile(
                tf.reshape(self.h0, (1, self.num_units)), (n_batches, 1))
            h0s.set_shape((None, self.num_units))
            c0s = tf.tile(
                tf.reshape(self.c0, (1, self.num_units)), (n_batches, 1))
            c0s.set_shape((None, self.num_units))
            state = (c0s, h0s)
            if self.horizon is not None:
                outputs = []
                for idx in range(self.horizon):
                    output, state = self.lstm(
                        input[:, idx, :], state,
                        scope=self.scope)  # self.name)
                    outputs.append(tf.expand_dims(output, 1))
                outputs = tf.concat(axis=1, values=outputs)
                return outputs
            else:
                n_steps = input_shape[1]
                input = tf.reshape(input, tf.stack([n_batches, n_steps, -1]))
                # flatten extra dimensions
                shuffled_input = tf.transpose(input, (1, 0, 2))
                shuffled_input.set_shape((None, None, self.input_shape[-1]))
                hcs = tf.scan(
                    self.step,
                    elems=shuffled_input,
                    initializer=tf.concat(axis=1, values=[h0s, c0s]),
                )
                shuffled_hcs = tf.transpose(hcs, (1, 0, 2))
                shuffled_hs = shuffled_hcs[:, :, :self.num_units]
                return shuffled_hs

    def get_output_shape_for(self, input_shape):
        n_batch, n_steps = input_shape[:2]
        return n_batch, n_steps, self.num_units

    def get_step_layer(self, l_in, l_prev_state, name=None):
        return LSTMStepLayer(
            incomings=[l_in, l_prev_state], recurrent_layer=self, name=name)


def get_all_layers(layer, treat_as_input=None):
    """
    :type layer: Layer | list[Layer]
    :rtype: list[Layer]
    """
    # We perform a depth-first search. We add a layer to the result list only
    # after adding all its incoming layers (if any) or when detecting a cycle.
    # We use a LIFO stack to avoid ever running into recursion depth limits.
    try:
        queue = deque(layer)
    except TypeError:
        queue = deque([layer])
    seen = set()
    done = set()
    result = []

    # If treat_as_input is given, we pretend we've already collected all their
    # incoming layers.
    if treat_as_input is not None:
        seen.update(treat_as_input)

    while queue:
        # Peek at the leftmost node in the queue.
        layer = queue[0]
        if layer is None:
            # Some node had an input_layer set to `None`. Just ignore it.
            queue.popleft()
        elif layer not in seen:
            # We haven't seen this node yet: Mark it and queue all incomings
            # to be processed first. If there are no incomings, the node will
            # be appended to the result list in the next iteration.
            seen.add(layer)
            if hasattr(layer, 'input_layers'):
                queue.extendleft(reversed(layer.input_layers))
            elif hasattr(layer, 'input_layer'):
                queue.appendleft(layer.input_layer)
        else:
            # We've been here before: Either we've finished all its incomings,
            # or we've detected a cycle. In both cases, we remove the layer
            # from the queue and append it to the result list.
            queue.popleft()
            if layer not in done:
                result.append(layer)
                done.add(layer)

    return result


class NonlinearityLayer(Layer):
    def __init__(self, incoming, nonlinearity=tf.nn.relu, **kwargs):
        super(NonlinearityLayer, self).__init__(incoming, **kwargs)
        self.nonlinearity = (tf.identity
                             if nonlinearity is None else nonlinearity)

    def get_output_for(self, input, **kwargs):
        with tf.name_scope(self.name, values=[input]):
            return self.nonlinearity(input)

    def get_output_shape_for(self, input_shape):
        return input_shape


class BatchNormLayer(Layer):
    def __init__(self,
                 incoming,
                 center=True,
                 scale=False,
                 epsilon=0.001,
                 decay=0.9,
                 beta=tf.zeros_initializer(),
                 gamma=tf.ones_initializer(),
                 moving_mean=tf.zeros_initializer(),
                 moving_variance=tf.ones_initializer(),
                 **kwargs):
        super(BatchNormLayer, self).__init__(incoming, **kwargs)

        self.center = center
        self.scale = scale
        self.epsilon = epsilon
        self.decay = decay

        input_shape = incoming.output_shape
        axis = list(range(len(input_shape) - 1))
        params_shape = input_shape[-1:]

        if center:
            self.beta = self.add_param(
                beta,
                shape=params_shape,
                name='beta',
                trainable=True,
                regularizable=False)
        else:
            self.beta = None
        if scale:
            self.gamma = self.add_param(
                gamma,
                shape=params_shape,
                name='gamma',
                trainable=True,
                regularizable=True)
        else:
            self.gamma = None

        self.moving_mean = self.add_param(
            moving_mean,
            shape=params_shape,
            name='moving_mean',
            trainable=False,
            regularizable=False)
        self.moving_variance = self.add_param(
            moving_variance,
            shape=params_shape,
            name='moving_variance',
            trainable=False,
            regularizable=False)
        self.axis = axis

    def get_output_for(self, input, phase='train', **kwargs):
        with tf.name_scope(self.name, values=[input]):
            if phase == 'train':
                # Calculate the moments based on the individual batch.
                mean, variance = tf.nn.moments(
                    input, self.axis, shift=self.moving_mean)
                # Update the moving_mean and moving_variance moments.
                update_moving_mean = moving_averages.assign_moving_average(
                    self.moving_mean, mean, self.decay)
                update_moving_variance = moving_averages.assign_moving_average(
                    self.moving_variance, variance, self.decay)
                # Make sure the updates are computed here.
                with tf.control_dependencies(
                        [update_moving_mean,
                         update_moving_variance]):  # yapf:disable
                    output = tf.nn.batch_normalization(input, mean, variance,
                                                       self.beta, self.gamma,
                                                       self.epsilon)
            else:
                output = tf.nn.batch_normalization(
                    input, self.moving_mean, self.moving_variance, self.beta,
                    self.gamma, self.epsilon)
            output.set_shape(self.input_shape)
            return output

    def get_output_shape_for(self, input_shape):
        return input_shape


def batch_norm(layer, **kwargs):
    nonlinearity = getattr(layer, 'nonlinearity', None)
    scale = True
    if nonlinearity is not None:
        layer.nonlinearity = tf.identity
        if nonlinearity is tf.nn.relu:
            scale = False
    if hasattr(layer, 'b') and layer.b is not None:
        del layer.params[layer.b]
        layer.b = None
    bn_name = (kwargs.pop('name', None)
               or (getattr(layer, 'name', None) and layer.name + '_bn'))
    layer = BatchNormLayer(layer, name=bn_name, scale=scale, **kwargs)
    if nonlinearity is not None:
        nonlin_name = bn_name and bn_name + '_nonlin'
        layer = NonlinearityLayer(
            layer, nonlinearity=nonlinearity, name=nonlin_name)
    return layer


class ElemwiseSumLayer(MergeLayer):
    def __init__(self, incomings, **kwargs):
        super(ElemwiseSumLayer, self).__init__(incomings, **kwargs)

    def get_output_for(self, inputs, **kwargs):
        with tf.name_scope(self.name, values=[inputs]):
            return functools.reduce(tf.add, inputs)

    def get_output_shape_for(self, input_shapes):
        assert len(set(input_shapes)) == 1
        return input_shapes[0]


def get_output(layer_or_layers, inputs=None, **kwargs):
    # track accepted kwargs used by get_output_for
    accepted_kwargs = {'deterministic'}
    # obtain topological ordering of all layers the output layer(s) depend on
    treat_as_input = list(inputs.keys()) if isinstance(inputs, dict) else []
    all_layers = get_all_layers(layer_or_layers, treat_as_input)
    # initialize layer-to-expression mapping from all input layers
    all_outputs = dict(
        (layer, layer.input_var) for layer in all_layers
        if isinstance(layer, InputLayer) and layer not in treat_as_input)
    # update layer-to-expression mapping from given input(s), if any
    if isinstance(inputs, dict):
        all_outputs.update((layer, tf.convert_to_tensor(expr))
                           for layer, expr in list(inputs.items()))
    elif inputs is not None:
        if len(all_outputs) > 1:
            raise ValueError('get_output() was called with a single input '
                             'expression on a network with multiple input '
                             'layers. Please call it with a dictionary of '
                             'input expressions instead.')
        for input_layer in all_outputs:
            all_outputs[input_layer] = tf.convert_to_tensor(inputs)
    # update layer-to-expression mapping by propagating the inputs
    for layer in all_layers:
        if layer not in all_outputs:
            try:
                if isinstance(layer, MergeLayer):
                    layer_inputs = [
                        all_outputs[input_layer]
                        for input_layer in layer.input_layers
                    ]
                else:
                    layer_inputs = all_outputs[layer.input_layer]
            except KeyError:
                # one of the input_layer attributes must have been `None`
                raise ValueError(
                    'get_output() was called without giving an '
                    'input expression for the free-floating '
                    'layer %r. Please call it with a dictionary '
                    'mapping this layer to an input expression.' % layer)
            all_outputs[layer] = layer.get_output_for(layer_inputs, **kwargs)
            try:
                names, _, _, defaults, _, _, _ = inspect.getfullargspec(
                    layer.get_output_for)
            except TypeError:
                # If introspection is not possible, skip it
                pass
            else:
                if defaults is not None:
                    accepted_kwargs |= set(names[-len(defaults):])
            accepted_kwargs |= set(layer.get_output_kwargs)
    unused_kwargs = set(kwargs.keys()) - accepted_kwargs
    if unused_kwargs:
        suggestions = []
        for kwarg in unused_kwargs:
            suggestion = get_close_matches(kwarg, accepted_kwargs)
            if suggestion:
                suggestions.append(
                    '%s (perhaps you meant %s)' % (kwarg, suggestion[0]))
            else:
                suggestions.append(kwarg)
        warn('get_output() was called with unused kwargs:\n\t%s' %
             '\n\t'.join(suggestions))
    # return the output(s) of the requested layer(s) only
    try:
        return [all_outputs[layer] for layer in layer_or_layers]
    except TypeError:
        return all_outputs[layer_or_layers]


def unique(l):
    """Filters duplicates of iterable.
    Create a new list from l with duplicate entries removed,
    while preserving the original order.
    Parameters
    ----------
    l : iterable
        Input iterable to filter of duplicates.
    Returns
    -------
    list
        A list of elements of `l` without duplicates and in the same order.
    """
    new_list = []
    seen = set()
    for el in l:
        if el not in seen:
            new_list.append(el)
            seen.add(el)

    return new_list


def get_all_params(layer, **tags):
    """
    :type layer: Layer|list[Layer]
    """
    layers = get_all_layers(layer)
    params = chain.from_iterable(l.get_params(**tags) for l in layers)
    return unique(params)
