# 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/feature_ops.py
# def build_quantile_map(quantile_file_dir: str):
# def discard_fn(obs: types.Float):
# def identity_fn(obs: types.Float):
# def get_normalize_fn(quantile: List[float],
#                      with_sqrt: bool,
#                      with_z_score_normalization: bool,
#                      eps: float = 1e-8,
#                      preprocessing_fn: Optional[Callable[[types.Tensor],
#                                                          types.Float]] = None):
#   def normalize(obs: types.Float):


import os
import numpy as np
import tensorflow as tf
from absl.testing import parameterized
from compiler_opt.rl import constant
from compiler_opt.rl import feature_ops
class FeatureUtilsTest(tf.test.TestCase, parameterized.TestCase):

  def setUp(self):
    self._quantile_file_dir = os.path.join(constant.BASE_DIR, 'testdata')
    super(FeatureUtilsTest, self).setUp()

  def test_build_quantile_map(self):
    quantile_map = feature_ops.build_quantile_map(self._quantile_file_dir)

    self.assertLen(quantile_map, 1)

    self.assertIn('edge_count', quantile_map)

    quantile = quantile_map['edge_count']

    # quantile
    self.assertLen(quantile, 9)
    self.assertEqual(2, quantile[0])
    self.assertEqual(8, quantile[6])

  def test_discard_fn(self):
    # obs in shape of [2, 1].
    obs = tf.constant(value=[[2.0], [8.0]])
    output = feature_ops.discard_fn(obs)

    self.assertAllEqual([2, 1, 0], output.shape)

  def test_identity_fn(self):
    # obs in shape of [2, 1].
    obs = tf.constant(value=[[2.0], [8.0]])

# 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/feature_ops.py
# def build_quantile_map(quantile_file_dir: str):
# def discard_fn(obs: types.Float):
# def identity_fn(obs: types.Float):
# def get_normalize_fn(quantile: List[float],
#                      with_sqrt: bool,
#                      with_z_score_normalization: bool,
#                      eps: float = 1e-8,
#                      preprocessing_fn: Optional[Callable[[types.Tensor],
#                                                          types.Float]] = None):
#   def normalize(obs: types.Float):


# Path: compiler_opt/rl/feature_ops_test.py
import os
import numpy as np
import tensorflow as tf
from absl.testing import parameterized
from compiler_opt.rl import constant
from compiler_opt.rl import feature_ops

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

  def setUp(self):
    self._quantile_file_dir = os.path.join(constant.BASE_DIR, 'testdata')
    super(FeatureUtilsTest, self).setUp()

  def test_build_quantile_map(self):
    quantile_map = feature_ops.build_quantile_map(self._quantile_file_dir)

    self.assertLen(quantile_map, 1)

    self.assertIn('edge_count', quantile_map)

    quantile = quantile_map['edge_count']

    # quantile
    self.assertLen(quantile, 9)
    self.assertEqual(2, quantile[0])
    self.assertEqual(8, quantile[6])

  def test_discard_fn(self):
    # obs in shape of [2, 1].
    obs = tf.constant(value=[[2.0], [8.0]])
    output = feature_ops.discard_fn(obs)

    self.assertAllEqual([2, 1, 0], output.shape)

  def test_identity_fn(self):
    # obs in shape of [2, 1].
    obs = tf.constant(value=[[2.0], [8.0]])

    output = feature_ops.identity_fn(obs)