# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/schedulers.py
# def n_step_scheduler(step_interval: int) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def bernoulli_step_scheduler(keep_probability: float) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def n_episode_scheduler(episode_interval: int) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:
# def bernoulli_episode_scheduler(keep_probability: float) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:
# def list_step_scheduler(
#     desired_steps: Union[List[int], np.ndarray]) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def list_episode_scheduler(
#     desired_episodes: Union[List[int], np.ndarray]) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:


from absl.testing import absltest
from absl.testing import parameterized
from envlogger import step_data
from envlogger.backends import schedulers
import dm_env
import numpy as np
  def test_bernoulli_steps_probability_99pct(self):
    """bernoulli_step_scheduler should return more True than False with p=0.99."""
    scheduler = schedulers.bernoulli_step_scheduler(
        keep_probability=0.99)
    num_true = 0
    for _ in range(1000):
      num_true += scheduler(None)
    num_false = 1000 - num_true
    self.assertGreater(num_true, num_false)

  @parameterized.named_parameters(
      ('negative_interval', -1),
      ('zero_interval', 0),
  )
  def test_n_episode_invalid_args(self, episode_interval):
    """n_episode_scheduler should raise an error if given invalid intervals."""
    self.assertRaises(
        ValueError,
        schedulers.n_episode_scheduler,
        episode_interval=episode_interval)

  def test_n_episode_interval_1(self):
    """n_episode_scheduler should always return True if interval is 1."""
    scheduler = schedulers.n_episode_scheduler(episode_interval=1)
    for _ in range(100):
      for timestep in _create_episode(num_transitions=np.random.randint(100)):
        self.assertTrue(scheduler(timestep))

  def test_n_episode_interval_2(self):
    """n_episode_scheduler should return True every other episode."""

# Path: envlogger/step_data.py
# class StepData(NamedTuple):
# 
# Path: envlogger/backends/schedulers.py
# def n_step_scheduler(step_interval: int) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def bernoulli_step_scheduler(keep_probability: float) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def n_episode_scheduler(episode_interval: int) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:
# def bernoulli_episode_scheduler(keep_probability: float) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:
# def list_step_scheduler(
#     desired_steps: Union[List[int], np.ndarray]) -> Scheduler:
#   def f(unused_data: step_data.StepData) -> bool:
# def list_episode_scheduler(
#     desired_episodes: Union[List[int], np.ndarray]) -> Scheduler:
#   def f(data: step_data.StepData) -> bool:


# Path: envlogger/backends/schedulers_test.py
from absl.testing import absltest
from absl.testing import parameterized
from envlogger import step_data
from envlogger.backends import schedulers
import dm_env
import numpy as np

  def test_bernoulli_steps_probability_99pct(self):
    """bernoulli_step_scheduler should return more True than False with p=0.99."""
    scheduler = schedulers.bernoulli_step_scheduler(
        keep_probability=0.99)
    num_true = 0
    for _ in range(1000):
      num_true += scheduler(None)
    num_false = 1000 - num_true
    self.assertGreater(num_true, num_false)

  @parameterized.named_parameters(
      ('negative_interval', -1),
      ('zero_interval', 0),
  )
  def test_n_episode_invalid_args(self, episode_interval):
    """n_episode_scheduler should raise an error if given invalid intervals."""
    self.assertRaises(
        ValueError,
        schedulers.n_episode_scheduler,
        episode_interval=episode_interval)

  def test_n_episode_interval_1(self):
    """n_episode_scheduler should always return True if interval is 1."""
    scheduler = schedulers.n_episode_scheduler(episode_interval=1)
    for _ in range(100):
      for timestep in _create_episode(num_transitions=np.random.randint(100)):
        self.assertTrue(scheduler(timestep))

  def test_n_episode_interval_2(self):
    """n_episode_scheduler should return True every other episode."""

    scheduler = schedulers.n_episode_scheduler(episode_interval=2)