# Path: ldif/util/net_util.py
# def twin_inference(inference_fun, observation, element_count,
#                    flat_element_length, model_config):
# def get_normalizer_and_mode(model_config):
# def conv_layer(inputs, depth, model_config):
# def maxpool2x2_layer(inputs):
# def down_layer(inputs, depth, model_config):
# def up_layer(inputs, spatial_dims, depth, model_config):
# def encode(inputs,
#            model_config,
#            conv_layer_depths,
#            fc_layer_depths,
#            name='Encoder'):
# def inputs_to_feature_vector(inputs, feature_length, model_config):
# def decode(z,
#            model_config,
#            spatial_dims,
#            fc_layer_depths,
#            conv_layer_depths,
#            output_dim,
#            name='Decoder'):
# def residual_layer(inputs, num_outputs, model_config):
# 
# Path: ldif/util/math_util.py
# def int_log2(i):
# def nonzero_mean(tensor):
# def increase_frequency(t, out_dim, flatten=False, interleave=True):


import tensorflow as tf
from tensorflow.contrib import layers
from ldif.util import net_util
from ldif.util import math_util


def remove_batch_dim(tensor):
  shape = tensor.get_shape().as_list()
  rank = len(shape)
  ensure_shape(tensor, [1] + (rank-1)*[-1])
  tensor = tf.reshape(tensor, shape[1:])
  return tensor


def add_batch_dim(tensor):
  shape = tensor.get_shape().as_list()
  tensor = tf.reshape(tensor, [1] + shape)
  return tensor


def occnet_decoder(embedding, samples, apply_sigmoid, model_config):
  """Computes the OccNet output for the input embedding and its sample batch.

  Args:
    embedding: Tensor with shape [batch_size, shape_embedding_length].
    samples: Tensor with shape [batch_size, sample_count, 3].
    apply_sigmoid: Boolean. Whether to apply a sigmoid layer to the final linear
      activations.
    model_config: A ModelConfig object.

  Returns:
    Tensor with shape [batch_size, sample_count, 1].
  """
  if model_config.hparams.hyo == 't':

# Path: ldif/util/net_util.py
# def twin_inference(inference_fun, observation, element_count,
#                    flat_element_length, model_config):
# def get_normalizer_and_mode(model_config):
# def conv_layer(inputs, depth, model_config):
# def maxpool2x2_layer(inputs):
# def down_layer(inputs, depth, model_config):
# def up_layer(inputs, spatial_dims, depth, model_config):
# def encode(inputs,
#            model_config,
#            conv_layer_depths,
#            fc_layer_depths,
#            name='Encoder'):
# def inputs_to_feature_vector(inputs, feature_length, model_config):
# def decode(z,
#            model_config,
#            spatial_dims,
#            fc_layer_depths,
#            conv_layer_depths,
#            output_dim,
#            name='Decoder'):
# def residual_layer(inputs, num_outputs, model_config):
# 
# Path: ldif/util/math_util.py
# def int_log2(i):
# def nonzero_mean(tensor):
# def increase_frequency(t, out_dim, flatten=False, interleave=True):


# Path: ldif/nets/occnet.py
import tensorflow as tf
from tensorflow.contrib import layers
from ldif.util import net_util
from ldif.util import math_util



def remove_batch_dim(tensor):
  shape = tensor.get_shape().as_list()
  rank = len(shape)
  ensure_shape(tensor, [1] + (rank-1)*[-1])
  tensor = tf.reshape(tensor, shape[1:])
  return tensor


def add_batch_dim(tensor):
  shape = tensor.get_shape().as_list()
  tensor = tf.reshape(tensor, [1] + shape)
  return tensor


def occnet_decoder(embedding, samples, apply_sigmoid, model_config):
  """Computes the OccNet output for the input embedding and its sample batch.

  Args:
    embedding: Tensor with shape [batch_size, shape_embedding_length].
    samples: Tensor with shape [batch_size, sample_count, 3].
    apply_sigmoid: Boolean. Whether to apply a sigmoid layer to the final linear
      activations.
    model_config: A ModelConfig object.

  Returns:
    Tensor with shape [batch_size, sample_count, 1].
  """
  if model_config.hparams.hyo == 't':

    samples = math_util.nerfify(samples, 10, flatten=True, interleave=True)