# Path: dreamer/tools/count_dataset.py
# def count_dataset(directory):
#   directory = os.path.expanduser(directory)
#   if not tf.gfile.Exists(directory):
#     message = "Data set directory '{}' does not exist."
#     raise ValueError(message.format(directory))
#   pattern = os.path.join(directory, '*.npz')
#   def func():
#     filenames = tf.gfile.Glob(pattern)
#     episodes = len(filenames)
#     episodes = np.array(episodes, dtype=np.int32)
#     return episodes
#   return tf.py_func(func, [], tf.int32)
# 
# Path: dreamer/tools/gif_summary.py
# def gif_summary(name, tensor, max_outputs, fps, collections=None, family=None):
#   tensor = tf.convert_to_tensor(tensor)
#   if tensor.dtype in (tf.float32, tf.float64):
#     tensor = tf.cast(255.0 * tensor, tf.uint8)
#   with summary_op_util.summary_scope(
#       name, family, values=[tensor]) as (tag, scope):
#     val = tf.py_func(
#         py_gif_summary,
#         [tag, tensor, max_outputs, fps],
#         tf.string,
#         stateful=False,
#         name=scope)
#     summary_op_util.collect(val, collections, [tf.GraphKeys.SUMMARIES])
#   return val
# 
# Path: dreamer/tools/image_strip_summary.py
# def image_strip_summary(name, images, max_length=100, max_batch=10):
#   if max_batch:
#     images = images[:max_batch]
#   if max_length:
#     images = images[:, :max_length]
#   if images.dtype == tf.uint8:
#     images = tf.cast(images, tf.float32) / 255.0
#   length, width = tf.shape(images)[1], tf.shape(images)[3]
#   channels = images.shape[-1].value
#   images = tf.transpose(images, [0, 2, 1, 3, 4])
#   images = tf.reshape(images, [1, -1, length * width, channels])
#   images = tf.clip_by_value(images, 0., 1.)
#   return tf.summary.image(name, images)
# 
# 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 matplotlib
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from dreamer.tools import count_dataset
from dreamer.tools import gif_summary
from dreamer.tools import image_strip_summary
from dreamer.tools import shape as shapelib
        summaries.append(tf.summary.scalar(name + '_inner_std', inner_std))
      if name in obs:
        log_prob = tf.reduce_mean(dist.log_prob(obs[name]))
        summaries.append(tf.summary.scalar(name + '_log_prob', log_prob))
        abs_error = tf.reduce_mean(tf.abs(mode - obs[name]))
        summaries.append(tf.summary.scalar(name + '_abs_error', abs_error))
  return summaries


def image_summaries(dist, target, name='image', max_batch=6):
  summaries = []
  with tf.variable_scope(name):
    image = dist.mode()[:max_batch]
    target = target[:max_batch]
    error = ((image - target) + 1) / 2
    # empty_frame = 0 * target[:max_batch, :1]
    # change = tf.concat([empty_frame, image[:, 1:] - image[:, :-1]], 1)
    # change = (change + 1) / 2
    summaries.append(image_strip_summary.image_strip_summary(
        'prediction', image))
    # summaries.append(image_strip_summary.image_strip_summary(
    #     'change', change))
    # summaries.append(image_strip_summary.image_strip_summary(
    #     'error', error))
    # Concat prediction and target vertically.
    frames = tf.concat([target, image, error], 2)
    # Stack batch entries horizontally.
    frames = tf.transpose(frames, [1, 2, 0, 3, 4])
    s = shapelib.shape(frames)
    frames = tf.reshape(frames, [s[0], s[1], s[2] * s[3], s[4]])

# Path: dreamer/tools/count_dataset.py
# def count_dataset(directory):
#   directory = os.path.expanduser(directory)
#   if not tf.gfile.Exists(directory):
#     message = "Data set directory '{}' does not exist."
#     raise ValueError(message.format(directory))
#   pattern = os.path.join(directory, '*.npz')
#   def func():
#     filenames = tf.gfile.Glob(pattern)
#     episodes = len(filenames)
#     episodes = np.array(episodes, dtype=np.int32)
#     return episodes
#   return tf.py_func(func, [], tf.int32)
# 
# Path: dreamer/tools/gif_summary.py
# def gif_summary(name, tensor, max_outputs, fps, collections=None, family=None):
#   tensor = tf.convert_to_tensor(tensor)
#   if tensor.dtype in (tf.float32, tf.float64):
#     tensor = tf.cast(255.0 * tensor, tf.uint8)
#   with summary_op_util.summary_scope(
#       name, family, values=[tensor]) as (tag, scope):
#     val = tf.py_func(
#         py_gif_summary,
#         [tag, tensor, max_outputs, fps],
#         tf.string,
#         stateful=False,
#         name=scope)
#     summary_op_util.collect(val, collections, [tf.GraphKeys.SUMMARIES])
#   return val
# 
# Path: dreamer/tools/image_strip_summary.py
# def image_strip_summary(name, images, max_length=100, max_batch=10):
#   if max_batch:
#     images = images[:max_batch]
#   if max_length:
#     images = images[:, :max_length]
#   if images.dtype == tf.uint8:
#     images = tf.cast(images, tf.float32) / 255.0
#   length, width = tf.shape(images)[1], tf.shape(images)[3]
#   channels = images.shape[-1].value
#   images = tf.transpose(images, [0, 2, 1, 3, 4])
#   images = tf.reshape(images, [1, -1, length * width, channels])
#   images = tf.clip_by_value(images, 0., 1.)
#   return tf.summary.image(name, images)
# 
# 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/summary.py
import matplotlib
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
from dreamer.tools import count_dataset
from dreamer.tools import gif_summary
from dreamer.tools import image_strip_summary
from dreamer.tools import shape as shapelib

        summaries.append(tf.summary.scalar(name + '_inner_std', inner_std))
      if name in obs:
        log_prob = tf.reduce_mean(dist.log_prob(obs[name]))
        summaries.append(tf.summary.scalar(name + '_log_prob', log_prob))
        abs_error = tf.reduce_mean(tf.abs(mode - obs[name]))
        summaries.append(tf.summary.scalar(name + '_abs_error', abs_error))
  return summaries


def image_summaries(dist, target, name='image', max_batch=6):
  summaries = []
  with tf.variable_scope(name):
    image = dist.mode()[:max_batch]
    target = target[:max_batch]
    error = ((image - target) + 1) / 2
    # empty_frame = 0 * target[:max_batch, :1]
    # change = tf.concat([empty_frame, image[:, 1:] - image[:, :-1]], 1)
    # change = (change + 1) / 2
    summaries.append(image_strip_summary.image_strip_summary(
        'prediction', image))
    # summaries.append(image_strip_summary.image_strip_summary(
    #     'change', change))
    # summaries.append(image_strip_summary.image_strip_summary(
    #     'error', error))
    # Concat prediction and target vertically.
    frames = tf.concat([target, image, error], 2)
    # Stack batch entries horizontally.
    frames = tf.transpose(frames, [1, 2, 0, 3, 4])
    s = shapelib.shape(frames)
    frames = tf.reshape(frames, [s[0], s[1], s[2] * s[3], s[4]])

    summaries.append(gif_summary.gif_summary(