# coding=utf-8
# Copyright 2018 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
#
#     http://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.

"""Stationary Stochastic Python Bandit environment with per-arm features."""
import gin
import numpy as np
import tensorflow as tf

from tf_agents.bandits.environments import bandit_py_environment
from tf_agents.bandits.specs import utils as bandit_spec_utils
from tf_agents.specs import array_spec
from tf_agents.trajectories import time_step as ts

GLOBAL_KEY = bandit_spec_utils.GLOBAL_FEATURE_KEY
PER_ARM_KEY = bandit_spec_utils.PER_ARM_FEATURE_KEY
NUM_ACTIONS_KEY = bandit_spec_utils.NUM_ACTIONS_FEATURE_KEY

VariableActionMethod = bandit_spec_utils.VariableActionMethod


@gin.configurable
class StationaryStochasticPerArmPyEnvironment(
    bandit_py_environment.BanditPyEnvironment):
  """Stationary Stochastic Bandit environment with per-arm features."""

  def __init__(
      self,
      global_context_sampling_fn,
      arm_context_sampling_fn,
      max_num_actions,
      reward_fn,
      num_actions_fn=None,
      batch_size=1,
      variable_action_method=VariableActionMethod.FIXED):
    """Initializes the environment.

    In each round, global context is generated by global_context_sampling_fn,
    per-arm contexts are generated by arm_context_sampling_fn. The reward_fn
    function takes the concatenation of a gloabl and a per-arm feature, and
    outputs a possibly random reward.
    In case `num_action_fn` is specified, the number of actions will be dynamic.
    The actual number of actions can be encoded in multiple ways, specified by
    `variable_action_method`. The observation spec constructed by the
    environment will also reflect the method used. The below list explains how
    the observations are built for all the methods.

    The different values of `variable_action_method` and the corresponding
    behavior:
    -- `FIXED` (default): The number of actions per sample is fixed. In this
       case, `num_actions_fn` should be `None`.
    -- 'MASK': The actually available actions are encoded by an action mask
       added to the observation in the format of
       `(observation, [1 1 ... 1 0 ... 0])`. The length of the mask, as well of
       the number of arm observations if `max_num_actions`.
    -- `NUM_ACTIONS_FEATURE`: An extra feature key `num_actions` is added to the
       observation, with an integer feature value indicating the number of
       available actions. The arm observation tensor has shape
       `[batch_size, max_num_actions, arm_feature_dim]`.
    -- `IN_BATCH_DIM`: The number of actions is folded into the batch dimension.
       In this case, the actual batch size should be 1, and the batch dimension
       is used to list all the actions for a sample. The global observation will
       internally be tiled to match this induced batch size. Also note that in
       this case, the `max_num_actions` parameter is ignored.


    Example:
      def global_context_sampling_fn():
        return np.random.randint(0, 10, [2])  # 2-dimensional global features.

      def arm_context_sampling_fn():
        return np.random.randint(-3, 4, [3])  # 3-dimensional arm features.

      def reward_fn(x):
        return sum(x)

      def num_actions_fn():
        return np.random.randint(2, 6)

      env = StationaryStochasticPerArmPyEnvironment(
          global_context_sampling_fn,
          arm_context_sampling_fn,
          5,
          reward_fn,
          num_actions_fn,
          VariableActionMethod.NUM_ACTIONS_FEATURE)

    Args:
      global_context_sampling_fn: A function that outputs a random 1d array or
        list of ints or floats. This output is the global context. Its shape and
        type must be consistent accross calls.
      arm_context_sampling_fn: A function that outputs a random 1 array or list
        of ints or floats (same type as the output of
        `global_context_sampling_fn`). This output is the per-arm context. Its
        shape must be consistent accross calls.
      max_num_actions: (int) the maximum number of actions in every sample. If
        `num_actions_fn` is not set, this many actions are available in every
        time step.
      reward_fn: A function that generates a reward when called with an
        observation.
      num_actions_fn: If set, it should be a function that outputs a single
        integer specifying the number of actions for a given time step. The
        value output by this function will be capped between 1 and
        `max_num_actions`. The number of actions will be encoded based on the
        method specified in `variable_action_method`. The different encodings
        are explained in the documentation above.
      batch_size: The batch size.
      variable_action_method: An instance of `VariableActionMethod`. Determines
        the way variable number of actions are handled.
    """
    self._global_context_sampling_fn = global_context_sampling_fn
    self._arm_context_sampling_fn = arm_context_sampling_fn
    self._max_num_actions = max_num_actions
    self._reward_fn = reward_fn
    self._batch_size = batch_size
    self._num_actions_fn = num_actions_fn
    self._variable_action_method = variable_action_method

    observation_spec = self._create_observation_spec()

    action_spec = array_spec.BoundedArraySpec(
        shape=(),
        dtype=np.int32,
        minimum=0,
        maximum=max_num_actions - 1,
        name='action')

    super(StationaryStochasticPerArmPyEnvironment,
          self).__init__(observation_spec, action_spec)

  def batched(self):
    return True

  def _reset(self):
    """Returns a time step of type FIRST containing an observation."""
    observation = self._observe()
    first_observation = tf.nest.flatten(observation)[0]
    batch_size = first_observation.shape[0]
    return ts.restart(observation, batch_size=batch_size,
                      reward_spec=self.reward_spec())

  @property
  def batch_size(self):
    return self._batch_size

  def _create_obs_spec_fixed(self):
    spec = {
        GLOBAL_KEY:
            array_spec.ArraySpec.from_array(self._global_context_sampling_fn()),
        PER_ARM_KEY:
            array_spec.add_outer_dims_nest(
                array_spec.ArraySpec.from_array(
                    self._arm_context_sampling_fn()), (self._max_num_actions,))
    }
    return spec

  def _create_obs_spec_featured(self):
    num_actions_spec = array_spec.BoundedArraySpec(
        shape=(),
        dtype=np.dtype(type(self._num_actions_fn())),
        minimum=1,
        maximum=self._max_num_actions)
    spec = {
        GLOBAL_KEY:
            array_spec.ArraySpec.from_array(self._global_context_sampling_fn()),
        PER_ARM_KEY:
            array_spec.add_outer_dims_nest(
                array_spec.ArraySpec.from_array(
                    self._arm_context_sampling_fn()), (self._max_num_actions,)),
        NUM_ACTIONS_KEY:
            num_actions_spec
    }
    return spec

  def _create_obs_spec_masked(self):
    mask_spec = array_spec.BoundedArraySpec(
        shape=(self._max_num_actions,),
        dtype=np.int32,
        minimum=0,
        maximum=1)
    spec = ({
        GLOBAL_KEY:
            array_spec.ArraySpec.from_array(self._global_context_sampling_fn()),
        PER_ARM_KEY:
            array_spec.add_outer_dims_nest(
                array_spec.ArraySpec.from_array(
                    self._arm_context_sampling_fn()), (self._max_num_actions,))
    }, mask_spec)
    return spec

  def _create_obs_spec_in_batch(self):
    spec = {
        GLOBAL_KEY:
            array_spec.ArraySpec.from_array(self._global_context_sampling_fn()),
        PER_ARM_KEY:
            array_spec.ArraySpec.from_array(self._arm_context_sampling_fn())
    }
    return spec

  def _create_observation_spec(self):
    """Creates an observation spec respecting the variable action method."""
    if self._variable_action_method == VariableActionMethod.FIXED:
      assert self._num_actions_fn is None, (
          'If `variable_action_method` is set to `fixed`, we can not have a '
          'function determining the number of actions.')
      return self._create_obs_spec_fixed()
    assert self._num_actions_fn is not None, (
        'If `variable_action_method` is anything other than `fixed`, '
        'a `num_actions_fn` must be specified')
    if self._variable_action_method == VariableActionMethod.NUM_ACTIONS_FEATURE:
      return self._create_obs_spec_featured()
    if self._variable_action_method == VariableActionMethod.MASK:
      return self._create_obs_spec_masked()
    if self._variable_action_method == VariableActionMethod.IN_BATCH_DIM:
      assert self._batch_size == 1, ('Number of actions can be folded in the '
                                     'batch dimension only if batch size is 1.')
      return self._create_obs_spec_in_batch()
    raise ValueError('Unknown variable_action_method %s' %
                     self._variable_action_method)

  def _create_observation(self):
    """Creates an observation respecting the variable action method."""
    global_observation = np.stack(
        [self._global_context_sampling_fn() for _ in range(self._batch_size)])

    if self._variable_action_method == VariableActionMethod.IN_BATCH_DIM:
      num_actions = self._num_actions_fn()
      arm_observation = np.reshape(
          [self._arm_context_sampling_fn() for _ in range(num_actions)],
          (num_actions, -1))
      tiled_global_obs = np.tile(global_observation, (num_actions, 1))
      return {GLOBAL_KEY: tiled_global_obs, PER_ARM_KEY: arm_observation}

    arm_observation = np.reshape([
        self._arm_context_sampling_fn()
        for _ in range(self._batch_size * self._max_num_actions)
    ], (self._batch_size, self._max_num_actions, -1))

    if self._variable_action_method == VariableActionMethod.FIXED:
      return {GLOBAL_KEY: global_observation, PER_ARM_KEY: arm_observation}

    num_actions = [self._num_actions_fn() for _ in range(self._batch_size)]
    num_actions = np.maximum(num_actions, 1)
    num_actions = np.minimum(num_actions, self._max_num_actions)

    if self._variable_action_method == VariableActionMethod.NUM_ACTIONS_FEATURE:
      return {
          GLOBAL_KEY: global_observation,
          PER_ARM_KEY: arm_observation,
          NUM_ACTIONS_KEY: num_actions
      }

    if self._variable_action_method == VariableActionMethod.MASK:
      mask = np.array(
          np.less(range(self._max_num_actions), [[i] for i in num_actions]),
          dtype=np.int32)
      return ({
          GLOBAL_KEY: global_observation,
          PER_ARM_KEY: arm_observation
      }, mask)

  def _observe(self):
    self._observation = self._create_observation()
    return self._observation

  def _apply_action(self, action):
    if action.shape[0] != self.batch_size:
      raise ValueError('Number of actions must match batch size.')
    observation = self._observation
    if self._variable_action_method == VariableActionMethod.MASK:
      observation = self._observation[0]
    global_obs = observation[GLOBAL_KEY]
    arm_obs = observation[PER_ARM_KEY]
    if self._variable_action_method == VariableActionMethod.IN_BATCH_DIM:
      # If the number of actions is folded into the batch size, we need to take
      # only the first global observation (keeping the outer dimension), and
      # create a batch dimension for the arm observations.
      global_obs = global_obs[0:1]
      arm_obs = np.expand_dims(arm_obs, axis=0)
    batch_size_range = range(self.batch_size)
    arm_obs = arm_obs[batch_size_range, action, :]
    reward = np.stack([
        self._reward_fn(np.concatenate((global_obs[b, :], arm_obs[b, :])))
        for b in batch_size_range
    ])
    if self._variable_action_method == VariableActionMethod.IN_BATCH_DIM:
      actual_num_actions = tf.nest.flatten(observation)[0].shape[0]
      reward = np.tile(reward, (actual_num_actions))
    return reward
