# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/riegeli_backend_reader.py
# class RiegeliBackendReader(backend_reader.BackendReader):
#   def __init__(self, data_directory: str):
#   def _copy(self) -> 'RiegeliBackendReader':
#   def close(self):
#   def _decode_step_data(self, data: Tuple[Any, Any, Any]) -> step_data.StepData:
#   def _get_num_steps(self):
#   def _get_num_episodes(self):
#   def _get_nth_step(self, i: int) -> step_data.StepData:
#   def _get_nth_episode_info(self,
#                             i: int,
#                             include_metadata: bool = False
#                            ) -> episode_info.EpisodeInfo:
#   def metadata(self):
# 
# Path: envlogger/backends/riegeli_backend_writer.py
# class RiegeliBackendWriter(backend_writer.BackendWriter):
#   def __init__(
#       self,
#       data_directory: str,
#       max_episodes_per_file: int = 10000,
#       writer_options: str = 'transpose,brotli:6,chunk_size:1M',
#       flush_scheduler: Optional[schedulers.Scheduler] = None,
#       **base_kwargs,
#   ):
#   def _record_step(self, data: step_data.StepData,
#                    is_new_episode: bool) -> None:
#   def set_episode_metadata(self, data: Any) -> None:
#   def close(self) -> None:
# 
# Path: envlogger/backends/schedulers.py
# def n_step_scheduler(step_interval: int) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def bernoulli_step_scheduler(keep_probability: float) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def n_episode_scheduler(episode_interval: int) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:
# def bernoulli_episode_scheduler(keep_probability: float) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:
# def list_step_scheduler(
#     desired_steps: Union[List[int], np.ndarray]) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def list_episode_scheduler(
#     desired_episodes: Union[List[int], np.ndarray]) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:
# 
# 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 operator
import dm_env
import numpy as np
from typing import Any, List, Optional, Tuple
from absl import logging
from absl.testing import absltest
from absl.testing import parameterized
from envlogger import step_data
from envlogger.backends import riegeli_backend_reader
from envlogger.backends import riegeli_backend_writer
from envlogger.backends import schedulers
from envlogger.testing import catch_env

  def _collect_episode_data(
      self,
      env: Optional[dm_env.Environment] = None,
      num_episodes: int = 2,
      metadata: Optional[Any] = None,
      scheduler: Optional[schedulers.Scheduler] = None,
      max_episodes_per_file: int = 1,
      writer_options: str = 'transpose,brotli:6,chunk_size:1M'
  ) -> Tuple[List[step_data.StepData], str]:

    if env is None:
      logging.info('Creating Catch environment...')
      env = catch_env.Catch()
      logging.info('Done creating Catch environment.')
    temp_dir = self.create_tempdir()
    data_directory = temp_dir.full_path

    backend = riegeli_backend_writer.RiegeliBackendWriter(
        data_directory=data_directory,
        max_episodes_per_file=max_episodes_per_file,
        metadata=metadata,
        scheduler=scheduler,
        writer_options=writer_options)

    logging.info('Training a random agent for %r episodes...', num_episodes)
    num_actions = 3
    episodes_data = []
    for _ in range(num_episodes):
      timestep = env.reset()

# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/riegeli_backend_reader.py
# class RiegeliBackendReader(backend_reader.BackendReader):
#   def __init__(self, data_directory: str):
#   def _copy(self) -> 'RiegeliBackendReader':
#   def close(self):
#   def _decode_step_data(self, data: Tuple[Any, Any, Any]) -> step_data.StepData:
#   def _get_num_steps(self):
#   def _get_num_episodes(self):
#   def _get_nth_step(self, i: int) -> step_data.StepData:
#   def _get_nth_episode_info(self,
#                             i: int,
#                             include_metadata: bool = False
#                            ) -> episode_info.EpisodeInfo:
#   def metadata(self):
# 
# Path: envlogger/backends/riegeli_backend_writer.py
# class RiegeliBackendWriter(backend_writer.BackendWriter):
#   def __init__(
#       self,
#       data_directory: str,
#       max_episodes_per_file: int = 10000,
#       writer_options: str = 'transpose,brotli:6,chunk_size:1M',
#       flush_scheduler: Optional[schedulers.Scheduler] = None,
#       **base_kwargs,
#   ):
#   def _record_step(self, data: step_data.StepData,
#                    is_new_episode: bool) -> None:
#   def set_episode_metadata(self, data: Any) -> None:
#   def close(self) -> None:
# 
# Path: envlogger/backends/schedulers.py
# def n_step_scheduler(step_interval: int) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def bernoulli_step_scheduler(keep_probability: float) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def n_episode_scheduler(episode_interval: int) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:
# def bernoulli_episode_scheduler(keep_probability: float) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:
# def list_step_scheduler(
#     desired_steps: Union[List[int], np.ndarray]) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def list_episode_scheduler(
#     desired_episodes: Union[List[int], np.ndarray]) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:
# 
# 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/riegeli_backend_writer_test.py
import operator
import dm_env
import numpy as np
from typing import Any, List, Optional, Tuple
from absl import logging
from absl.testing import absltest
from absl.testing import parameterized
from envlogger import step_data
from envlogger.backends import riegeli_backend_reader
from envlogger.backends import riegeli_backend_writer
from envlogger.backends import schedulers
from envlogger.testing import catch_env


  def _collect_episode_data(
      self,
      env: Optional[dm_env.Environment] = None,
      num_episodes: int = 2,
      metadata: Optional[Any] = None,
      scheduler: Optional[schedulers.Scheduler] = None,
      max_episodes_per_file: int = 1,
      writer_options: str = 'transpose,brotli:6,chunk_size:1M'
  ) -> Tuple[List[step_data.StepData], str]:

    if env is None:
      logging.info('Creating Catch environment...')
      env = catch_env.Catch()
      logging.info('Done creating Catch environment.')
    temp_dir = self.create_tempdir()
    data_directory = temp_dir.full_path

    backend = riegeli_backend_writer.RiegeliBackendWriter(
        data_directory=data_directory,
        max_episodes_per_file=max_episodes_per_file,
        metadata=metadata,
        scheduler=scheduler,
        writer_options=writer_options)

    logging.info('Training a random agent for %r episodes...', num_episodes)
    num_actions = 3
    episodes_data = []
    for _ in range(num_episodes):
      timestep = env.reset()

      data = step_data.StepData(timestep, None, None)