# 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/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/converters/codec.py
# _SCALAR_DIM_SIZE = -438
# def _python_int_to_bytes(py_int: int) -> bytes:
# def _set_datum_values_from_scalar(scalar: Union[ScalarNumber, bool, str, bytes],
#                                   datum: storage_pb2.Datum) -> bool:
# def _set_datum_values_from_array(array: np.ndarray,
#                                  values: storage_pb2.Datum.Values) -> None:
# def encode(
#     user_data: Union[np.ndarray, List[Any], Tuple[Any, ...], Dict[str, Any]]
# ) -> storage_pb2.Data:
# def decode_datum(
#     datum: storage_pb2.Datum
# ) -> Union[np.ndarray, ScalarNumber, bool, str, bytes]:
#     def from_bigint(int_bytes):
# def decode(
#     user_data: storage_pb2.Data
# ) -> Optional[Union[ScalarNumber, bool, str, bytes, np.ndarray, List[Any],


from typing import Any, Optional
from absl import logging
from envlogger import step_data
from envlogger.backends import backend_writer
from envlogger.backends import schedulers
from envlogger.backends.python import riegeli_dataset_writer
from envlogger.converters import codec
from pybind11_abseil import status
      self._flush_scheduler = schedulers.bernoulli_step_scheduler(
          1.0 / 5000)
    else:
      self._flush_scheduler = flush_scheduler
    self._data_writer = riegeli_dataset_writer.RiegeliDatasetWriter()
    logging.info('self._data_directory: %r', self._data_directory)

    metadata = self._metadata or {}

    try:
      self._data_writer.init(
          data_dir=data_directory,
          metadata=codec.encode(metadata),
          max_episodes_per_shard=max_episodes_per_file,
          writer_options=writer_options)
    except status.StatusNotOk as e:
      logging.exception('exception: %r', e)

  def _record_step(self, data: step_data.StepData,
                   is_new_episode: bool) -> None:
    encoded_data = codec.encode(data)
    if not self._data_writer.add_step(encoded_data, is_new_episode):
      raise RuntimeError(
          'Failed to write `data`. Please see logs for more details.')

    if self._flush_scheduler is not None and not self._flush_scheduler(data):
      return
    self._data_writer.flush()

  def set_episode_metadata(self, data: Any) -> None:

# 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/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/converters/codec.py
# _SCALAR_DIM_SIZE = -438
# def _python_int_to_bytes(py_int: int) -> bytes:
# def _set_datum_values_from_scalar(scalar: Union[ScalarNumber, bool, str, bytes],
#                                   datum: storage_pb2.Datum) -> bool:
# def _set_datum_values_from_array(array: np.ndarray,
#                                  values: storage_pb2.Datum.Values) -> None:
# def encode(
#     user_data: Union[np.ndarray, List[Any], Tuple[Any, ...], Dict[str, Any]]
# ) -> storage_pb2.Data:
# def decode_datum(
#     datum: storage_pb2.Datum
# ) -> Union[np.ndarray, ScalarNumber, bool, str, bytes]:
#     def from_bigint(int_bytes):
# def decode(
#     user_data: storage_pb2.Data
# ) -> Optional[Union[ScalarNumber, bool, str, bytes, np.ndarray, List[Any],


# Path: envlogger/backends/riegeli_backend_writer.py
from typing import Any, Optional
from absl import logging
from envlogger import step_data
from envlogger.backends import backend_writer
from envlogger.backends import schedulers
from envlogger.backends.python import riegeli_dataset_writer
from envlogger.converters import codec
from pybind11_abseil import status

      self._flush_scheduler = schedulers.bernoulli_step_scheduler(
          1.0 / 5000)
    else:
      self._flush_scheduler = flush_scheduler
    self._data_writer = riegeli_dataset_writer.RiegeliDatasetWriter()
    logging.info('self._data_directory: %r', self._data_directory)

    metadata = self._metadata or {}

    try:
      self._data_writer.init(
          data_dir=data_directory,
          metadata=codec.encode(metadata),
          max_episodes_per_shard=max_episodes_per_file,
          writer_options=writer_options)
    except status.StatusNotOk as e:
      logging.exception('exception: %r', e)

  def _record_step(self, data: step_data.StepData,
                   is_new_episode: bool) -> None:
    encoded_data = codec.encode(data)
    if not self._data_writer.add_step(encoded_data, is_new_episode):
      raise RuntimeError(
          'Failed to write `data`. Please see logs for more details.')

    if self._flush_scheduler is not None and not self._flush_scheduler(data):
      return
    self._data_writer.flush()

  def set_episode_metadata(self, data: Any) -> None:

    encoded_data = codec.encode(data)