# 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
  num_bytes: int = 0

  def add_episode(self, serialized_episode: str) -> None:
    self.writer.write(serialized_episode)
    self.num_episodes += 1
    self.num_bytes += len(serialized_episode)

  def close_writer(self) -> None:
    self.writer.flush()
    self.writer.close()


@dataclasses.dataclass
class Episode(object):
  """Episode that is being constructed."""
  prev_step: step_data.StepData
  steps: Optional[List[rlds_utils.Step]] = None
  metadata: Optional[Dict[str, Any]] = None

  def add_step(self, step: step_data.StepData) -> None:
    rlds_step = rlds_utils.to_rlds_step(self.prev_step, step)
    if self.steps is None:
      self.steps = []
    self.steps.append(rlds_step)
    self.prev_step = step

  def serialize_episode(
      self, features: tfds.features.FeaturesDict,
      serializer: tfds.core.example_serializer.ExampleSerializer) -> str:
    """Serializes an episode."""

# 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

  num_bytes: int = 0

  def add_episode(self, serialized_episode: str) -> None:
    self.writer.write(serialized_episode)
    self.num_episodes += 1
    self.num_bytes += len(serialized_episode)

  def close_writer(self) -> None:
    self.writer.flush()
    self.writer.close()


@dataclasses.dataclass
class Episode(object):
  """Episode that is being constructed."""
  prev_step: step_data.StepData
  steps: Optional[List[rlds_utils.Step]] = None
  metadata: Optional[Dict[str, Any]] = None

  def add_step(self, step: step_data.StepData) -> None:
    rlds_step = rlds_utils.to_rlds_step(self.prev_step, step)
    if self.steps is None:
      self.steps = []
    self.steps.append(rlds_step)
    self.prev_step = step

  def serialize_episode(
      self, features: tfds.features.FeaturesDict,
      serializer: tfds.core.example_serializer.ExampleSerializer) -> str:
    """Serializes an episode."""

    last_step = rlds_utils.to_rlds_step(self.prev_step, None)