# Path: planet/tools/nested.py
# def zip_(*structures, **kwargs):
# def map_(function, *structures, **kwargs):
#   def impl(function, *structures):
# def flatten_(structure):
# def filter_(predicate, *structures, **kwargs):
#   def impl(predicate, *structures):
# 
# Path: planet/tools/shape.py
# def shape(tensor):
#   static = tensor.get_shape().as_list()
#   dynamic = tf.unstack(tf.shape(tensor))
#   assert len(static) == len(dynamic)
#   combined = [d if s is None else s for s, d in zip(static, dynamic)]
#   return combined


import functools
import numpy as np
import tensorflow as tf
from planet.tools import nested
from planet.tools import shape

def _merge_dims(tensor, dims):
  """Flatten consecutive axes of a tensor trying to preserve static shapes."""
  if isinstance(tensor, (list, tuple, dict)):
    return nested.map(tensor, lambda x: _merge_dims(x, dims))
  tensor = tf.convert_to_tensor(tensor)
  if (np.array(dims) - min(dims) != np.arange(len(dims))).all():
    raise ValueError('Dimensions to merge must all follow each other.')
  start, end = dims[0], dims[-1]
  output = tf.reshape(tensor, tf.concat([
      tf.shape(tensor)[:start],
      [tf.reduce_prod(tf.shape(tensor)[start: end + 1])],
      tf.shape(tensor)[end + 1:]], axis=0))
  merged = tensor.shape[start: end + 1].as_list()
  output.set_shape(
      tensor.shape[:start].as_list() +
      [None if None in merged else np.prod(merged)] +
      tensor.shape[end + 1:].as_list())
  return output


def _pad_dims(tensor, rank):
  """Append empty dimensions to the tensor until it is of the given rank."""
  for _ in range(rank - tensor.shape.ndims):
    tensor = tensor[..., None]
  return tensor


def _restore_batch_dim(tensor, batch_size):
  """Split batch dimension out of the first dimension of a tensor."""

# Path: planet/tools/nested.py
# def zip_(*structures, **kwargs):
# def map_(function, *structures, **kwargs):
#   def impl(function, *structures):
# def flatten_(structure):
# def filter_(predicate, *structures, **kwargs):
#   def impl(predicate, *structures):
# 
# Path: planet/tools/shape.py
# def shape(tensor):
#   static = tensor.get_shape().as_list()
#   dynamic = tf.unstack(tf.shape(tensor))
#   assert len(static) == len(dynamic)
#   combined = [d if s is None else s for s, d in zip(static, dynamic)]
#   return combined


# Path: planet/tools/overshooting.py
import functools
import numpy as np
import tensorflow as tf
from planet.tools import nested
from planet.tools import shape


def _merge_dims(tensor, dims):
  """Flatten consecutive axes of a tensor trying to preserve static shapes."""
  if isinstance(tensor, (list, tuple, dict)):
    return nested.map(tensor, lambda x: _merge_dims(x, dims))
  tensor = tf.convert_to_tensor(tensor)
  if (np.array(dims) - min(dims) != np.arange(len(dims))).all():
    raise ValueError('Dimensions to merge must all follow each other.')
  start, end = dims[0], dims[-1]
  output = tf.reshape(tensor, tf.concat([
      tf.shape(tensor)[:start],
      [tf.reduce_prod(tf.shape(tensor)[start: end + 1])],
      tf.shape(tensor)[end + 1:]], axis=0))
  merged = tensor.shape[start: end + 1].as_list()
  output.set_shape(
      tensor.shape[:start].as_list() +
      [None if None in merged else np.prod(merged)] +
      tensor.shape[end + 1:].as_list())
  return output


def _pad_dims(tensor, rank):
  """Append empty dimensions to the tensor until it is of the given rank."""
  for _ in range(rank - tensor.shape.ndims):
    tensor = tensor[..., None]
  return tensor


def _restore_batch_dim(tensor, batch_size):
  """Split batch dimension out of the first dimension of a tensor."""

  initial = shape.shape(tensor)