# Copyright 2020 DeepMind Technologies Limited. 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.
# ==============================================================================
"""Tests for `asserts.py`."""

import functools
import re

from absl.testing import absltest
from absl.testing import parameterized
from chex._src import asserts
from chex._src import asserts_internal
from chex._src import variants
import jax
import jax.numpy as jnp
import numpy as np

_get_err_regex = asserts_internal.get_err_regex
_num_devices_available = asserts_internal.num_devices_available


def as_arrays(arrays):
  return [np.asarray(a) for a in arrays]


def array_from_shape(*shape):
  return np.ones(shape=shape)


def emplace(arrays):
  return arrays


class AssertsSwitchTest(parameterized.TestCase):
  """Tests for enable/disable_asserts."""

  def test_enable_disable_asserts(self):
    with self.assertRaisesRegex(AssertionError, _get_err_regex('scalar')):
      asserts.assert_scalar('test')

    asserts.disable_asserts()
    asserts.assert_scalar('test')

    asserts.enable_asserts()
    with self.assertRaisesRegex(AssertionError, _get_err_regex('scalar')):
      asserts.assert_scalar('test')

    asserts.disable_asserts()
    asserts.assert_is_divisible(13, 5)

    # To avoid side effects.
    asserts.enable_asserts()


class AssertMaxTracesTest(variants.TestCase):

  def setUp(self):
    super().setUp()
    asserts.clear_trace_counter()

  def _init(self, fn_, init_type, max_traces, kwargs, static_arg):
    """Initializes common test cases."""
    variant_kwargs = {}
    if static_arg:
      variant_kwargs['static_argnums'] = 1

    if kwargs:
      args, kwargs = [], {'n': max_traces}
    else:
      args, kwargs = [max_traces], {}

    if init_type == 't1':

      @asserts.assert_max_traces(*args, **kwargs)
      def fn(x, y):
        if static_arg:
          self.assertNotIsInstance(y, jax.core.Tracer)
        return fn_(x, y)

      fn_jitted = self.variant(fn, **variant_kwargs)
    elif init_type == 't2':

      def fn(x, y):
        if static_arg:
          self.assertNotIsInstance(y, jax.core.Tracer)
        return fn_(x, y)

      fn = asserts.assert_max_traces(fn, *args, **kwargs)
      fn_jitted = self.variant(fn, **variant_kwargs)
    elif init_type == 't3':

      def fn(x, y):
        if static_arg:
          self.assertNotIsInstance(y, jax.core.Tracer)
        return fn_(x, y)

      @self.variant(**variant_kwargs)
      @asserts.assert_max_traces(*args, **kwargs)
      def fn_jitted(x, y):
        self.assertIsInstance(x, jax.core.Tracer)
        return fn_(x, y)
    else:
      raise ValueError(f'Unknown type {init_type}.')

    return fn, fn_jitted

  @variants.variants(with_jit=True, with_pmap=True)
  @parameterized.named_parameters(
      variants.params_product((
          ('type1', 't1'),
          ('type2', 't2'),
          ('type3', 't3'),
      ), (
          ('args', False),
          ('kwargs', True),
      ), (
          ('no_static_arg', False),
          ('with_static_arg', True),
      ), (
          ('max_traces_0', 0),
          ('max_traces_1', 1),
          ('max_traces_2', 2),
          ('max_traces_10', 10),
      ),
                              named=True))
  def test_assert(self, init_type, kwargs, static_arg, max_traces):
    fn_ = lambda x, y: x + y
    fn, fn_jitted = self._init(fn_, init_type, max_traces, kwargs, static_arg)

    # Original function.
    for _ in range(max_traces + 3):
      self.assertEqual(fn(1, 2), 3)

    # Every call results in re-tracing because arguments' shapes are different.
    for i in range(max_traces):
      for k in range(5):
        arg = jnp.zeros(i + 1) + k
        np.testing.assert_array_equal(fn_jitted(arg, 2), arg + 2)

    # Original function.
    for _ in range(max_traces + 3):
      self.assertEqual(fn(1, 2), 3)
      self.assertEqual(fn([1], [2]), [1, 2])
      self.assertEqual(fn('a', 'b'), 'ab')

    # (max_traces + 1)-th re-tracing.
    with self.assertRaisesRegex(AssertionError,
                                _get_err_regex('fn.* is traced > .* times!')):
      arg = jnp.zeros(max_traces + 1)
      fn_jitted(arg, 2)

  def test_incorrect_ordering(self):
    # pylint:disable=g-error-prone-assert-raises,unused-variable
    with self.assertRaisesRegex(ValueError, 'change wrappers ordering'):

      @asserts.assert_max_traces(1)
      @jax.jit
      def fn(_):
        pass

    def dummy_wrapper(fn):

      @functools.wraps(fn)
      def fn_wrapped():
        return fn()

      return fn_wrapped

    with self.assertRaisesRegex(ValueError, 'change wrappers ordering'):

      @asserts.assert_max_traces(1)
      @dummy_wrapper
      @jax.jit
      def fn_2():
        pass

    # pylint:enable=g-error-prone-assert-raises,unused-variable

  def test_redefined_traced_function(self):

    def outer_fn(x):

      @jax.jit
      @asserts.assert_max_traces(3)
      def inner_fn(y):
        return y.sum()

      return inner_fn(2 * x)

    self.assertEqual(outer_fn(1), 2)
    self.assertEqual(outer_fn(2), 4)
    self.assertEqual(outer_fn(3), 6)

    # Fails since the traced inner function is redefined at each call.
    with self.assertRaisesRegex(AssertionError,
                                _get_err_regex('fn.* is traced > .* times!')):
      outer_fn(4)

    asserts.clear_trace_counter()
    for i in range(10):
      if i > 2:
        with self.assertRaisesRegex(
            AssertionError, _get_err_regex('fn.* is traced > .* times!')):
          outer_fn(1)
      else:
        outer_fn(1)

  def test_nested_functions(self):

    @jax.jit
    def jitted_outer_fn(x):

      @jax.jit
      @asserts.assert_max_traces(1)
      def inner_fn(y):
        return y.sum()

      return inner_fn(2 * x)

    # Inner assert_max_traces have no effect since the outer_fn is traced once.
    for i in range(10):
      self.assertEqual(jitted_outer_fn(i), 2 * i)


class ScalarAssertTest(parameterized.TestCase):

  def test_scalar(self):
    asserts.assert_scalar(1)
    asserts.assert_scalar(1.)
    with self.assertRaisesRegex(AssertionError,
                                _get_err_regex('must be a scalar')):
      asserts.assert_scalar(np.array(1.))  # pytype: disable=wrong-arg-types

  def test_scalar_positive(self):
    asserts.assert_scalar_positive(0.5)
    with self.assertRaisesRegex(AssertionError,
                                _get_err_regex('must be positive')):
      asserts.assert_scalar_positive(-0.5)

  def test_scalar_non_negative(self):
    asserts.assert_scalar_non_negative(0.5)
    asserts.assert_scalar_non_negative(0.)
    with self.assertRaisesRegex(AssertionError,
                                _get_err_regex('must be non-negative')):
      asserts.assert_scalar_non_negative(-0.5)

  def test_scalar_negative(self):
    asserts.assert_scalar_negative(-0.5)
    with self.assertRaisesRegex(AssertionError,
                                _get_err_regex('argument must be negative')):
      asserts.assert_scalar_negative(0.5)

  def test_scalar_in(self):
    asserts.assert_scalar_in(0.5, 0, 1)
    with self.assertRaisesRegex(AssertionError,
                                _get_err_regex('argument must be in')):
      asserts.assert_scalar_in(-0.5, 0, 1)
    with self.assertRaisesRegex(AssertionError,
                                _get_err_regex('argument must be in')):
      asserts.assert_scalar_in(1.5, 0, 1)

  def test_scalar_in_excluded(self):
    asserts.assert_scalar_in(0.5, 0, 1, included=False)
    with self.assertRaisesRegex(AssertionError,
                                _get_err_regex('argument must be in')):
      asserts.assert_scalar_in(0, 0, 1, included=False)
    with self.assertRaisesRegex(AssertionError,
                                _get_err_regex('argument must be in')):