# 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.

"""Test for tf_agents.train.tf_metrics."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from absl.testing import parameterized
import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import
from tf_agents.metrics import tf_metrics
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import trajectory

from tensorflow.python.eager import context  # pylint: disable=g-direct-tensorflow-import  # TF internal


class TFDequeTest(tf.test.TestCase):

  def evaluate_last_func_if_graph_mode(self):
    if not tf.executing_eagerly():
      # Latest op in graph is the call op for fn so evaluate it.
      self.evaluate(tf.compat.v1.get_default_graph().get_operations()[-1])

  def test_data_is_zero(self):
    d = tf_metrics.TFDeque(3, tf.int32)
    self.evaluate(tf.compat.v1.global_variables_initializer())
    self.assertAllEqual([], self.evaluate(d.data))

  def test_rolls_over(self):
    d = tf_metrics.TFDeque(3, tf.int32)
    self.evaluate(tf.compat.v1.global_variables_initializer())

    d.add(1)
    self.evaluate_last_func_if_graph_mode()
    d.add(2)
    self.evaluate_last_func_if_graph_mode()
    d.add(3)
    self.evaluate_last_func_if_graph_mode()
    self.assertAllEqual([1, 2, 3], self.evaluate(d.data))

    d.add(4)
    self.evaluate_last_func_if_graph_mode()
    self.assertAllEqual([4, 2, 3], self.evaluate(d.data))

  def test_clear(self):
    d = tf_metrics.TFDeque(3, tf.int32)
    self.evaluate(tf.compat.v1.global_variables_initializer())

    d.add(1)
    self.evaluate_last_func_if_graph_mode()
    d.add(2)
    self.evaluate_last_func_if_graph_mode()
    d.add(3)
    self.evaluate_last_func_if_graph_mode()
    self.assertAllEqual([1, 2, 3], self.evaluate(d.data))

    d.clear()
    self.evaluate_last_func_if_graph_mode()
    self.assertAllEqual([], self.evaluate(d.data))

    d.add(4)
    self.evaluate_last_func_if_graph_mode()
    d.add(5)
    self.evaluate_last_func_if_graph_mode()
    self.assertAllEqual([4, 5], self.evaluate(d.data))

  def test_mean_not_full(self):
    d = tf_metrics.TFDeque(3, tf.int32)
    self.evaluate(tf.compat.v1.global_variables_initializer())

    d.add(2)
    self.evaluate_last_func_if_graph_mode()
    d.add(4)
    self.evaluate_last_func_if_graph_mode()
    self.assertEqual(3, self.evaluate(d.mean()))
    self.assertEqual(tf.int32, d.mean().dtype)

  def test_queue_empty(self):
    d = tf_metrics.TFDeque(3, tf.int32)
    self.evaluate(tf.compat.v1.global_variables_initializer())

    self.assertEqual(0, self.evaluate(d.mean()))
    self.assertEqual(tf.int32.min, self.evaluate(d.max()))
    self.assertEqual(tf.int32.max, self.evaluate(d.min()))
    self.assertEqual(tf.int32, d.mean().dtype)
    self.assertEqual(tf.int32, d.min().dtype)
    self.assertEqual(tf.int32, d.max().dtype)

  def test_mean_roll_over(self):
    d = tf_metrics.TFDeque(3, tf.float32)
    self.evaluate(tf.compat.v1.global_variables_initializer())

    d.add(1)
    self.evaluate_last_func_if_graph_mode()
    d.add(2)
    self.evaluate_last_func_if_graph_mode()
    d.add(3)
    self.evaluate_last_func_if_graph_mode()
    d.add(4)
    self.evaluate_last_func_if_graph_mode()

    self.assertEqual(3.0, self.evaluate(d.mean()))
    self.assertEqual(tf.float32, d.mean().dtype)

  def test_extend(self):
    d = tf_metrics.TFDeque(3, tf.float32)
    self.evaluate(tf.compat.v1.global_variables_initializer())

    d.extend([1, 2, 3, 4])
    self.evaluate_last_func_if_graph_mode()
    self.assertEqual(3.0, self.evaluate(d.mean()))
    self.assertEqual(tf.float32, d.mean().dtype)

  def test_min(self):
    d = tf_metrics.TFDeque(4, tf.float32)
    self.evaluate(tf.compat.v1.global_variables_initializer())
    d.extend([1, 2, 3, 4])
    self.evaluate_last_func_if_graph_mode()
    self.assertEqual(1.0, self.evaluate(d.min()))

  def test_max(self):
    d = tf_metrics.TFDeque(4, tf.float32)
    self.evaluate(tf.compat.v1.global_variables_initializer())
    d.extend([1, 2, 3, 4])
    self.evaluate_last_func_if_graph_mode()
    self.assertEqual(4.0, self.evaluate(d.max()))


class TFShapedDequeTest(tf.test.TestCase):

  def evaluate_last_func_if_graph_mode(self):
    if not tf.executing_eagerly():
      self.evaluate(tf.compat.v1.get_default_graph().get_operations()[-1])

  def test_data_is_zero(self):
    d = tf_metrics.TFDeque(3, tf.int32, shape=(2,))
    self.evaluate(tf.compat.v1.global_variables_initializer())
    values = self.evaluate(d.data)
    self.assertAllEqual((0, 2), values.shape)

  def test_rolls_over(self):
    d = tf_metrics.TFDeque(3, tf.int32, shape=(2,))
    self.evaluate(tf.compat.v1.global_variables_initializer())

    d.add([1, 1])
    self.evaluate_last_func_if_graph_mode()
    d.add([2, 2])
    self.evaluate_last_func_if_graph_mode()
    d.add([3, 3])
    self.evaluate_last_func_if_graph_mode()
    self.assertAllEqual([[1, 1], [2, 2], [3, 3]], self.evaluate(d.data))

    d.add([4, 4])
    self.evaluate_last_func_if_graph_mode()
    self.assertAllEqual([[4, 4], [2, 2], [3, 3]], self.evaluate(d.data))

  def test_clear(self):
    d = tf_metrics.TFDeque(3, tf.int32, shape=(2,))
    self.evaluate(tf.compat.v1.global_variables_initializer())

    d.add([1, 1])
    self.evaluate_last_func_if_graph_mode()

    d.clear()
    self.evaluate_last_func_if_graph_mode()
    self.assertAllEqual((0, 2), self.evaluate(d.data).shape)

    d.add([2, 2])
    self.evaluate_last_func_if_graph_mode()
    d.add([3, 3])
    self.evaluate_last_func_if_graph_mode()
    self.assertAllEqual([[2, 2], [3, 3]], self.evaluate(d.data))

  def test_mean_full(self):
    d = tf_metrics.TFDeque(3, tf.int32, shape=(2,))
    self.evaluate(tf.compat.v1.global_variables_initializer())

    d.add([1, 1])
    self.evaluate_last_func_if_graph_mode()
    d.add([2, 2])
    self.evaluate_last_func_if_graph_mode()
    d.add([3, 3])
    self.evaluate_last_func_if_graph_mode()
    self.assertAllEqual([2, 2], self.evaluate(d.mean()))
    self.assertEqual(tf.int32, d.mean().dtype)

  def test_mean_not_full(self):
    d = tf_metrics.TFDeque(3, tf.int32, shape=(2,))
    self.evaluate(tf.compat.v1.global_variables_initializer())

    d.add([1, 1])
    self.evaluate_last_func_if_graph_mode()
    d.add([3, 3])
    self.evaluate_last_func_if_graph_mode()

    self.assertAllEqual([2, 2], self.evaluate(d.mean()))
    self.assertEqual(tf.int32, d.mean().dtype)

  def test_mean_empty(self):
    d = tf_metrics.TFDeque(3, tf.int32, shape=(2,))
    self.evaluate(tf.compat.v1.global_variables_initializer())

    self.assertAllEqual([0, 0], self.evaluate(d.mean()))
    self.assertEqual(tf.int32, d.mean().dtype)

  def test_mean_roll_over(self):
    d = tf_metrics.TFDeque(3, tf.float32, shape=(2,))
    self.evaluate(tf.compat.v1.global_variables_initializer())

    d.add([1, 1])
    self.evaluate_last_func_if_graph_mode()
    d.add([2, 2])
    self.evaluate_last_func_if_graph_mode()
    d.add([3, 3])
    self.evaluate_last_func_if_graph_mode()
    d.add([4, 4])
    self.evaluate_last_func_if_graph_mode()

    self.assertAllEqual([3.0, 3.0], self.evaluate(d.mean()))
    self.assertEqual(tf.float32, d.mean().dtype)

  def test_extend(self):
    d = tf_metrics.TFDeque(3, tf.int32, shape=(2,))
    self.evaluate(tf.compat.v1.global_variables_initializer())

    d.extend([[1, 1], [2, 2], [3, 3], [4, 4]])
    self.evaluate_last_func_if_graph_mode()

    self.assertAllEqual([[4, 4], [2, 2], [3, 3]], self.evaluate(d.data))
    self.assertAllEqual([3, 3], self.evaluate(d.mean()))


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

  def evaluate_last_func_if_graph_mode(self):
    if not tf.executing_eagerly():
      self.evaluate(tf.compat.v1.get_default_graph().get_operations()[-1])

  def _create_trajectories(self):
    def _concat_nested_tensors(nest1, nest2):
      return tf.nest.map_structure(
          lambda t1, t2: tf.concat([t1, t2], axis=0), nest1, nest2
      )

    # Order of args for trajectory methods:
    # observation, action, policy_info, reward, discount
    ts0 = _concat_nested_tensors(
        trajectory.boundary(
            (),
            tf.constant([1]),
            (),
            tf.constant([0.0], dtype=tf.float32),
            [1.0],
        ),
        trajectory.boundary(
            (),
            tf.constant([2]),
            (),
            tf.constant([0.0], dtype=tf.float32),
            [1.0],
        ),
    )
    ts1 = _concat_nested_tensors(
        trajectory.first(
            (),
            tf.constant([2]),
            (),
            tf.constant([1.0], dtype=tf.float32),
            [1.0],
        ),
        trajectory.first(
            (),
            tf.constant([1]),
            (),
            tf.constant([2.0], dtype=tf.float32),
            [1.0],
        ),
    )
    ts2 = _concat_nested_tensors(
        trajectory.last(
            (),
            tf.constant([1]),
            (),
            tf.constant([3.0], dtype=tf.float32),
            [1.0],
        ),
        trajectory.last(
            (),
            tf.constant([1]),
            (),
            tf.constant([4.0], dtype=tf.float32),
            [1.0],
        ),
    )
    ts3 = _concat_nested_tensors(
        trajectory.boundary(
            (),
            tf.constant([2]),
            (),
            tf.constant([0.0], dtype=tf.float32),
            [1.0],
        ),
        trajectory.boundary(
            (),
            tf.constant([0]),
            (),
            tf.constant([0.0], dtype=tf.float32),
            [1.0],
        ),
    )
    ts4 = _concat_nested_tensors(
        trajectory.first(
            (),
            tf.constant([1]),
            (),
            tf.constant([5.0], dtype=tf.float32),
            [1.0],
        ),
        trajectory.first(
            (),
            tf.constant([1]),
            (),
            tf.constant([6.0], dtype=tf.float32),
            [1.0],
        ),
    )
    ts5 = _concat_nested_tensors(
        trajectory.last(
            (),
            tf.constant([1]),
            (),
            tf.constant([7.0], dtype=tf.float32),
            [1.0],
        ),
        trajectory.last(
            (),
            tf.constant([1]),
            (),
            tf.constant([8.0], dtype=tf.float32),
            [1.0],
        ),
    )

    return [ts0, ts1, ts2, ts3, ts4, ts5]

  def _create_misaligned_trajectories(self):
    def _concat_nested_tensors(nest1, nest2):
      return tf.nest.map_structure(
          lambda t1, t2: tf.concat([t1, t2], axis=0), nest1, nest2
      )

    # Order of args for trajectory methods:
    # observation, action, policy_info, reward, discount
    ts1 = _concat_nested_tensors(
        trajectory.first(
            (),
            tf.constant([2]),
            (),
            tf.constant([1.0], dtype=tf.float32),
            [1.0],
        ),
        trajectory.boundary(
            (),
            tf.constant([1]),
            (),
            tf.constant([0.0], dtype=tf.float32),
            [1.0],
        ),
    )
    ts2 = _concat_nested_tensors(
        trajectory.last(
            (),
            tf.constant([1]),
            (),
            tf.constant([3.0], dtype=tf.float32),
            [1.0],
        ),
        trajectory.first(
            (),
            tf.constant([1]),
            (),
            tf.constant([2.0], dtype=tf.float32),
            [1.0],
        ),
    )
    ts3 = _concat_nested_tensors(
        trajectory.boundary(
            (),
            tf.constant([2]),
            (),
            tf.constant([0.0], dtype=tf.float32),
            [1.0],
        ),
        trajectory.last(
            (),
            tf.constant([1]),
            (),
            tf.constant([4.0], dtype=tf.float32),
            [1.0],
        ),
    )

    return [ts1, ts2, ts3]

  @parameterized.named_parameters([
      (
          'testEnvironmentStepsGraph',
          context.graph_mode,
          tf_metrics.EnvironmentSteps,
          5,
          6,
          0.0,
      ),
      (
          'testNumberOfEpisodesGraph',
          context.graph_mode,
          tf_metrics.NumberOfEpisodes,
          4,
          2,
          0.0,
      ),
      (
          'testAverageReturnGraph',
          context.graph_mode,
          tf_metrics.AverageReturnMetric,
          6,
          9.0,
          0.0,
      ),
      (
          'testMaxReturnGraph',
          context.graph_mode,
          tf_metrics.MaxReturnMetric,
          6,
          14.0,
          tf.float32.min,
      ),
      (
          'testMinReturnGraph',
          context.graph_mode,
          tf_metrics.MinReturnMetric,
          6,
          4.0,
          tf.float32.max,
      ),
      (
          'testAverageEpisodeLengthGraph',
          context.graph_mode,
          tf_metrics.AverageEpisodeLengthMetric,
          6,
          2.0,
          0.0,
      ),
      (
          'testEnvironmentStepsEager',
          context.eager_mode,
          tf_metrics.EnvironmentSteps,
          5,
          6,
          0.0,
      ),
      (
          'testNumberOfEpisodesEager',
          context.eager_mode,
          tf_metrics.NumberOfEpisodes,
          4,
          2,
          0.0,
      ),
      (
          'testAverageReturnEager',
          context.eager_mode,
          tf_metrics.AverageReturnMetric,
          6,
          9.0,
          0.0,
      ),
      (
          'testMaxReturnEager',
          context.eager_mode,
          tf_metrics.MaxReturnMetric,
          6,
          14.0,
          tf.float32.min,
      ),
      (
          'testMinReturnEager',
          context.eager_mode,
          tf_metrics.MinReturnMetric,
          6,
          4.0,
          tf.float32.max,
      ),
      (
          'testAverageEpisodeLengthEager',
          context.eager_mode,
          tf_metrics.AverageEpisodeLengthMetric,
          6,
          2.0,
          0.0,
      ),
  ])
  def testMetric(
      self,
      run_mode,
      metric_class,
      num_trajectories,
      expected_result,
      empty_queue_expected_result,
  ):
    with run_mode():
      trajectories = self._create_trajectories()
      if metric_class in [
          tf_metrics.AverageReturnMetric,
          tf_metrics.MaxReturnMetric,
          tf_metrics.MinReturnMetric,
          tf_metrics.AverageEpisodeLengthMetric,
      ]:
        metric = metric_class(batch_size=2)
      else:
        metric = metric_class()
      self.evaluate(tf.compat.v1.global_variables_initializer())
      self.evaluate(metric.init_variables())
      for i in range(num_trajectories):
        self.evaluate(metric(trajectories[i]))

      self.assertEqual(expected_result, self.evaluate(metric.result()))
      metric.reset()
      self.evaluate_last_func_if_graph_mode()
      self.assertEqual(
          empty_queue_expected_result, self.evaluate(metric.result())
      )

  @parameterized.named_parameters([
      ('testActionRelativeFreqGraph', context.graph_mode),
      ('testActionRelativeFreqEager', context.eager_mode),
  ])
  def testChosenActionHistogram(self, run_mode):
    with run_mode():
      trajectories = self._create_trajectories()
      num_trajectories = 5
      expected_result = [1, 2, 2, 1, 1, 1, 2, 0, 1, 1]
      metric = tf_metrics.ChosenActionHistogram(buffer_size=10)
      self.evaluate(tf.compat.v1.global_variables_initializer())
      self.evaluate(metric.init_variables())
      for i in range(num_trajectories):
        self.evaluate(metric(trajectories[i]))

      self.assertAllEqual(expected_result, self.evaluate(metric.result()))
      metric.reset()
      self.evaluate_last_func_if_graph_mode()
      self.assertEmpty(self.evaluate(metric.result()))

  @parameterized.named_parameters([
      (
          'testAverageReturnMetricVectorGraph',
          context.graph_mode,
          6,
          tensor_spec.TensorSpec((2,), tf.float32, 'r'),
          18.0,
      ),
      (
          'testAverageReturnMetricVectorEager',
          context.eager_mode,
          6,
          tensor_spec.TensorSpec((5,), tf.float32, 'r'),
          45.0,
      ),
  ])
  def testAverageReturnMetricVector(
      self, run_mode, num_trajectories, reward_spec, expected_result
  ):
    with run_mode():
      trajectories = self._create_trajectories()
      multi_trajectories = []
      for traj in trajectories:
        new_reward = tf.stack(
            [traj.reward] * reward_spec.shape.as_list()[0], axis=1
        )
        new_traj = trajectory.Trajectory(
            step_type=traj.step_type,
            observation=traj.observation,
            action=traj.action,
            policy_info=traj.policy_info,
            next_step_type=traj.next_step_type,
            reward=new_reward,
            discount=traj.discount,
        )
        multi_trajectories.append(new_traj)

      metric = tf_metrics.AverageReturnMetric(batch_size=2)
      self.evaluate(tf.compat.v1.global_variables_initializer())
      self.evaluate(metric.init_variables())
      for i in range(num_trajectories):
        self.evaluate(metric(multi_trajectories[i]))

      self.assertAllClose(expected_result, self.evaluate(metric.result()))

  @parameterized.named_parameters([
      (
          'testAverageReturnMultiMetricGraph',
          context.graph_mode,
          6,
          tensor_spec.TensorSpec((2,), tf.float32, 'r'),
          [9.0, 9.0],
      ),
      (
          'testAverageReturnMultiMetricEager',
          context.eager_mode,
          6,
          tensor_spec.TensorSpec((2,), tf.float32, 'r'),
          [9.0, 9.0],
      ),
      (
          'testAverageReturnMultiMetricRewardSpecListGraph',
          context.graph_mode,
          6,
          [
              tensor_spec.TensorSpec((), tf.float32, 'r1'),
              tensor_spec.TensorSpec((), tf.float32, 'r2'),
          ],
          [9.0, 9.0],
      ),
      (
          'testAverageReturnMultiMetricRewardSpecListEager',
          context.eager_mode,
          6,
          [
              tensor_spec.TensorSpec((), tf.float32, 'r1'),
              tensor_spec.TensorSpec((), tf.float32, 'r2'),
          ],
          [9.0, 9.0],
      ),
      (
          'testAverageReturnMultiMetricRewardSpecDictEager',
          context.eager_mode,
          6,
          {
              'a': tensor_spec.TensorSpec((), tf.float32, 'r1'),
              'b': tensor_spec.TensorSpec((), tf.float32, 'r2'),
          },
          {'a': 9.0, 'b': 9.0},
      ),
  ])
  def testAverageReturnMultiMetric(
      self, run_mode, num_trajectories, reward_spec, expected_result
  ):
    with run_mode():
      trajectories = self._create_trajectories()
      multi_trajectories = []
      for traj in trajectories:
        if isinstance(reward_spec, list):
          new_reward = [traj.reward, traj.reward]
        elif isinstance(reward_spec, dict):
          new_reward = {'a': traj.reward, 'b': traj.reward}
        else:
          new_reward = tf.stack([traj.reward, traj.reward], axis=1)
        new_traj = trajectory.Trajectory(
            step_type=traj.step_type,
            observation=traj.observation,
            action=traj.action,
            policy_info=traj.policy_info,
            next_step_type=traj.next_step_type,
            reward=new_reward,
            discount=traj.discount,
        )
        multi_trajectories.append(new_traj)

      metric = tf_metrics.AverageReturnMultiMetric(reward_spec, batch_size=2)
      self.evaluate(tf.compat.v1.global_variables_initializer())
      self.evaluate(metric.init_variables())
      for i in range(num_trajectories):
        self.evaluate(metric(multi_trajectories[i]))

      self.assertAllClose(expected_result, self.evaluate(metric.result()))
      metric.reset()
      self.evaluate_last_func_if_graph_mode()
      reset_result = self.evaluate(
          tf.nest.map_structure(tf.zeros_like, expected_result)
      )
      self.assertAllClose(reset_result, self.evaluate(metric.result()))

  @parameterized.named_parameters([
      (
          'testAverageReturnMultiMetricTimeMisalignedGraph',
          context.graph_mode,
          3,
          tensor_spec.TensorSpec((2,), tf.float32, 'r'),
          [5.0, 5.0],
      ),
      (
          'testAverageReturnMultiMetricTimeMisalignedEager',
          context.eager_mode,
          3,
          tensor_spec.TensorSpec((2,), tf.float32, 'r'),
          [5.0, 5.0],
      ),
      (
          'testAverageReturnMultiMetricRewardSpecListTimeMisalignedGraph',
          context.graph_mode,
          3,
          [
              tensor_spec.TensorSpec((), tf.float32, 'r1'),
              tensor_spec.TensorSpec((), tf.float32, 'r2'),
          ],
          [5.0, 5.0],
      ),
      (
          'testAverageReturnMultiMetricRewardSpecListTimeMisalignedEager',
          context.eager_mode,
          3,
          [
              tensor_spec.TensorSpec((), tf.float32, 'r1'),
              tensor_spec.TensorSpec((), tf.float32, 'r2'),
          ],
          [5.0, 5.0],
      ),
  ])
  def testAverageReturnMultiMetricTimeMisalignment(
      self, run_mode, num_trajectories, reward_spec, expected_result
  ):
    with run_mode():
      trajectories = self._create_misaligned_trajectories()
      multi_trajectories = []
      for traj in trajectories:
        if isinstance(reward_spec, list):
          new_reward = [traj.reward, traj.reward]
        else:
          new_reward = tf.stack([traj.reward, traj.reward], axis=1)
        new_traj = trajectory.Trajectory(
            step_type=traj.step_type,
            observation=traj.observation,
            action=traj.action,
            policy_info=traj.policy_info,
            next_step_type=traj.next_step_type,
            reward=new_reward,
            discount=traj.discount,
        )
        multi_trajectories.append(new_traj)

      metric = tf_metrics.AverageReturnMultiMetric(reward_spec, batch_size=2)
      self.evaluate(tf.compat.v1.global_variables_initializer())
      self.evaluate(metric.init_variables())
      for i in range(num_trajectories):
        self.evaluate(metric(multi_trajectories[i]))

      self.assertAllEqual(expected_result, self.evaluate(metric.result()))
      metric.reset()
      self.evaluate_last_func_if_graph_mode()
      self.assertAllEqual([0.0, 0.0], self.evaluate(metric.result()))


if __name__ == '__main__':
  tf.test.main()
