# Path: planet/tools/attr_dict.py
# class AttrDict(dict):  # collections.OrderedDict
#   def __init__(self, *args, **kwargs):
#   def __getattr__(self, name):
#   def __setattr__(self, name, value):
#   def __getitem__(self, name):
#   def __setitem__(self, name, value):
#   def __repr__(self):
#   def get(self, key, default=None):
#   def untouched(self):
#   def unlocked(self):
#   def lock(self):
#   def unlock(self):
#   def summarize(self):
#   def update(self, mapping):
#   def copy(self, _unlocked=False):
#   def save(self, filename):
#   def load(cls, filename):
#   def _format_value(self, value):
# 
# Path: planet/tools/chunk_sequence.py
# def chunk_sequence(sequence, chunk_length, randomize=True, num_chunks=None):
#   """Split a nested dict of sequence tensors into a batch of chunks.
# 
#   This function does not expect a batch of sequences, but a single sequence. A
#   `length` key is added if it did not exist already. When `randomize` is set,
#   up to `chunk_length - 1` initial frames will be discarded. Final frames that
#   do not fit into a chunk are always discarded.
# 
#   Args:
#     sequence: Nested dict of tensors with time dimension.
#     chunk_length: Size of chunks the sequence will be split into.
#     randomize: Start chunking from a random offset in the sequence,
#         enforcing that at least one chunk is generated.
#     num_chunks: Optionally specify the exact number of chunks to be extracted
#         from the sequence. Requires input to be long enough.
# 
#   Returns:
#     Nested dict of sequence tensors with chunk dimension.
#   """
#   with tf.device('/cpu:0'):
#     if 'length' in sequence:
#       length = sequence.pop('length')
#     else:
#       length = tf.shape(nested.flatten(sequence)[0])[0]
#     if randomize:
#       if num_chunks is None:
#         num_chunks = tf.maximum(1, length // chunk_length - 1)
#       else:
#         num_chunks = num_chunks + 0 * length
#       used_length = num_chunks * chunk_length
#       max_offset = length - used_length
#       offset = tf.random_uniform((), 0, max_offset + 1, dtype=tf.int32)
#     else:
#       if num_chunks is None:
#         num_chunks = length // chunk_length
#       else:
#         num_chunks = num_chunks + 0 * length
#       used_length = num_chunks * chunk_length
#       max_offset = 0
#       offset = 0
#     clipped = nested.map(
#         lambda tensor: tensor[offset: offset + used_length],
#         sequence)
#     chunks = nested.map(
#         lambda tensor: tf.reshape(
#             tensor, [num_chunks, chunk_length] + tensor.shape[1:].as_list()),
#         clipped)
#     chunks['length'] = chunk_length * tf.ones((num_chunks,), dtype=tf.int32)
#     return chunks


import functools
import os
import random
import numpy as np
import tensorflow as tf
from scipy.ndimage import interpolation
from planet.tools import attr_dict
from planet.tools import chunk_sequence


def numpy_episodes(
    train_dir, test_dir, shape, reader=None, loader=None,
    num_chunks=None, preprocess_fn=None):
  """Read sequences stored as compressed Numpy files as a TensorFlow dataset.

  Args:
    train_dir: Directory containing NPZ files of the training dataset.
    test_dir: Directory containing NPZ files of the testing dataset.
    shape: Tuple of batch size and chunk length for the datasets.
    reader: Callable that reads an episode from a NPZ filename.
    loader: Generator that yields episodes.

  Returns:
    Structured data from numpy episodes as Tensors.
  """
  reader = reader or episode_reader
  loader = loader or cache_loader
  try:
    dtypes, shapes = _read_spec(reader, train_dir)
  except ZeroDivisionError:
    dtypes, shapes = _read_spec(reader, test_dir)
  train = tf.data.Dataset.from_generator(
      functools.partial(loader, reader, train_dir, shape[0]),
      dtypes, shapes)
  test = tf.data.Dataset.from_generator(
      functools.partial(loader, reader, test_dir, shape[0]),
      dtypes, shapes)
  chunking = lambda x: tf.data.Dataset.from_tensor_slices(

# Path: planet/tools/attr_dict.py
# class AttrDict(dict):  # collections.OrderedDict
#   def __init__(self, *args, **kwargs):
#   def __getattr__(self, name):
#   def __setattr__(self, name, value):
#   def __getitem__(self, name):
#   def __setitem__(self, name, value):
#   def __repr__(self):
#   def get(self, key, default=None):
#   def untouched(self):
#   def unlocked(self):
#   def lock(self):
#   def unlock(self):
#   def summarize(self):
#   def update(self, mapping):
#   def copy(self, _unlocked=False):
#   def save(self, filename):
#   def load(cls, filename):
#   def _format_value(self, value):
# 
# Path: planet/tools/chunk_sequence.py
# def chunk_sequence(sequence, chunk_length, randomize=True, num_chunks=None):
#   """Split a nested dict of sequence tensors into a batch of chunks.
# 
#   This function does not expect a batch of sequences, but a single sequence. A
#   `length` key is added if it did not exist already. When `randomize` is set,
#   up to `chunk_length - 1` initial frames will be discarded. Final frames that
#   do not fit into a chunk are always discarded.
# 
#   Args:
#     sequence: Nested dict of tensors with time dimension.
#     chunk_length: Size of chunks the sequence will be split into.
#     randomize: Start chunking from a random offset in the sequence,
#         enforcing that at least one chunk is generated.
#     num_chunks: Optionally specify the exact number of chunks to be extracted
#         from the sequence. Requires input to be long enough.
# 
#   Returns:
#     Nested dict of sequence tensors with chunk dimension.
#   """
#   with tf.device('/cpu:0'):
#     if 'length' in sequence:
#       length = sequence.pop('length')
#     else:
#       length = tf.shape(nested.flatten(sequence)[0])[0]
#     if randomize:
#       if num_chunks is None:
#         num_chunks = tf.maximum(1, length // chunk_length - 1)
#       else:
#         num_chunks = num_chunks + 0 * length
#       used_length = num_chunks * chunk_length
#       max_offset = length - used_length
#       offset = tf.random_uniform((), 0, max_offset + 1, dtype=tf.int32)
#     else:
#       if num_chunks is None:
#         num_chunks = length // chunk_length
#       else:
#         num_chunks = num_chunks + 0 * length
#       used_length = num_chunks * chunk_length
#       max_offset = 0
#       offset = 0
#     clipped = nested.map(
#         lambda tensor: tensor[offset: offset + used_length],
#         sequence)
#     chunks = nested.map(
#         lambda tensor: tf.reshape(
#             tensor, [num_chunks, chunk_length] + tensor.shape[1:].as_list()),
#         clipped)
#     chunks['length'] = chunk_length * tf.ones((num_chunks,), dtype=tf.int32)
#     return chunks


# Path: planet/tools/numpy_episodes.py
import functools
import os
import random
import numpy as np
import tensorflow as tf
from scipy.ndimage import interpolation
from planet.tools import attr_dict
from planet.tools import chunk_sequence



def numpy_episodes(
    train_dir, test_dir, shape, reader=None, loader=None,
    num_chunks=None, preprocess_fn=None):
  """Read sequences stored as compressed Numpy files as a TensorFlow dataset.

  Args:
    train_dir: Directory containing NPZ files of the training dataset.
    test_dir: Directory containing NPZ files of the testing dataset.
    shape: Tuple of batch size and chunk length for the datasets.
    reader: Callable that reads an episode from a NPZ filename.
    loader: Generator that yields episodes.

  Returns:
    Structured data from numpy episodes as Tensors.
  """
  reader = reader or episode_reader
  loader = loader or cache_loader
  try:
    dtypes, shapes = _read_spec(reader, train_dir)
  except ZeroDivisionError:
    dtypes, shapes = _read_spec(reader, test_dir)
  train = tf.data.Dataset.from_generator(
      functools.partial(loader, reader, train_dir, shape[0]),
      dtypes, shapes)
  test = tf.data.Dataset.from_generator(
      functools.partial(loader, reader, test_dir, shape[0]),
      dtypes, shapes)
  chunking = lambda x: tf.data.Dataset.from_tensor_slices(

      chunk_sequence.chunk_sequence(x, shape[1], True, num_chunks))