# Path: dreamer/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: dreamer/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 dreamer.tools import nested
from dreamer.tools import shape
# Copyright 2019 The Dreamer Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function




def closed_loop(cell, embedded, prev_action, debug=False):
  use_obs = tf.ones(tf.shape(embedded[:, :, :1])[:3], tf.bool)
  (prior, posterior), _ = tf.nn.dynamic_rnn(
      cell, (embedded, prev_action, use_obs), dtype=tf.float32)
  if debug:
    with tf.control_dependencies([tf.assert_equal(
        tf.shape(nested.flatten(posterior)[0])[1], tf.shape(embedded)[1])]):
      prior = nested.map(tf.identity, prior)

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

# Copyright 2019 The Dreamer Authors. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function




def closed_loop(cell, embedded, prev_action, debug=False):
  use_obs = tf.ones(tf.shape(embedded[:, :, :1])[:3], tf.bool)
  (prior, posterior), _ = tf.nn.dynamic_rnn(
      cell, (embedded, prev_action, use_obs), dtype=tf.float32)
  if debug:
    with tf.control_dependencies([tf.assert_equal(
        tf.shape(nested.flatten(posterior)[0])[1], tf.shape(embedded)[1])]):
      prior = nested.map(tf.identity, prior)

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