# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/backend_writer.py
# class BackendWriter(metaclass=abc.ABCMeta):
#   def __init__(self,
#                metadata: Optional[Dict[str, Any]] = None,
#                scheduler: Optional[schedulers.Scheduler] = None):
#   def record_step(self, data: step_data.StepData, is_new_episode: bool) -> None:
#   def set_episode_metadata(self, data: Any) -> None:
#   def _record_step(self, data: step_data.StepData,
#                    is_new_episode: bool) -> None:
#   def close(self) -> None:
#   def __del__(self):
#   def metadata(self):
# 
# Path: envlogger/backends/tfds_backend_writer.py
# class RLDSEnvloggerBuilder(tfds.core.GeneratorBasedBuilder):
# class Shard(object):
# class Episode(object):
# class Split(object):
# class TFDSBackendWriter(backend_writer.BackendWriter):
#   def _info(self) -> tfds.core.DatasetInfo:
#   def _split_generators(self, dl_manager: tfds.download.DownloadManager):
#   def _generate_examples(self, path):
#   def add_episode(self, serialized_episode: str) -> None:
#   def close_writer(self) -> None:
#   def add_step(self, step: step_data.StepData) -> None:
#   def serialize_episode(
#       self, features: tfds.features.FeaturesDict,
#       serializer: tfds.core.example_serializer.ExampleSerializer) -> str:
#   def update(self, shard: Shard) -> None:
#   def get_shard_path(self) -> str:
#   def get_split_dict(self) -> tfds.core.splits.SplitDict:
# def initialize_split(split_name: Optional[str], data_directory: Any,
#                      builder_name: str) -> Split:
#   def __init__(self,
#                data_directory: str,
#                ds_config: tfds.rlds.rlds_base.DatasetConfig,
#                max_episodes_per_file: int = 1000,
#                split_name: Optional[str] = None,
#                version: str = '0.0.1',
#                **base_kwargs):
#   def _write_split_metadata(self) -> None:
#   def _finalize_shard_and_update_split(self) -> None:
#   def _finalize_episode_and_update_shard(self) -> None:
#   def _write_and_reset_episode(self, is_last_episode: bool = False) -> None:
#   def _record_step(self, data: step_data.StepData,
#                    is_new_episode: bool) -> None:
#   def set_episode_metadata(self, data: Dict[str, Any]) -> None:
#   def close(self) -> None:
# 
# Path: envlogger/testing/catch_env.py
# _ACTIONS = (-1, 0, 1)  # Left, no-op, right.
# class Catch(dm_env.Environment):
#   def __init__(self, rows=10, columns=5, seed=1):
#   def reset(self):
#   def step(self, action):
#   def observation_spec(self):
#   def action_spec(self):
#   def _observation(self):


import time
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
from typing import List, Optional
from absl import logging
from envlogger import step_data
from envlogger.backends import backend_writer
from envlogger.backends import tfds_backend_writer
from envlogger.testing import catch_env

def catch_env_tfds_config(
    name: str = 'catch_example') -> tfds.rlds.rlds_base.DatasetConfig:
  """Creates a TFDS DatasetConfig for the Catch environment."""
  return tfds.rlds.rlds_base.DatasetConfig(
      name=name,
      observation_info=tfds.features.Tensor(
          shape=(10, 5), dtype=tf.float32,
          encoding=tfds.features.Encoding.ZLIB),
      action_info=tf.int64,
      reward_info=tf.float64,
      discount_info=tf.float64,
      step_metadata_info={'timestamp': tf.int64},
      episode_metadata_info={'episode_id': tf.int64})


def tfds_backend_catch_env(
    data_directory: str,
    max_episodes_per_file: int = 1,
    split_name: Optional[str] = None) -> tfds_backend_writer.TFDSBackendWriter:
  """Creates a TFDS Backend Writer for the Catch Environment.

  Args:
    data_directory: directory where the data will be created (it has to exist).
    max_episodes_per_file: maximum number of episodes per file.
    split_name: number of the TFDS split to create.

  Returns:
    TFDS backend writer.
  """

# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/backend_writer.py
# class BackendWriter(metaclass=abc.ABCMeta):
#   def __init__(self,
#                metadata: Optional[Dict[str, Any]] = None,
#                scheduler: Optional[schedulers.Scheduler] = None):
#   def record_step(self, data: step_data.StepData, is_new_episode: bool) -> None:
#   def set_episode_metadata(self, data: Any) -> None:
#   def _record_step(self, data: step_data.StepData,
#                    is_new_episode: bool) -> None:
#   def close(self) -> None:
#   def __del__(self):
#   def metadata(self):
# 
# Path: envlogger/backends/tfds_backend_writer.py
# class RLDSEnvloggerBuilder(tfds.core.GeneratorBasedBuilder):
# class Shard(object):
# class Episode(object):
# class Split(object):
# class TFDSBackendWriter(backend_writer.BackendWriter):
#   def _info(self) -> tfds.core.DatasetInfo:
#   def _split_generators(self, dl_manager: tfds.download.DownloadManager):
#   def _generate_examples(self, path):
#   def add_episode(self, serialized_episode: str) -> None:
#   def close_writer(self) -> None:
#   def add_step(self, step: step_data.StepData) -> None:
#   def serialize_episode(
#       self, features: tfds.features.FeaturesDict,
#       serializer: tfds.core.example_serializer.ExampleSerializer) -> str:
#   def update(self, shard: Shard) -> None:
#   def get_shard_path(self) -> str:
#   def get_split_dict(self) -> tfds.core.splits.SplitDict:
# def initialize_split(split_name: Optional[str], data_directory: Any,
#                      builder_name: str) -> Split:
#   def __init__(self,
#                data_directory: str,
#                ds_config: tfds.rlds.rlds_base.DatasetConfig,
#                max_episodes_per_file: int = 1000,
#                split_name: Optional[str] = None,
#                version: str = '0.0.1',
#                **base_kwargs):
#   def _write_split_metadata(self) -> None:
#   def _finalize_shard_and_update_split(self) -> None:
#   def _finalize_episode_and_update_shard(self) -> None:
#   def _write_and_reset_episode(self, is_last_episode: bool = False) -> None:
#   def _record_step(self, data: step_data.StepData,
#                    is_new_episode: bool) -> None:
#   def set_episode_metadata(self, data: Dict[str, Any]) -> None:
#   def close(self) -> None:
# 
# Path: envlogger/testing/catch_env.py
# _ACTIONS = (-1, 0, 1)  # Left, no-op, right.
# class Catch(dm_env.Environment):
#   def __init__(self, rows=10, columns=5, seed=1):
#   def reset(self):
#   def step(self, action):
#   def observation_spec(self):
#   def action_spec(self):
#   def _observation(self):


# Path: envlogger/backends/tfds_backend_testlib.py
import time
import numpy as np
import tensorflow as tf
import tensorflow_datasets as tfds
from typing import List, Optional
from absl import logging
from envlogger import step_data
from envlogger.backends import backend_writer
from envlogger.backends import tfds_backend_writer
from envlogger.testing import catch_env


def catch_env_tfds_config(
    name: str = 'catch_example') -> tfds.rlds.rlds_base.DatasetConfig:
  """Creates a TFDS DatasetConfig for the Catch environment."""
  return tfds.rlds.rlds_base.DatasetConfig(
      name=name,
      observation_info=tfds.features.Tensor(
          shape=(10, 5), dtype=tf.float32,
          encoding=tfds.features.Encoding.ZLIB),
      action_info=tf.int64,
      reward_info=tf.float64,
      discount_info=tf.float64,
      step_metadata_info={'timestamp': tf.int64},
      episode_metadata_info={'episode_id': tf.int64})


def tfds_backend_catch_env(
    data_directory: str,
    max_episodes_per_file: int = 1,
    split_name: Optional[str] = None) -> tfds_backend_writer.TFDSBackendWriter:
  """Creates a TFDS Backend Writer for the Catch Environment.

  Args:
    data_directory: directory where the data will be created (it has to exist).
    max_episodes_per_file: maximum number of episodes per file.
    split_name: number of the TFDS split to create.

  Returns:
    TFDS backend writer.
  """

  return tfds_backend_writer.TFDSBackendWriter(