# 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.policies.batched_py_policy."""

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

from absl.testing import parameterized

import numpy as np
import tensorflow as tf  # pylint: disable=g-explicit-tensorflow-version-import

from tf_agents.policies import batched_py_policy
from tf_agents.policies import py_policy
from tf_agents.specs import array_spec
from tf_agents.trajectories import policy_step as ps
from tf_agents.trajectories import time_step as ts
from tf_agents.typing import types

MT_PARAMETERS = ({'multithreading': False},
                 {'multithreading': True})
NP_PARAMETERS = ({'multithreading': False, 'num_policies': 1},
                 {'multithreading': True, 'num_policies': 1},
                 {'multithreading': False, 'num_policies': 5},
                 {'multithreading': True, 'num_policies': 5},
                )


class MockPyPolicy(py_policy.PyPolicy):

  def __init__(self,
               time_step_spec: ts.TimeStep,
               action_spec: types.NestedArraySpec,
               policy_state_spec: types.NestedArraySpec = (),
               info_spec: types.NestedArraySpec = ()):
    seed = 987654321
    self._rng = np.random.RandomState(seed)
    super(MockPyPolicy, self).__init__(
        time_step_spec=time_step_spec,
        action_spec=action_spec,
        policy_state_spec=policy_state_spec)

  def _action(self,  # pytype: disable=signature-mismatch  # overriding-parameter-count-checks
              time_step: ts.TimeStep,
              policy_state: types.NestedArray) -> ps.PolicyStep:
    random_action = array_spec.sample_spec_nest(
        self._action_spec, self._rng)

    return ps.PolicyStep(random_action, policy_state)


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

  @property
  def time_step_spec(self):
    return ts.time_step_spec(
        observation_spec=array_spec.ArraySpec((1,), np.int32))

  @property
  def action_spec(self):
    return array_spec.BoundedArraySpec(
        [7], dtype=np.float32, minimum=-1.0, maximum=1.0)

  @property
  def policy_state_spec(self):
    return array_spec.BoundedArraySpec(
        [3], dtype=np.int16, minimum=-7.0, maximum=7.0)

  def _make_batched_py_policy(self, multithreading, num_policies=3):
    policies = []
    for _ in range(num_policies):
      policies.append(MockPyPolicy(self.time_step_spec,
                                   self.action_spec,
                                   self.policy_state_spec))
    return batched_py_policy.BatchedPyPolicy(
        policies=policies, multithreading=multithreading)

  @parameterized.parameters(*MT_PARAMETERS)
  def test_close_no_hang_after_init(self, multithreading):
    self._make_batched_py_policy(multithreading)

  @parameterized.parameters(*MT_PARAMETERS)
  def test_get_specs(self, multithreading):
    policy = self._make_batched_py_policy(multithreading)
    self.assertEqual(self.time_step_spec, policy.time_step_spec)
    self.assertEqual(self.action_spec, policy.action_spec)
    self.assertEqual(self.policy_state_spec, policy.policy_state_spec)

  @parameterized.parameters(*NP_PARAMETERS)
  def test_get_initial_state(self, multithreading, num_policies):
    policy = self._make_batched_py_policy(multithreading,
                                          num_policies=num_policies)
    policy_state = policy.get_initial_state()
    # Expect policy_state.shape[0] to be batch_size aka num_policies.
    # The remaining dimensions should match the policy_state_spec.
    correct_shape = (num_policies,) + self.policy_state_spec.shape
    self.assertEqual(correct_shape, policy_state.shape)

  @parameterized.parameters(*NP_PARAMETERS)
  def test_action(self, multithreading, num_policies):
    policy = self._make_batched_py_policy(multithreading,
                                          num_policies=num_policies)
    time_steps = np.array([
        ts.restart(observation=np.array([1]))
        for _ in range(num_policies)], dtype=object)

    # Call policy.action() and assert PolicySteps are batched correctly.
    policy_step = policy.action(time_steps)
    self.assertEqual(num_policies, policy_step.action.shape[0])

    # Take another step and assert that actions have the same shape.
    policy_step2 = policy.action(time_steps)
    self.assertAllEqual(policy_step.action.shape[0],
                        policy_step2.action.shape[0])


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