# coding=utf-8
# Copyright 2020 The TF-Agents Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Tests for tf_agents.replay_buffers.rlds_to_reverb."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from typing import Dict, List, Tuple

from absl.testing import parameterized
import reverb
from rlds import rlds_types
import tensorflow as tf
from tf_agents.replay_buffers import reverb_replay_buffer
from tf_agents.replay_buffers import reverb_utils
from tf_agents.replay_buffers import rlds_to_reverb
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import trajectory
from tf_agents.utils import test_utils

_SEQUENCE_LENGTH = 2
_STRIDE_LENGTH = 1
_REVERB_TABLE_SIZE = 10000
_REVERB_RATE_LIMITER = 1
_OBSERVATIONS = [[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]
_ACTIONS = [[11.0], [21.0], [31.0]]
_REWARDS = [1.0, 2.0, 3.0]
_DISCOUNTS = [1.0, 1.0, 1.0]
_EXTRA_PARAM = 'extra_param'


def generate_valid_episodes() -> (
    Dict[str, Tuple[tf.data.Dataset, List[trajectory.Trajectory]]]
):
  """Get test data for valid RLDS datasets to be used across differrent tests.

  Returns:
    A dict representing all valid cases as keys and tuples of valid input RLDS
    and expected TF-Agents trajectories as values.
  """
  complete_steps = tf.data.Dataset.from_tensor_slices({
      rlds_types.OBSERVATION: _OBSERVATIONS,
      rlds_types.ACTION: _ACTIONS,
      rlds_types.REWARD: _REWARDS,
      rlds_types.DISCOUNT: _DISCOUNTS,
      rlds_types.IS_TERMINAL: [False, False, True],
      rlds_types.IS_LAST: [False, False, True],
      rlds_types.IS_FIRST: [True, False, False],
  })
  complete_episode_trajectories = [
      trajectory.Trajectory(0, [1.0, 2.0], [11.0], (), 1, 1.0, 1.0),
      trajectory.Trajectory(1, [3.0, 4.0], [21.0], (), 2, 2.0, 1.0),
      trajectory.Trajectory(2, [5.0, 6.0], [31.0], (), 0, 3.0, 0.0),
  ]

  truncated_steps = tf.data.Dataset.from_tensor_slices({
      rlds_types.OBSERVATION: _OBSERVATIONS,
      rlds_types.ACTION: _ACTIONS,
      rlds_types.REWARD: _REWARDS,
      rlds_types.DISCOUNT: _DISCOUNTS,
      rlds_types.IS_TERMINAL: [False, False, False],
      rlds_types.IS_LAST: [False, False, True],
      rlds_types.IS_FIRST: [True, False, False],
  })
  truncated_episode_trajectories = [
      trajectory.Trajectory(0, [1.0, 2.0], [11.0], (), 1, 1.0, 1.0),
      trajectory.Trajectory(1, [3.0, 4.0], [21.0], (), 2, 2.0, 1.0),
      trajectory.Trajectory(2, [5.0, 6.0], [31.0], (), 0, 3.0, 1.0),
  ]

  single_step = tf.data.Dataset.from_tensor_slices({
      rlds_types.OBSERVATION: [[1.0, 2.0]],
      rlds_types.ACTION: [[11.0]],
      rlds_types.REWARD: [1.0],
      rlds_types.DISCOUNT: [1.0],
      rlds_types.IS_TERMINAL: [False],
      rlds_types.IS_LAST: [True],
      rlds_types.IS_FIRST: [True],
  })
  single_step_episode_trajectories = [
      trajectory.Trajectory(2, [1.0, 2.0], [11.0], (), 2, 1.0, 1.0),
  ]

  extra_param_steps = tf.data.Dataset.from_tensor_slices({
      rlds_types.OBSERVATION: [[1.0, 2.0]],
      rlds_types.ACTION: [[11.0]],
      rlds_types.REWARD: [1.0],
      rlds_types.DISCOUNT: [1.0],
      rlds_types.IS_TERMINAL: [False],
      rlds_types.IS_LAST: [True],
      rlds_types.IS_FIRST: [True],
      _EXTRA_PARAM: [True],
  })

  return {
      'complete_episode': (
          tf.data.Dataset.from_tensor_slices({
              rlds_types.STEPS: [complete_steps],
          }),
          complete_episode_trajectories,
      ),
      'truncated_episode': (
          tf.data.Dataset.from_tensor_slices({
              rlds_types.STEPS: [truncated_steps],
          }),
          truncated_episode_trajectories,
      ),
      'single_step_episode': (
          tf.data.Dataset.from_tensor_slices({
              rlds_types.STEPS: [single_step],
          }),
          single_step_episode_trajectories,
      ),
      'extra_param_episode': (
          tf.data.Dataset.from_tensor_slices({
              rlds_types.STEPS: [extra_param_steps],
          }),
          single_step_episode_trajectories,
      ),
      'multiple_episodes': (
          tf.data.Dataset.from_tensor_slices({
              rlds_types.STEPS: [complete_steps, single_step, truncated_steps],
          }),
          complete_episode_trajectories
          + single_step_episode_trajectories
          + truncated_episode_trajectories,
      ),
  }


def generate_invalid_episodes() -> Dict[str, Tuple[tf.data.Dataset, str]]:
  """Get test data for invalid RLDS datasets to be used across differrent tests.

  Returns:
    A dict representing all invalid cases as keys and tuples of invalid input
    RLDS and expected error messages as values.
  """
  return {
      'incorrect_ending': (
          tf.data.Dataset.from_tensor_slices({
              rlds_types.STEPS: [
                  tf.data.Dataset.from_tensor_slices({
                      rlds_types.OBSERVATION: _OBSERVATIONS,
                      rlds_types.ACTION: _ACTIONS,
                      rlds_types.REWARD: _REWARDS,
                      rlds_types.DISCOUNT: _DISCOUNTS,
                      rlds_types.IS_TERMINAL: [False, False, False],
                      rlds_types.IS_LAST: [False, False, False],
                      rlds_types.IS_FIRST: [True, False, True],
                  })
              ],
          }),
          'Mid step should not be followed by a first step.',
      ),
      'incorrect_termination': (
          tf.data.Dataset.from_tensor_slices({
              rlds_types.STEPS: [
                  tf.data.Dataset.from_tensor_slices({
                      rlds_types.OBSERVATION: _OBSERVATIONS,
                      rlds_types.ACTION: _ACTIONS,
                      rlds_types.REWARD: _REWARDS,
                      rlds_types.DISCOUNT: _DISCOUNTS,
                      rlds_types.IS_TERMINAL: [False, False, True],
                      rlds_types.IS_LAST: [False, False, False],
                      rlds_types.IS_FIRST: [True, False, False],
                  })
              ],
          }),
          'Terminal step must be the last step of an episode.',
      ),
      'incorrect_beginning': (
          tf.data.Dataset.from_tensor_slices({
              rlds_types.STEPS: [
                  tf.data.Dataset.from_tensor_slices({
                      rlds_types.OBSERVATION: _OBSERVATIONS,
                      rlds_types.ACTION: _ACTIONS,
                      rlds_types.REWARD: _REWARDS,
                      rlds_types.DISCOUNT: _DISCOUNTS,
                      rlds_types.IS_TERMINAL: [False, True, False],
                      rlds_types.IS_LAST: [False, True, False],
                      rlds_types.IS_FIRST: [True, False, False],
                  })
              ],
          }),
          'Last step of an episode must be followed by a first step.',
      ),
      'different_spec_episode': (
          tf.data.Dataset.from_tensor_slices({
              rlds_types.STEPS: [
                  tf.data.Dataset.from_tensor_slices({
                      rlds_types.OBSERVATION: [[1.0], [3.0], [5.0]],
                      rlds_types.ACTION: _ACTIONS,
                      rlds_types.REWARD: _REWARDS,
                      rlds_types.DISCOUNT: _DISCOUNTS,
                      rlds_types.IS_TERMINAL: [False, True, False],
                      rlds_types.IS_LAST: [False, True, False],
                      rlds_types.IS_FIRST: [True, False, False],
                  })
              ],
          }),
          'Replay buffer table signature spec should match RLDS data spec.',
      ),
      'no_step_episode': (
          tf.data.Dataset.from_tensors({'random': True}),
          f'No dataset representing RLDS {rlds_types.STEPS} exist in the data.',
      ),
      'incorrect_step_spec': (
          tf.data.Dataset.from_tensors(
              {
                  rlds_types.STEPS: tf.data.Dataset.from_tensors(
                      {'random1': [True, True], 'random2': [False, False]}
                  )
              }
          ),
          (
              'Invalid RLDS step spec. Features expected are'
              f' {rlds_to_reverb.get_rlds_step_features()}, but found'
              " ['random1', 'random2']"
          ),
      ),
  }


def get_policy_info_test_fn(
    current_step: Dict[str, tf.Tensor], next_step: Dict[str, tf.Tensor]
) -> Dict[str, tf.Tensor]:
  """Returns policy info for test function for policy info tests."""
  del next_step  # unused
  return {rlds_types.ACTION: current_step[rlds_types.ACTION]}


def get_policy_with_extra_param_test_fn(
    current_step: Dict[str, tf.Tensor], next_step: Dict[str, tf.Tensor]
) -> Dict[str, tf.Tensor]:
  """Returns policy info for test function with extra parameter."""
  del next_step
  return {
      rlds_types.ACTION: current_step[rlds_types.ACTION],
      _EXTRA_PARAM: current_step[_EXTRA_PARAM],
  }


class RldsToReverbTest(parameterized.TestCase, test_utils.TestCase):

  def setUp(self):
    super(RldsToReverbTest, self).setUp()

    self._valid_episodes = generate_valid_episodes()
    self._invalid_episodes = generate_invalid_episodes()

    # Data spec corresponding to our test data. This data spec is used for
    #  1) Validation of create_trajectory_data_spec.
    #  2) Initializing Reverb server and Reverb Replay Buffer.
    self._data_spec = trajectory.Trajectory(
        tensor_spec.TensorSpec(shape=(), dtype=tf.int32, name='step_type'),
        tensor_spec.TensorSpec(shape=(2,), dtype=tf.float32, name=None),
        tensor_spec.TensorSpec(shape=(1,), dtype=tf.float32, name=None),
        (),
        tensor_spec.TensorSpec(shape=(), dtype=tf.int32, name='step_type'),
        tensor_spec.TensorSpec(shape=(), dtype=tf.float32, name=None),
        tensor_spec.BoundedTensorSpec(
            shape=(),
            dtype=tf.float32,
            name='discount',
            minimum=[0.0],
            maximum=[1.0],
        ),
    )

    self._table_name = 'test_table'

    # Initialize Reverb server, Reverb Replay Buffer and Reverb Observer.
    uniform_table = reverb.Table(
        name=self._table_name,
        max_size=_REVERB_TABLE_SIZE,
        sampler=reverb.selectors.Uniform(),
        remover=reverb.selectors.Fifo(),
        rate_limiter=reverb.rate_limiters.MinSize(_REVERB_RATE_LIMITER),
        signature=tensor_spec.add_outer_dim(self._data_spec),
    )

    self._server = reverb.Server([uniform_table])

    replay_buffer = reverb_replay_buffer.ReverbReplayBuffer(
        self._data_spec,
        sequence_length=_SEQUENCE_LENGTH,
        table_name=self._table_name,
        local_server=self._server,
    )

    self._reverb_observer = reverb_utils.ReverbAddTrajectoryObserver(
        replay_buffer.py_client,
        self._table_name,
        sequence_length=_SEQUENCE_LENGTH,
        stride_length=_STRIDE_LENGTH,
    )

  def tearDown(self):
    if self._server:
      # Stop the Reverb server if it is running.
      self._server.stop()
      self._server = None
    super(RldsToReverbTest, self).tearDown()

  @parameterized.named_parameters(
      ('_complete_episode', 'complete_episode'),
      ('_truncated_episode', 'truncated_episode'),
      ('_single_step_episode', 'single_step_episode'),
      ('_extra_param_episode', 'extra_param_episode'),
      ('_multiple_episodes', 'multiple_episodes'),
  )
  def test_trajectory_data_spec_valid_episodes(self, episode):
    rlds_data, _ = self._valid_episodes[episode]
    self.assertEqual(
        rlds_to_reverb.create_trajectory_data_spec(rlds_data), self._data_spec
    )

  @parameterized.named_parameters(
      ('_no_step_episode', 'no_step_episode'),
      ('_incorrect_step_spec', 'incorrect_step_spec'),
  )
  def test_trajectory_data_spec_no_step_episode(self, episode):
    rlds_data, error_message = self._invalid_episodes[episode]
    with self.assertRaises(ValueError) as err:
      rlds_to_reverb.create_trajectory_data_spec(rlds_data)
    self.assertEqual(str(err.exception), error_message)

  @parameterized.named_parameters(
      ('_complete_episode', 'complete_episode', 3),
      ('_truncated_episode', 'truncated_episode', 3),
      ('_single_step_episode', 'single_step_episode', 1),
      ('_extra_param_episode', 'extra_param_episode', 1),
      ('_multiple_episodes', 'multiple_episodes', 7),
  )
  def test_push_to_reverb_valid_episodes(
      self, episode, expected_trajectories_pushed
  ):
    rlds_data, _ = self._valid_episodes[episode]
    trajectories_pushed = rlds_to_reverb.push_rlds_to_reverb(
        rlds_data, self._reverb_observer
    )  # type: int
    self.assertEqual(trajectories_pushed, expected_trajectories_pushed)

  @parameterized.named_parameters(
      ('_incorrect_ending', 'incorrect_ending'),
      ('_incorrect_termination', 'incorrect_termination'),
      ('_incorrect_beginning', 'incorrect_beginning'),
  )
  def test_push_to_reverb_invalid_episodes(self, episode):
    rlds_data, error_message = self._invalid_episodes[episode]
    with self.assertRaises(tf.errors.InvalidArgumentError) as err:
      rlds_to_reverb.push_rlds_to_reverb(rlds_data, self._reverb_observer)
    self.assertRegex(str(err.exception), error_message)

  @parameterized.named_parameters(
      ('_different_spec_episode', 'different_spec_episode'),
      ('_no_step_episode', 'no_step_episode'),
      ('_incorrect_step_spec', 'incorrect_step_spec'),
  )
  def test_push_to_reverb_invalid_episodes_value_errors(self, episode):
    rlds_data, error_message = self._invalid_episodes[episode]
    with self.assertRaises(ValueError) as err:
      rlds_to_reverb.push_rlds_to_reverb(rlds_data, self._reverb_observer)
    self.assertEqual(str(err.exception), error_message)

  @parameterized.named_parameters(
      ('_complete_episode', 'complete_episode', 3),
      ('_truncated_episode', 'truncated_episode', 3),
      ('_single_step_episode', 'single_step_episode', 1),
      ('_extra_param_episode', 'extra_param_episode', 1),
      ('_multiple_episodes', 'multiple_episodes', 7),
  )
  def test_push_to_reverb_with_policy(
      self, episode, expected_trajectories_pushed
  ):
    rlds_data, _ = self._valid_episodes[episode]
    trajectories_pushed = rlds_to_reverb.push_rlds_to_reverb(
        rlds_data, self._reverb_observer, get_policy_info_test_fn
    )  # type: int
    self.assertEqual(trajectories_pushed, expected_trajectories_pushed)


class RldsToTrajectoriesTest(parameterized.TestCase, test_utils.TestCase):

  def setUp(self):
    super(RldsToTrajectoriesTest, self).setUp()

    self._valid_episodes = generate_valid_episodes()
    self._invalid_episodes = generate_invalid_episodes()

  @parameterized.named_parameters(
      ('_complete_episode', 'complete_episode'),
      ('_truncated_episode', 'truncated_episode'),
      ('_single_step_episode', 'single_step_episode'),
      ('_extra_param_episode', 'extra_param_episode'),
      ('_multiple_episodes', 'multiple_episodes'),
  )
  def test_conversion_valid_episodes(self, episode):
    rlds_data, expected_trajectories = self._valid_episodes[episode]
    generated_trajectories = rlds_to_reverb.convert_rlds_to_trajectories(
        rlds_data
    )  # type: tf.data.Dataset
    for generated_trajectory, expected_trajectory in zip(
        list(generated_trajectories.as_numpy_iterator()), expected_trajectories
    ):
      self.assertEqual(
          generated_trajectory.step_type, expected_trajectory.step_type
      )
      self.assertEqual(
          generated_trajectory.next_step_type,
          expected_trajectory.next_step_type,
      )
      self.assertAllEqual(
          generated_trajectory.observation, expected_trajectory.observation
      )
      self.assertAllEqual(
          generated_trajectory.action, expected_trajectory.action
      )
      self.assertAllEqual(
          generated_trajectory.discount, expected_trajectory.discount
      )
      self.assertAllEqual(
          generated_trajectory.reward, expected_trajectory.reward
      )

  @parameterized.named_parameters(
      ('_incorrect_ending', 'incorrect_ending'),
      ('_incorrect_termination', 'incorrect_termination'),
      ('_incorrect_beginning', 'incorrect_beginning'),
  )
  def test_conversion_invalid_episodes(self, episode):
    rlds_data, error_message = self._invalid_episodes[episode]
    with self.assertRaises(tf.errors.InvalidArgumentError) as err:
      list(
          rlds_to_reverb.convert_rlds_to_trajectories(
              rlds_data
          ).as_numpy_iterator()
      )
    self.assertRegex(str(err.exception), error_message)

  @parameterized.named_parameters(
      ('_no_step_episode', 'no_step_episode'),
      ('_incorrect_step_spec', 'incorrect_step_spec'),
  )
  def test_conversion_no_step_episodes(self, episode):
    rlds_data, error_message = self._invalid_episodes[episode]
    with self.assertRaises(ValueError) as err:
      rlds_to_reverb.convert_rlds_to_trajectories(rlds_data)
    self.assertEqual(str(err.exception), error_message)

  @parameterized.named_parameters(
      ('_complete_episode', 'complete_episode'),
      ('_truncated_episode', 'truncated_episode'),
      ('_single_step_episode', 'single_step_episode'),
      ('_extra_param_episode', 'extra_param_episode'),
      ('_multiple_episodes', 'multiple_episodes'),
  )
  def test_conversion_with_policy(self, episode):
    rlds_data, expected_trajectories = self._valid_episodes[episode]
    generated_trajectories = rlds_to_reverb.convert_rlds_to_trajectories(
        rlds_data, get_policy_info_test_fn
    )  # type: tf.data.Dataset
    for generated_trajectory, expected_trajectory in zip(
        list(generated_trajectories.as_numpy_iterator()), expected_trajectories
    ):
      self.assertEqual(
          generated_trajectory.step_type, expected_trajectory.step_type
      )
      self.assertEqual(
          generated_trajectory.next_step_type,
          expected_trajectory.next_step_type,
      )
      self.assertAllEqual(
          generated_trajectory.observation, expected_trajectory.observation
      )
      self.assertAllEqual(
          generated_trajectory.action, expected_trajectory.action
      )
      self.assertAllEqual(
          generated_trajectory.discount, expected_trajectory.discount
      )
      self.assertAllEqual(
          generated_trajectory.reward, expected_trajectory.reward
      )
      self.assertAllEqual(
          generated_trajectory.policy_info,
          {rlds_types.ACTION: expected_trajectory.action},
      )

  def test_conversion_extra_param_policy(self):
    rlds_data, expected_trajectories = self._valid_episodes[
        'extra_param_episode'
    ]
    generated_trajectories = rlds_to_reverb.convert_rlds_to_trajectories(
        rlds_data, get_policy_with_extra_param_test_fn
    )  # type: tf.data.Dataset
    for generated_trajectory, expected_trajectory in zip(
        list(generated_trajectories.as_numpy_iterator()), expected_trajectories
    ):
      self.assertEqual(
          generated_trajectory.step_type, expected_trajectory.step_type
      )
      self.assertEqual(
          generated_trajectory.next_step_type,
          expected_trajectory.next_step_type,
      )
      self.assertAllEqual(
          generated_trajectory.observation, expected_trajectory.observation
      )
      self.assertAllEqual(
          generated_trajectory.action, expected_trajectory.action
      )
      self.assertAllEqual(
          generated_trajectory.discount, expected_trajectory.discount
      )
      self.assertAllEqual(
          generated_trajectory.reward, expected_trajectory.reward
      )
      self.assertAllEqual(
          generated_trajectory.policy_info,
          {rlds_types.ACTION: expected_trajectory.action, _EXTRA_PARAM: True},
      )


if __name__ == '__main__':
  test_utils.main()
