# Copyright 2019 DeepMind Technologies Ltd. All rights reserved.
#
# 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.

"""Computes the value of a given policy."""

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

from typing import Iterable, List, Union

import numpy as np

from open_spiel.python import policy

PROBABILITY_THRESHOLD = 0


def _child(state, action):
  """Returns a child state, handling the simultaneous node case."""
  if isinstance(action, Iterable):
    child = state.clone()
    child.apply_actions(action)
    return child
  else:
    return state.child(action)


def _transitions(state, policies):
  """Returns iterator over (action, prob) from the given state."""
  if state.is_chance_node():
    return state.chance_outcomes()
  elif state.is_simultaneous_node():
    return policy.joint_action_probabilities(state, policies)
  else:
    player = state.current_player()
    return policies[player].action_probabilities(state).items()


def policy_value(state, policies: Union[List[policy.Policy], policy.Policy]):
  """Returns the expected values for the state for players following `policies`.

  Computes the expected value of the`state` for each player, assuming player `i`
  follows the policy given in `policies[i]`.

  Args:
    state: A `pyspiel.State`.
    policies: A `list` of `policy.Policy` objects, one per player for sequential
      games, one policy for simulatenous games.

  Returns:
    A `numpy.array` containing the expected value for each player.
  """
  if state.is_terminal():
    return np.array(state.returns())
  else:
    return sum(prob * policy_value(_child(state, action), policies)
               for action, prob in _transitions(state, policies)
               if prob > PROBABILITY_THRESHOLD)

# Copyright 2019 DeepMind Technologies Ltd. All rights reserved.
#
# 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.

"""Computes the value of a given policy."""

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

from typing import Iterable, List, Union

import numpy as np

from open_spiel.python import policy

PROBABILITY_THRESHOLD = 0


def _child(state, action):
  """Returns a child state, handling the simultaneous node case."""
  if isinstance(action, Iterable):
    child = state.clone()
    child.apply_actions(action)
    return child
  else:
    return state.child(action)


def _transitions(state, policies):
  """Returns iterator over (action, prob) from the given state."""
  if state.is_chance_node():
    return state.chance_outcomes()
  elif state.is_simultaneous_node():
    return policy.joint_action_probabilities(state, policies)
  else:
    player = state.current_player()
    return policies[player].action_probabilities(state).items()


def policy_value(state, policies: Union[List[policy.Policy], policy.Policy],
                 probability_threshold: float = PROBABILITY_THRESHOLD):
  """Returns the expected values for the state for players following `policies`.

  Computes the expected value of the`state` for each player, assuming player `i`
  follows the policy given in `policies[i]`.

  Args:
    state: A `pyspiel.State`.
    policies: A `list` of `policy.Policy` objects, one per player for sequential
      games, one policy for simulatenous games.

  Returns:
    A `numpy.array` containing the expected value for each player.
  """
  if state.is_terminal():
    return np.array(state.returns())
  else:
    return sum(prob * policy_value(_child(state, action), policies)
               for action, prob in _transitions(state, policies)
               if prob > probability_threshold)

Enable using probability_threshold