# 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/rlds_utils.py
# def to_rlds_step(prev_step: step_data.StepData,
#                  step: Optional[step_data.StepData]) -> Step:
# def _find_extra_shard(split_info: tfds.core.SplitInfo) -> Optional[Any]:
# def maybe_recover_last_shard(builder: tfds.core.DatasetBuilder):


import dataclasses
import os
import tensorflow as tf
import tensorflow_datasets as tfds
from typing import Any, Dict, List, Optional
from absl import logging
from envlogger import step_data
from envlogger.backends import backend_writer
from envlogger.backends import rlds_utils
    # Either writes the first metadata, or overwrites it.
    self._ds_info.write_to_directory(self._data_directory)

  def _finalize_shard_and_update_split(self) -> None:
    self._current_shard.close_writer()
    self._split.update(self._current_shard)
    self._current_shard = None
    self._write_split_metadata()

  def _finalize_episode_and_update_shard(self) -> None:
    serialized_example = self._current_episode.serialize_episode(
        self._ds_info.features, self._serializer)
    self._current_episode = None

    if self._current_shard is None:
      path = os.path.join(self._data_directory, self._split.get_shard_path())
      self._current_shard = Shard(writer=tf.io.TFRecordWriter(os.fspath(path)))

    self._current_shard.add_episode(serialized_example)

  def _write_and_reset_episode(self, is_last_episode: bool = False) -> None:
    if self._current_episode is None and self._current_shard is None:
      return

    if self._current_episode is not None:
      self._finalize_episode_and_update_shard()

    if is_last_episode or self._current_shard.num_episodes >= self._max_episodes_per_shard:
      self._finalize_shard_and_update_split()


# 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/rlds_utils.py
# def to_rlds_step(prev_step: step_data.StepData,
#                  step: Optional[step_data.StepData]) -> Step:
# def _find_extra_shard(split_info: tfds.core.SplitInfo) -> Optional[Any]:
# def maybe_recover_last_shard(builder: tfds.core.DatasetBuilder):


# Path: envlogger/backends/tfds_backend_writer.py
import dataclasses
import os
import tensorflow as tf
import tensorflow_datasets as tfds
from typing import Any, Dict, List, Optional
from absl import logging
from envlogger import step_data
from envlogger.backends import backend_writer
from envlogger.backends import rlds_utils

    # Either writes the first metadata, or overwrites it.
    self._ds_info.write_to_directory(self._data_directory)

  def _finalize_shard_and_update_split(self) -> None:
    self._current_shard.close_writer()
    self._split.update(self._current_shard)
    self._current_shard = None
    self._write_split_metadata()

  def _finalize_episode_and_update_shard(self) -> None:
    serialized_example = self._current_episode.serialize_episode(
        self._ds_info.features, self._serializer)
    self._current_episode = None

    if self._current_shard is None:
      path = os.path.join(self._data_directory, self._split.get_shard_path())
      self._current_shard = Shard(writer=tf.io.TFRecordWriter(os.fspath(path)))

    self._current_shard.add_episode(serialized_example)

  def _write_and_reset_episode(self, is_last_episode: bool = False) -> None:
    if self._current_episode is None and self._current_shard is None:
      return

    if self._current_episode is not None:
      self._finalize_episode_and_update_shard()

    if is_last_episode or self._current_shard.num_episodes >= self._max_episodes_per_shard:
      self._finalize_shard_and_update_split()


  def _record_step(self, data: step_data.StepData,