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

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

import os

import tensorflow as tf
import tensorflow_probability as tfp
from tf_agents.bandits.policies import mixture_policy
from tf_agents.policies import policy_saver
from tf_agents.policies import tf_policy
from tf_agents.specs import tensor_spec
from tf_agents.trajectories import policy_step
from tf_agents.trajectories import time_step as ts
from tf_agents.utils import test_utils

tfd = tfp.distributions


class ConstantPolicy(tf_policy.TFPolicy):
  """A policy that outputs a constant action, for testing purposes."""

  def __init__(self, action_spec, time_step_spec, action):
    self._constant_action = action
    super(ConstantPolicy, self).__init__(
        time_step_spec=time_step_spec,
        action_spec=action_spec,
        info_spec={
            'a': tensor_spec.TensorSpec(shape=(), dtype=tf.int32),
            'b': tensor_spec.TensorSpec(shape=(4,), dtype=tf.float32),
        },
    )

  def _variables(self):
    return []

  def _distribution(self, time_step, policy_state):
    raise NotImplementedError(
        'This policy outputs an action and not a distribution.'
    )

  def _action(self, time_step, policy_state, seed=None):
    batch_size = tf.compat.dimension_value(tf.shape(time_step.observation)[0])
    return policy_step.PolicyStep(
        tf.fill([batch_size], self._constant_action),
        policy_state,
        {
            'a': tf.range(batch_size, dtype=tf.int32),
            'b': tf.ones(shape=[batch_size, 4], dtype=tf.float32),
        },
    )


class MixturePolicyTest(test_utils.TestCase):

  def testMixturePolicyInconsistentSpecs(self):
    context_dim = 11
    observation_spec = tensor_spec.TensorSpec([context_dim], tf.float32)
    time_step_spec = ts.time_step_spec(observation_spec)
    action_spec = tensor_spec.BoundedTensorSpec(
        shape=(), dtype=tf.int32, minimum=0, maximum=9, name='action'
    )
    sub_policies = [
        ConstantPolicy(action_spec, time_step_spec, i) for i in range(9)
    ]
    wrong_obs_spec = tensor_spec.TensorSpec([context_dim + 1], tf.float32)
    wrong_time_step_spec = ts.time_step_spec(wrong_obs_spec)
    wrong_policy = ConstantPolicy(action_spec, wrong_time_step_spec, 9)
    sub_policies.append(wrong_policy)
    weights = [0, 0, 0.2, 0, 0, -0.3, 0, 0, 0.5, 0]
    dist = tfd.Categorical(probs=weights)
    with self.assertRaisesRegexp(
        AssertionError, 'Inconsistent time step specs'
    ):
      mixture_policy.MixturePolicy(dist, sub_policies)

  def testMixturePolicyChoices(self):
    context_dim = 34
    observation_spec = tensor_spec.TensorSpec([context_dim], tf.float32)
    time_step_spec = ts.time_step_spec(observation_spec)
    action_spec = tensor_spec.BoundedTensorSpec(
        shape=(), dtype=tf.int32, minimum=0, maximum=9, name='action'
    )
    sub_policies = [
        ConstantPolicy(action_spec, time_step_spec, i) for i in range(10)
    ]
    weights = [0, 0, 0.2, 0, 0, 0.3, 0, 0, 0.5, 0]
    dist = tfd.Categorical(probs=weights)
    policy = mixture_policy.MixturePolicy(dist, sub_policies)
    batch_size = 15
    time_step = ts.TimeStep(
        tf.constant(
            ts.StepType.FIRST,
            dtype=tf.int32,
            shape=[batch_size],
            name='step_type',
        ),
        tf.constant(0.0, dtype=tf.float32, shape=[batch_size], name='reward'),
        tf.constant(1.0, dtype=tf.float32, shape=[batch_size], name='discount'),
        tf.constant(
            list(range(batch_size * context_dim)),
            dtype=tf.float32,
            shape=[batch_size, context_dim],
            name='observation',
        ),
    )
    action_step = policy.action(time_step)
    actions, infos = self.evaluate([action_step.action, action_step.info])
    tf.nest.assert_same_structure(policy.info_spec, infos)
    self.assertAllEqual(actions.shape, [batch_size])
    self.assertAllInSet(actions, [2, 5, 8])

  def testMixturePolicyDynamicBatchSize(self):
    context_dim = 35
    observation_spec = tensor_spec.TensorSpec([context_dim], tf.float32)
    time_step_spec = ts.time_step_spec(observation_spec)
    action_spec = tensor_spec.BoundedTensorSpec(
        shape=(), dtype=tf.int32, minimum=0, maximum=9, name='action'
    )
    sub_policies = [
        ConstantPolicy(action_spec, time_step_spec, i) for i in range(10)
    ]
    weights = [0, 0, 0.2, 0, 0, 0.3, 0, 0, 0.5, 0]
    dist = tfd.Categorical(probs=weights)

    policy = mixture_policy.MixturePolicy(dist, sub_policies)
    batch_size = tf.random.uniform(
        shape=(), minval=10, maxval=15, dtype=tf.int32
    )
    time_step = ts.TimeStep(
        tf.fill(
            tf.expand_dims(batch_size, axis=0),
            ts.StepType.FIRST,
            name='step_type',
        ),
        tf.zeros(shape=[batch_size], dtype=tf.float32, name='reward'),
        tf.ones(shape=[batch_size], dtype=tf.float32, name='discount'),
        tf.reshape(
            tf.range(
                tf.cast(batch_size * context_dim, dtype=tf.float32),
                dtype=tf.float32,
            ),
            shape=[-1, context_dim],
            name='observation',
        ),
    )
    action_step = policy.action(time_step)
    actions, bsize = self.evaluate([action_step.action, batch_size])
    self.assertAllEqual(actions.shape, [bsize])
    self.assertAllInSet(actions, [2, 5, 8])

    train_step = tf.compat.v1.train.get_or_create_global_step()
    saver = policy_saver.PolicySaver(policy, train_step=train_step)
    location = os.path.join(self.get_temp_dir(), 'saved_policy')
    if not tf.executing_eagerly():
      with self.cached_session():
        self.evaluate(tf.compat.v1.global_variables_initializer())
        saver.save(location)
    else:
      saver.save(location)
    loaded_policy = tf.compat.v2.saved_model.load(location)
    new_batch_size = 3
    new_time_step = ts.TimeStep(
        tf.fill(
            tf.expand_dims(new_batch_size, axis=0),
            ts.StepType.FIRST,
            name='step_type',
        ),
        tf.zeros(shape=[new_batch_size], dtype=tf.float32, name='reward'),
        tf.ones(shape=[new_batch_size], dtype=tf.float32, name='discount'),
        tf.reshape(
            tf.range(
                tf.cast(new_batch_size * context_dim, dtype=tf.float32),
                dtype=tf.float32,
            ),
            shape=[-1, context_dim],
            name='observation',
        ),
    )
    new_action = self.evaluate(loaded_policy.action(new_time_step).action)
    self.assertAllEqual(new_action.shape, [new_batch_size])
    self.assertAllInSet(new_action, [2, 5, 8])


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