# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/backend_reader.py
# T = TypeVar('T')
# class _SequenceAdapter(Generic[T], Sequence[T]):
# class BackendReader(metaclass=abc.ABCMeta):
#   def __init__(self, count: int, get_nth_item: Callable[[int], T]):
#   def __getitem__(self, index: Union[int, slice]) -> Union[T, List[T]]:
#   def __len__(self) -> int:
#   def __iter__(self) -> Iterator[T]:
#   def __next__(self) -> T:
#   def __init__(self):
#   def copy(self) -> 'BackendReader':
#   def _copy(self) -> 'BackendReader':
#   def _init_visitors(self):
#   def _get_nth_step(self, i: int) -> step_data.StepData:
#   def _get_num_steps(self) -> int:
#   def _get_num_episodes(self) -> int:
#   def _get_nth_episode_info(self,
#                             i: int,
#                             include_metadata: bool = False
#                            ) -> episode_info.EpisodeInfo:
#   def _get_nth_episode(self, i: int) -> Sequence[step_data.StepData]:
#     def get_nth_step_from_episode(j: int):
#   def _get_nth_episode_metadata(self, i: int) -> Optional[Any]:
#   def __enter__(self):
#   def __exit__(self, exc_type, exc_value, tb):
#   def __del__(self):
#   def close(self) -> None:
#   def metadata(self) -> Dict[str, Any]:
#   def episodes(self) -> Sequence[Sequence[step_data.StepData]]:
#   def episode_metadata(self) -> Sequence[Optional[Any]]:
#   def steps(self) -> Sequence[step_data.StepData]:
# 
# 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],


import copy
import dm_env
from typing import Any, Dict, Tuple
from absl import logging
from envlogger import step_data
from envlogger.backends import backend_reader
from envlogger.backends.python import episode_info
from envlogger.backends.python import riegeli_dataset_reader
from envlogger.converters import codec
from envlogger.proto import storage_pb2
from pybind11_abseil import status

    self._metadata = codec.decode(self._reader.metadata()) or {}
    super().__init__()

  def _copy(self) -> 'RiegeliBackendReader':
    c = copy.copy(self)

    c._metadata = copy.deepcopy(self._metadata)
    c._reader = self._reader.clone()

    return c

  def close(self):
    if self._reader is not None:
      self._reader.close()
    self._reader = None

  def _decode_step_data(self, data: Tuple[Any, Any, Any]) -> step_data.StepData:
    """Recovers dm_env.TimeStep from logged data (either dict or tuple)."""
    # Recover the TimeStep from the first tuple element.
    timestep = dm_env.TimeStep(
        dm_env.StepType(data[0][0]), data[0][1], data[0][2], data[0][3])
    return step_data.StepData(timestep, data[1], data[2])

  def _get_num_steps(self):
    return self._reader.num_steps

  def _get_num_episodes(self):
    return self._reader.num_episodes


# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/backend_reader.py
# T = TypeVar('T')
# class _SequenceAdapter(Generic[T], Sequence[T]):
# class BackendReader(metaclass=abc.ABCMeta):
#   def __init__(self, count: int, get_nth_item: Callable[[int], T]):
#   def __getitem__(self, index: Union[int, slice]) -> Union[T, List[T]]:
#   def __len__(self) -> int:
#   def __iter__(self) -> Iterator[T]:
#   def __next__(self) -> T:
#   def __init__(self):
#   def copy(self) -> 'BackendReader':
#   def _copy(self) -> 'BackendReader':
#   def _init_visitors(self):
#   def _get_nth_step(self, i: int) -> step_data.StepData:
#   def _get_num_steps(self) -> int:
#   def _get_num_episodes(self) -> int:
#   def _get_nth_episode_info(self,
#                             i: int,
#                             include_metadata: bool = False
#                            ) -> episode_info.EpisodeInfo:
#   def _get_nth_episode(self, i: int) -> Sequence[step_data.StepData]:
#     def get_nth_step_from_episode(j: int):
#   def _get_nth_episode_metadata(self, i: int) -> Optional[Any]:
#   def __enter__(self):
#   def __exit__(self, exc_type, exc_value, tb):
#   def __del__(self):
#   def close(self) -> None:
#   def metadata(self) -> Dict[str, Any]:
#   def episodes(self) -> Sequence[Sequence[step_data.StepData]]:
#   def episode_metadata(self) -> Sequence[Optional[Any]]:
#   def steps(self) -> Sequence[step_data.StepData]:
# 
# 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_reader.py
import copy
import dm_env
from typing import Any, Dict, Tuple
from absl import logging
from envlogger import step_data
from envlogger.backends import backend_reader
from envlogger.backends.python import episode_info
from envlogger.backends.python import riegeli_dataset_reader
from envlogger.converters import codec
from envlogger.proto import storage_pb2
from pybind11_abseil import status


    self._metadata = codec.decode(self._reader.metadata()) or {}
    super().__init__()

  def _copy(self) -> 'RiegeliBackendReader':
    c = copy.copy(self)

    c._metadata = copy.deepcopy(self._metadata)
    c._reader = self._reader.clone()

    return c

  def close(self):
    if self._reader is not None:
      self._reader.close()
    self._reader = None

  def _decode_step_data(self, data: Tuple[Any, Any, Any]) -> step_data.StepData:
    """Recovers dm_env.TimeStep from logged data (either dict or tuple)."""
    # Recover the TimeStep from the first tuple element.
    timestep = dm_env.TimeStep(
        dm_env.StepType(data[0][0]), data[0][1], data[0][2], data[0][3])
    return step_data.StepData(timestep, data[1], data[2])

  def _get_num_steps(self):
    return self._reader.num_steps

  def _get_num_episodes(self):
    return self._reader.num_episodes


  def _get_nth_step(self, i: int) -> step_data.StepData: