# Path: models/conv_lstm.py
# @add_arg_scope
# def basic_conv_lstm_cell(inputs,
#                          state,
#                          num_channels,
#                          initializer,
#                          filter_size=5,
#                          forget_bias=1.0,
#                          scope=None,
#                          reuse=None):
#   """Basic LSTM recurrent network cell, with 2D convolution connctions.
#   We add forget_bias (default: 1) to the biases of the forget gate in order to
#   reduce the scale of forgetting in the beginning of the training.
#   It does not allow cell clipping, a projection layer, and does not
#   use peep-hole connections: it is the basic baseline.
#   Args:
#     inputs: input Tensor, 4D, batch x height x width x channels.
#     state: state Tensor, 4D, batch x height x width x channels.
#     num_channels: the number of output channels in the layer.
#     filter_size: the shape of the each convolution filter.
#     forget_bias: the initial value of the forget biases.
#     scope: Optional scope for variable_scope.
#     reuse: whether or not the layer and the variables should be reused.
#   Returns:
#      a tuple of tensors representing output and the new state.
#   """
#   spatial_size = inputs.get_shape()[1:3]
#   if state is None:
#     state = init_state(inputs, list(spatial_size) + [2 * num_channels])
#   with tf.variable_scope(scope, default_name='BasicConvLstmCell', values=[inputs, state], reuse=reuse):
#     inputs.get_shape().assert_has_rank(4)
#     state.get_shape().assert_has_rank(4)
#     c, h = tf.split(3, 2, state)
#     inputs_h = tf.concat(3, [inputs, h])
#     # Parameters of gates are concatenated into one conv for efficiency.
#     i_j_f_o = layers.conv2d(inputs_h,
#                             4 * num_channels, [filter_size, filter_size],
#                             stride=1,
#                             activation_fn=None,
#                             scope='Gates', weights_initializer=initializer)
# 
#     # i = input_gate, j = new_input, f = forget_gate, o = output_gate
#     i, j, f, o = tf.split(3, 4, i_j_f_o)
# 
#     new_c = c * tf.sigmoid(f + forget_bias) + tf.sigmoid(i) * tf.tanh(j)
#     new_h = tf.tanh(new_c) * tf.sigmoid(o)
# 
#     return new_h, tf.concat(3, [new_c, new_h])


import numpy as np
import tensorflow as tf
import tensorflow.contrib.slim as slim
import tensorflow.contrib.layers
from tensorflow.contrib.layers.python import layers as tf_layers
from models.conv_lstm import basic_conv_lstm_cell
    initializer: specifies the initialization type (default: contrib.slim.layers uses Xavier init with uniform data)
    fc_conv_layer: adds an fc layer at the end of the encoder
  Returns:
    hidden4: hidden state of highest ConvLSTM layer
    fc_conv_layer: indicated whether a Fully Convolutional (8x8x16 -> 1x1x1024) shall be added
  """

  lstm_state1, lstm_state2, lstm_state3, lstm_state4 = None, None, None, None

  for i in range(sequence_length):


    frame = frames[:,i,:,:,:]

    reuse = (i > 0)

    with tf.variable_scope(scope, reuse=reuse):
      #LAYER 1: conv1
      conv1 = slim.layers.conv2d(frame, 16, [5, 5], stride=2, scope='conv1', normalizer_fn=tf_layers.layer_norm, weights_initializer=initializer,
          normalizer_params={'scope': 'layer_norm1'})

      #LAYER 2: convLSTM1
      hidden1, lstm_state1 = basic_conv_lstm_cell(conv1, lstm_state1, 16, initializer, filter_size=5, scope='convlstm1')
      hidden1 = tf_layers.layer_norm(hidden1, scope='layer_norm2')

      #LAYER 3: conv2
      conv2 = slim.layers.conv2d(hidden1, hidden1.get_shape()[3], [5, 5], stride=2, scope='conv2', normalizer_fn=tf_layers.layer_norm, weights_initializer=initializer,
                                  normalizer_params={'scope': 'layer_norm3'})

      #LAYER 4: convLSTM2

# Path: models/conv_lstm.py
# @add_arg_scope
# def basic_conv_lstm_cell(inputs,
#                          state,
#                          num_channels,
#                          initializer,
#                          filter_size=5,
#                          forget_bias=1.0,
#                          scope=None,
#                          reuse=None):
#   """Basic LSTM recurrent network cell, with 2D convolution connctions.
#   We add forget_bias (default: 1) to the biases of the forget gate in order to
#   reduce the scale of forgetting in the beginning of the training.
#   It does not allow cell clipping, a projection layer, and does not
#   use peep-hole connections: it is the basic baseline.
#   Args:
#     inputs: input Tensor, 4D, batch x height x width x channels.
#     state: state Tensor, 4D, batch x height x width x channels.
#     num_channels: the number of output channels in the layer.
#     filter_size: the shape of the each convolution filter.
#     forget_bias: the initial value of the forget biases.
#     scope: Optional scope for variable_scope.
#     reuse: whether or not the layer and the variables should be reused.
#   Returns:
#      a tuple of tensors representing output and the new state.
#   """
#   spatial_size = inputs.get_shape()[1:3]
#   if state is None:
#     state = init_state(inputs, list(spatial_size) + [2 * num_channels])
#   with tf.variable_scope(scope, default_name='BasicConvLstmCell', values=[inputs, state], reuse=reuse):
#     inputs.get_shape().assert_has_rank(4)
#     state.get_shape().assert_has_rank(4)
#     c, h = tf.split(3, 2, state)
#     inputs_h = tf.concat(3, [inputs, h])
#     # Parameters of gates are concatenated into one conv for efficiency.
#     i_j_f_o = layers.conv2d(inputs_h,
#                             4 * num_channels, [filter_size, filter_size],
#                             stride=1,
#                             activation_fn=None,
#                             scope='Gates', weights_initializer=initializer)
# 
#     # i = input_gate, j = new_input, f = forget_gate, o = output_gate
#     i, j, f, o = tf.split(3, 4, i_j_f_o)
# 
#     new_c = c * tf.sigmoid(f + forget_bias) + tf.sigmoid(i) * tf.tanh(j)
#     new_h = tf.tanh(new_c) * tf.sigmoid(o)
# 
#     return new_h, tf.concat(3, [new_c, new_h])


# Path: models/model_zoo/model_conv4_fc512.py
import numpy as np
import tensorflow as tf
import tensorflow.contrib.slim as slim
import tensorflow.contrib.layers
from tensorflow.contrib.layers.python import layers as tf_layers
from models.conv_lstm import basic_conv_lstm_cell

    initializer: specifies the initialization type (default: contrib.slim.layers uses Xavier init with uniform data)
    fc_conv_layer: adds an fc layer at the end of the encoder
  Returns:
    hidden4: hidden state of highest ConvLSTM layer
    fc_conv_layer: indicated whether a Fully Convolutional (8x8x16 -> 1x1x1024) shall be added
  """

  lstm_state1, lstm_state2, lstm_state3, lstm_state4 = None, None, None, None

  for i in range(sequence_length):


    frame = frames[:,i,:,:,:]

    reuse = (i > 0)

    with tf.variable_scope(scope, reuse=reuse):
      #LAYER 1: conv1
      conv1 = slim.layers.conv2d(frame, 16, [5, 5], stride=2, scope='conv1', normalizer_fn=tf_layers.layer_norm, weights_initializer=initializer,
          normalizer_params={'scope': 'layer_norm1'})

      #LAYER 2: convLSTM1
      hidden1, lstm_state1 = basic_conv_lstm_cell(conv1, lstm_state1, 16, initializer, filter_size=5, scope='convlstm1')
      hidden1 = tf_layers.layer_norm(hidden1, scope='layer_norm2')

      #LAYER 3: conv2
      conv2 = slim.layers.conv2d(hidden1, hidden1.get_shape()[3], [5, 5], stride=2, scope='conv2', normalizer_fn=tf_layers.layer_norm, weights_initializer=initializer,
                                  normalizer_params={'scope': 'layer_norm3'})

      #LAYER 4: convLSTM2

      hidden2, lstm_state2 = basic_conv_lstm_cell(conv2, lstm_state2, 16, initializer, filter_size=5, scope='convlstm2')