# Copyright 2019 DeepMind Technologies Limited
#
# 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.

"""Tests for open_spiel.python.algorithms.adidas_utils.helpers.symmetric.exploitability.

Computing the exploitability of a tsallis-entropy regularized game is more
involved, so we include a derivation here of an example test case using the
prisoner's dilemma (see pd np.array below). Note that the tsallis-entropy
setting assumes non-negative payoffs so we add 3 to the array. We assume p=1
for the tsallis entropy in this example.

pd    dist   grad     br       payoff(br)     payoff(dist)
[2 0] [.5] = [1] --> [1/3] --> 5/3        --> 3/2
[3 1] [.5]   [2]     [2/3]

s = sum(grad) = 3

tsallis-entr(br) = s / (p + 1) * (1 - br_1^2 - br_2^2)
                 = 3 / 2 * (1 - 1/9 - 4/9) = 2/3

tsallis-entr(dist) = s / (p + 1) * (1 - dist_1^2 - dist_2^2)
                   = 3 / 2 * (1 - 1/4 - 1/4) = 3/4

u(br) - u(dist) = 5/3 + 2/3 - 3/2 - 3/4 = 7 / 3 - 9 / 4
"""

from absl import logging  # pylint:disable=unused-import
from absl.testing import absltest
from absl.testing import parameterized

import numpy as np

from open_spiel.python.algorithms.adidas_utils.helpers.symmetric import exploitability


test_seed = 12345

# prisoner's dilemma test case
# pylint:disable=bad-whitespace
pt_r = np.array([[-1, -3],
                 [0,  -2]])
# pylint:enable=bad-whitespace
pt_r -= pt_r.min()
pt_c = pt_r.T
pd = np.stack((pt_r, pt_c), axis=0)
pd_nash = np.array([0, 1])
pd_non_nash_1 = np.array([1, 0])
pd_non_nash_exp_1 = 1.
pd_non_nash_ate_exp_1 = pd_non_nash_exp_1
pd_non_nash_2 = np.array([0.5, 0.5])
pd_non_nash_exp_2 = 0.5
pd_non_nash_ate_exp_2 = 7. / 3. - 9. / 4.

qre_br = np.exp([1, 2]) / np.exp([1, 2]).sum()
entr_br = -np.sum(qre_br * np.log(qre_br))
entr_non_nash_2 = -np.sum(pd_non_nash_2 * np.log(pd_non_nash_2))
u_br_minus_non_nash = (qre_br - pd_non_nash_2).dot([1, 2])
pd_non_nash_qre_exp_2 = u_br_minus_non_nash + (entr_br - entr_non_nash_2)

# rock-paper-scissors test case
# pylint:disable=bad-whitespace
pt_r = np.array([[0, -1,  1],
                 [1,  0, -1],
                 [-1, 1,  0]])
# pylint:enable=bad-whitespace
pt_r -= pt_r.min()
pt_c = pt_r.T
rps = np.stack((pt_r, pt_c), axis=0)
rps_nash = np.ones(3) / 3.
rps_non_nash_1 = np.array([1, 0, 0])
rps_non_nash_exp_1 = 1.
rps_non_nash_2 = np.array([0, 1, 0])
rps_non_nash_exp_2 = 1.
rps_non_nash_3 = np.array([0, 0, 1])
rps_non_nash_exp_3 = 1.


class ExploitabilityTest(parameterized.TestCase):

  @parameterized.named_parameters(
      ('PD_nash', pd, pd_nash),
      ('RPS_nash', rps, rps_nash),
      )
  def test_unreg_exploitability_of_nash(self, payoff_tensor, nash):
    # assumes symmetric games
    exp = exploitability.unreg_exploitability(nash, payoff_tensor)
    self.assertEqual(exp, 0., 'nash should have zero exploitability')

  @parameterized.named_parameters(
      ('PD_non_nash_1', pd, pd_non_nash_1, pd_non_nash_exp_1),
      ('PD_non_nash_2', pd, pd_non_nash_2, pd_non_nash_exp_2),
      ('RPS_non_nash_1', rps, rps_non_nash_1, rps_non_nash_exp_1),
      ('RPS_non_nash_2', rps, rps_non_nash_2, rps_non_nash_exp_2),
      ('RPS_non_nash_3', rps, rps_non_nash_3, rps_non_nash_exp_3),
      )
  def test_unreg_exploitability_of_non_nash(self, payoff_tensor, dist, exp):
    # assumes symmetric games
    exp_pred = exploitability.unreg_exploitability(dist, payoff_tensor)
    self.assertEqual(exp_pred, exp, 'dist should have the given exploitability')

  @parameterized.named_parameters(
      ('PD_rand', pd, test_seed),
      ('RPS_rand', rps, test_seed),
      )
  def test_unreg_exploitability_of_rand(self, payoff_tensor, seed=None):
    trials = 100
    random = np.random.RandomState(seed)
    num_strategies = payoff_tensor.shape[-1]
    dists = random.rand(trials, num_strategies)
    dists /= np.sum(dists, axis=1, keepdims=True)
    exploitable = []
    for dist in dists:
      exp = exploitability.unreg_exploitability(dist, payoff_tensor)
      exploitable.append(exp > 0.)
    perc = 100 * np.mean(exploitable)
    logging.info('rand strat exploitable rate out of %d is %f', trials, perc)
    self.assertEqual(perc, 100., 'found rand strat that was nash')

  @parameterized.named_parameters(
      ('RPS_nash_p=0', rps, rps_nash, 0.),
      ('RPS_nash_p=0.1', rps, rps_nash, 0.1),
      ('RPS_nash_p=1', rps, rps_nash, 1.),
      )
  def test_ate_exploitability_of_nash(self, payoff_tensor, nash, p):
    # assumes symmetric games
    exp = exploitability.ate_exploitability(nash, payoff_tensor, p)
    self.assertGreaterEqual(0., exp,
                            'uniform nash should have zero exploitability')

  @parameterized.named_parameters(
      ('PD_non_nash_p=0', pd, 0., pd_non_nash_1, pd_non_nash_exp_1),
      ('PD_non_nash_p=1', pd, 1., pd_non_nash_2, pd_non_nash_ate_exp_2),
      )
  def test_ate_exploitability_of_non_nash(self, payoff_tensor, p, dist, exp):
    # assumes symmetric games
    exp_pred = exploitability.ate_exploitability(dist, payoff_tensor, p)
    self.assertAlmostEqual(exp_pred, exp,
                           msg='dist should have the given exploitability')

  @parameterized.named_parameters(
      ('RPS_rand_p=0', rps, 0., test_seed),
      ('RPS_rand_p=0.1', rps, 0.1, test_seed),
      ('RPS_rand_p=1', rps, 1., test_seed),
      )
  def test_ate_exploitability_of_rand(self, payoff_tensor, p, seed=None):
    trials = 100
    random = np.random.RandomState(seed)
    num_strategies = payoff_tensor.shape[-1]
    dists = random.rand(trials, num_strategies)
    dists /= np.sum(dists, axis=1, keepdims=True)
    exploitable = []
    for dist in dists:
      exp = exploitability.ate_exploitability(dist, payoff_tensor, p)
      exploitable.append(exp > 0.)
    perc = 100 * np.mean(exploitable)
    logging.info('rand strat exploitable rate out of %d is %f', trials, perc)
    self.assertEqual(perc, 100., 'found rand strat that was nash')

  @parameterized.named_parameters(
      ('RPS_nash_tau=0', rps, rps_nash, 0.),
      ('RPS_nash_tau=0.1', rps, rps_nash, 0.1),
      ('RPS_nash_tau=1', rps, rps_nash, 1.),
      )
  def test_qre_exploitability_of_nash(self, payoff_tensor, nash, temperature):
    # assumes symmetric games
    exp = exploitability.qre_exploitability(nash, payoff_tensor, temperature)
    self.assertGreaterEqual(1e-10, exp,
                            'uniform nash should have zero exploitability')

  @parameterized.named_parameters(
      ('PD_non_nash_tau=0', pd, 0., pd_non_nash_1, pd_non_nash_exp_1),
      ('PD_non_nash_tau=1', pd, 1., pd_non_nash_2, pd_non_nash_qre_exp_2),
      )
  def test_qre_exploitability_of_non_nash(self, payoff_tensor, temperature,
                                          dist, exp):
    # assumes symmetric games
    exp_pred = exploitability.qre_exploitability(dist, payoff_tensor,
                                                 temperature)
    self.assertAlmostEqual(exp_pred, exp,
                           msg='dist should have the given exploitability')

  @parameterized.named_parameters(
      ('RPS_rand_tau=0', rps, 0., test_seed),
      ('RPS_rand_tau=0.1', rps, 0.1, test_seed),
      ('RPS_rand_tau=1', rps, 1., test_seed),
      )
  def test_qre_exploitability_of_rand(self, payoff_tensor, temperature,
                                      seed=None):
    trials = 100
    random = np.random.RandomState(seed)
    num_strategies = payoff_tensor.shape[-1]
    dists = random.rand(trials, num_strategies)
    dists /= np.sum(dists, axis=1, keepdims=True)
    exploitable = []
    for dist in dists:
      exp = exploitability.qre_exploitability(dist, payoff_tensor, temperature)
      exploitable.append(exp > 0.)
    perc = 100 * np.mean(exploitable)
    logging.info('rand strat exploitable rate out of %d is %f', trials, perc)
    self.assertEqual(perc, 100., 'found rand strat that was nash')


if __name__ == '__main__':
  absltest.main()
