# Path: compiler_opt/rl/constant.py
# BASE_DIR = 'compiler_opt/rl'
# BASE_MODULE_DIR = 'compiler_opt.rl'
# DELTA = 0.01
#   BEHAVIORAL_CLONE = 0
#   DQN = 1
#   PPO = 2
# class AgentName(enum.Enum):
# 
# Path: compiler_opt/rl/data_reader.py
# def _get_policy_info_parsing_dict(agent_name, action_spec):
# def _process_parsed_sequence_and_get_policy_info(parsed_sequence, agent_name,
#                                                  action_spec):
# def create_parser_fn(
#     agent_name: constant.AgentName, time_step_spec: types.NestedSpec,
#     action_spec: types.NestedSpec) -> Callable[[str], trajectory.Trajectory]:
#   def _parser_fn(serialized_proto):
# def create_sequence_example_dataset_fn(
#     agent_name: constant.AgentName, time_step_spec: types.NestedSpec,
#     action_spec: types.NestedSpec, batch_size: int, train_sequence_length: int
# ) -> Callable[[List[str]], tf.data.Dataset]:
#   def _sequence_example_dataset_fn(sequence_examples):
# def create_file_dataset_fn(
#     agent_name: constant.AgentName, time_step_spec: types.NestedSpec,
#     action_spec: types.NestedSpec, batch_size: int, train_sequence_length: int,
#     input_dataset) -> Callable[[List[str]], tf.data.Dataset]:
#   def _file_dataset_fn(data_path):
# def create_tfrecord_dataset_fn(
#     agent_name: constant.AgentName, time_step_spec: types.NestedSpec,
#     action_spec: types.NestedSpec, batch_size: int, train_sequence_length: int
# ) -> Callable[[List[str]], tf.data.Dataset]:


import os
import tensorflow as tf
from absl.testing import parameterized
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import time_step
from tf_agents.trajectories import trajectory
from compiler_opt.rl import constant
from compiler_opt.rl import data_reader
      ).int64_list.value.append(0)
    else:
      example.feature_lists.feature_list['live_interval_weight'].feature.add(
      ).float_list.value.append(1.23)
    example.feature_lists.feature_list['reward'].feature.add(
    ).float_list.value.append(2.3)
    if agent_name == constant.AgentName.PPO:
      if is_action_discrete:
        example.feature_lists.feature_list[
            'CategoricalProjectionNetwork_logits'].feature.add(
            ).float_list.value.extend([1.2, 3.4])
      else:
        example.feature_lists.feature_list[
            'NormalProjectionNetwork_scale'].feature.add(
            ).float_list.value.extend([1.2])
        example.feature_lists.feature_list[
            'NormalProjectionNetwork_loc'].feature.add(
            ).float_list.value.extend([3.4])
  return example


def _write_tmp_tfrecord(file_path, example, num_examples):
  with tf.io.TFRecordWriter(file_path) as file_writer:
    for _ in range(num_examples):
      file_writer.write(example.SerializeToString())


class DataReaderTest(tf.test.TestCase, parameterized.TestCase):

  def setUp(self):

# Path: compiler_opt/rl/constant.py
# BASE_DIR = 'compiler_opt/rl'
# BASE_MODULE_DIR = 'compiler_opt.rl'
# DELTA = 0.01
#   BEHAVIORAL_CLONE = 0
#   DQN = 1
#   PPO = 2
# class AgentName(enum.Enum):
# 
# Path: compiler_opt/rl/data_reader.py
# def _get_policy_info_parsing_dict(agent_name, action_spec):
# def _process_parsed_sequence_and_get_policy_info(parsed_sequence, agent_name,
#                                                  action_spec):
# def create_parser_fn(
#     agent_name: constant.AgentName, time_step_spec: types.NestedSpec,
#     action_spec: types.NestedSpec) -> Callable[[str], trajectory.Trajectory]:
#   def _parser_fn(serialized_proto):
# def create_sequence_example_dataset_fn(
#     agent_name: constant.AgentName, time_step_spec: types.NestedSpec,
#     action_spec: types.NestedSpec, batch_size: int, train_sequence_length: int
# ) -> Callable[[List[str]], tf.data.Dataset]:
#   def _sequence_example_dataset_fn(sequence_examples):
# def create_file_dataset_fn(
#     agent_name: constant.AgentName, time_step_spec: types.NestedSpec,
#     action_spec: types.NestedSpec, batch_size: int, train_sequence_length: int,
#     input_dataset) -> Callable[[List[str]], tf.data.Dataset]:
#   def _file_dataset_fn(data_path):
# def create_tfrecord_dataset_fn(
#     agent_name: constant.AgentName, time_step_spec: types.NestedSpec,
#     action_spec: types.NestedSpec, batch_size: int, train_sequence_length: int
# ) -> Callable[[List[str]], tf.data.Dataset]:


# Path: compiler_opt/rl/data_reader_test.py
import os
import tensorflow as tf
from absl.testing import parameterized
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import time_step
from tf_agents.trajectories import trajectory
from compiler_opt.rl import constant
from compiler_opt.rl import data_reader

      ).int64_list.value.append(0)
    else:
      example.feature_lists.feature_list['live_interval_weight'].feature.add(
      ).float_list.value.append(1.23)
    example.feature_lists.feature_list['reward'].feature.add(
    ).float_list.value.append(2.3)
    if agent_name == constant.AgentName.PPO:
      if is_action_discrete:
        example.feature_lists.feature_list[
            'CategoricalProjectionNetwork_logits'].feature.add(
            ).float_list.value.extend([1.2, 3.4])
      else:
        example.feature_lists.feature_list[
            'NormalProjectionNetwork_scale'].feature.add(
            ).float_list.value.extend([1.2])
        example.feature_lists.feature_list[
            'NormalProjectionNetwork_loc'].feature.add(
            ).float_list.value.extend([3.4])
  return example


def _write_tmp_tfrecord(file_path, example, num_examples):
  with tf.io.TFRecordWriter(file_path) as file_writer:
    for _ in range(num_examples):
      file_writer.write(example.SerializeToString())


class DataReaderTest(tf.test.TestCase, parameterized.TestCase):

  def setUp(self):

    self._agent_name = constant.AgentName.DQN