# Path: network/Util.py
# def conv2d(x, W, strides=(1, 1), padding="SAME"):
# def conv3d(x, W, strides=(1, 1), padding="SAME"):
# def conv2d_transpose(x, W, strides=(1, 1), padding="SAME"):
# def conv2d_dilated(x, W, dilation, padding="SAME"):
# def create_batch_norm_vars(n_out, tower_setup, scope_name="bn"):
# def get_activation(act_str):
# def prepare_input(inputs):
# def apply_dropout(inp, dropout):
# def max_pool(x, shape, strides=None, padding="SAME"):
# def max_pool3d(x, shape, strides=None, padding="SAME"):
# def bootstrapped_ce_loss(raw_ce, fraction, n_valid_pixels_per_im):
#   def bootstrapped_ce_for_one_img(args):
# def class_balanced_ce_loss(raw_ce, targets, n_classes):
#   def class_balanced_ce_for_one_img(args):


import tensorflow as tf
from network.Util import smart_shape
  # inputs should have shape (time, batch, height, width, feature)
  input_shape = smart_shape(inputs)
  num_units = cell.num_units()
  h, final_state = tf.nn.dynamic_rnn(cell, inputs, sequence_length, initial_state, dtype, parallel_iterations,
                                     swap_memory, time_major, scope)
  h = tf.reshape(h, tf.stack([input_shape[0], input_shape[1], input_shape[2], input_shape[3], num_units]))
  return h, final_state


# similar to https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/rnn/python/ops/core_rnn_cell_impl.py
# for maximal flexibility we allow to pass the weights externally
class ConvLSTMCell(RNNCell):
  def __init__(self, num_units, height, width, filter_size, forget_bias=1.0, activation=tf.tanh, W=None, b=None):
    self._num_units = num_units
    self._height = height
    self._width = width
    self._size = num_units * height * width
    self._forget_bias = forget_bias
    self._activation = activation
    self._filter_size = list(filter_size)
    if W is not None:
      W_shape = W.get_shape().as_list()
      assert len(W_shape) == 4
      assert W_shape[:2] == self._filter_size
      assert W_shape[-1] == 4 * self._num_units
      self._W = W
    else:
      self._W = None
    if b is not None:
      b_shape = b.get_shape().as_list()

# Path: network/Util.py
# def conv2d(x, W, strides=(1, 1), padding="SAME"):
# def conv3d(x, W, strides=(1, 1), padding="SAME"):
# def conv2d_transpose(x, W, strides=(1, 1), padding="SAME"):
# def conv2d_dilated(x, W, dilation, padding="SAME"):
# def create_batch_norm_vars(n_out, tower_setup, scope_name="bn"):
# def get_activation(act_str):
# def prepare_input(inputs):
# def apply_dropout(inp, dropout):
# def max_pool(x, shape, strides=None, padding="SAME"):
# def max_pool3d(x, shape, strides=None, padding="SAME"):
# def bootstrapped_ce_loss(raw_ce, fraction, n_valid_pixels_per_im):
#   def bootstrapped_ce_for_one_img(args):
# def class_balanced_ce_loss(raw_ce, targets, n_classes):
#   def class_balanced_ce_for_one_img(args):


# Path: network/layer_implementations/ConvLSTMCell.py
import tensorflow as tf
from network.Util import smart_shape

  # inputs should have shape (time, batch, height, width, feature)
  input_shape = smart_shape(inputs)
  num_units = cell.num_units()
  h, final_state = tf.nn.dynamic_rnn(cell, inputs, sequence_length, initial_state, dtype, parallel_iterations,
                                     swap_memory, time_major, scope)
  h = tf.reshape(h, tf.stack([input_shape[0], input_shape[1], input_shape[2], input_shape[3], num_units]))
  return h, final_state


# similar to https://github.com/tensorflow/tensorflow/blob/master/tensorflow/contrib/rnn/python/ops/core_rnn_cell_impl.py
# for maximal flexibility we allow to pass the weights externally
class ConvLSTMCell(RNNCell):
  def __init__(self, num_units, height, width, filter_size, forget_bias=1.0, activation=tf.tanh, W=None, b=None):
    self._num_units = num_units
    self._height = height
    self._width = width
    self._size = num_units * height * width
    self._forget_bias = forget_bias
    self._activation = activation
    self._filter_size = list(filter_size)
    if W is not None:
      W_shape = W.get_shape().as_list()
      assert len(W_shape) == 4
      assert W_shape[:2] == self._filter_size
      assert W_shape[-1] == 4 * self._num_units
      self._W = W
    else:
      self._W = None
    if b is not None:
      b_shape = b.get_shape().as_list()

      assert len(b_shape) == 1