# 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 tensorflow as tf
from planet.tools import nested
from planet.tools import shape
      cell, (0 * embedded[:, context:], prev_action[:, context:], use_obs),
      initial_state=last_state)
  state = nested.map(
      lambda x, y: tf.concat([x, y], 1),
      closed_state, open_state)
  if debug:
    with tf.control_dependencies([tf.assert_equal(
        tf.shape(nested.flatten(state)[0])[1], tf.shape(embedded)[1])]):
      state = nested.map(tf.identity, state)
  return state


def planned(
    cell, objective_fn, embedded, prev_action, planner, context=1, length=20,
    amount=1000, debug=False):
  use_obs = tf.ones(tf.shape(embedded[:, :context, :1])[:3], tf.bool)
  (_, closed_state), last_state = tf.nn.dynamic_rnn(
      cell, (embedded[:, :context], prev_action[:, :context], use_obs),
      dtype=tf.float32)
  _, plan_state, return_ = planner(
      cell, objective_fn, last_state,
      obs_shape=shape.shape(embedded)[2:],
      action_shape=shape.shape(prev_action)[2:],
      horizon=length, amount=amount)
  state = nested.map(
      lambda x, y: tf.concat([x, y], 1),
      closed_state, plan_state)
  if debug:
    with tf.control_dependencies([tf.assert_equal(
        tf.shape(nested.flatten(state)[0])[1], context + length)]):

# 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/unroll.py
import tensorflow as tf
from planet.tools import nested
from planet.tools import shape

      cell, (0 * embedded[:, context:], prev_action[:, context:], use_obs),
      initial_state=last_state)
  state = nested.map(
      lambda x, y: tf.concat([x, y], 1),
      closed_state, open_state)
  if debug:
    with tf.control_dependencies([tf.assert_equal(
        tf.shape(nested.flatten(state)[0])[1], tf.shape(embedded)[1])]):
      state = nested.map(tf.identity, state)
  return state


def planned(
    cell, objective_fn, embedded, prev_action, planner, context=1, length=20,
    amount=1000, debug=False):
  use_obs = tf.ones(tf.shape(embedded[:, :context, :1])[:3], tf.bool)
  (_, closed_state), last_state = tf.nn.dynamic_rnn(
      cell, (embedded[:, :context], prev_action[:, :context], use_obs),
      dtype=tf.float32)
  _, plan_state, return_ = planner(
      cell, objective_fn, last_state,
      obs_shape=shape.shape(embedded)[2:],
      action_shape=shape.shape(prev_action)[2:],
      horizon=length, amount=amount)
  state = nested.map(
      lambda x, y: tf.concat([x, y], 1),
      closed_state, plan_state)
  if debug:
    with tf.control_dependencies([tf.assert_equal(
        tf.shape(nested.flatten(state)[0])[1], context + length)]):

      state = nested.map(tf.identity, state)