# Copyright 2023 The Flax 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.

from flax import errors
from flax.core import Scope, scope, freeze, lazy_init, init, apply, nn
from flax.core.scope import LazyRng
from flax.configurations import temp_flip_flag

import jax
from jax import random
from jax import numpy as jnp

import numpy as np


from absl.testing import absltest


class ScopeTest(absltest.TestCase):

  def test_rng(self):
    def f(scope):
      self.assertTrue(scope.has_rng('params'))
      self.assertFalse(scope.has_rng('dropout'))
      rng = scope.make_rng('params')
      self.assertTrue(
          np.all(rng == LazyRng.create(random.key(0), 1).as_jax_rng())
      )

    init(f)(random.key(0))

  def test_in_filter(self):
    filter_true = lambda x, y: self.assertTrue(scope.in_filter(x, y))
    filter_false = lambda x, y: self.assertFalse(scope.in_filter(x, y))

    filter_true(True, 'any_string1')
    filter_false(False, 'any_string2')
    filter_true('exact_match', 'exact_match')
    filter_false('no_match1', 'no_match2')
    filter_true(['one', 'two'], 'one')
    filter_false(['one', 'two'], 'three')
    filter_false([], 'one')
    filter_false([], None)

  def test_union_filter(self):
    def union_check(a, b, ans):
      self.assertEqual(scope.union_filters(a, b), ans)
      self.assertEqual(scope.union_filters(b, a), ans)

    union_check(['a', 'b'], ['b', 'c'], set(['a', 'b', 'c']))
    union_check(True, False, True)
    union_check(False, False, set())
    union_check(True, True, True)
    union_check(
        scope.DenyList(['a', 'b']),
        scope.DenyList(['b', 'c']),
        scope.DenyList(set(['b'])),
    )
    union_check(
        scope.DenyList(['a', 'b']), ['b', 'c'], scope.DenyList(set(['a']))
    )

  def test_intersect_filter(self):
    def intersect_check(a, b, ans):
      self.assertEqual(scope.intersect_filters(a, b), ans)
      self.assertEqual(scope.intersect_filters(b, a), ans)

    intersect_check(['a', 'b'], ['b', 'c'], set(['b']))
    intersect_check(True, False, False)
    intersect_check(False, False, set())
    intersect_check(True, True, True)
    intersect_check(
        scope.DenyList(['a', 'b']),
        scope.DenyList(['b', 'c']),
        scope.DenyList(set(['a', 'b', 'c'])),
    )
    intersect_check(scope.DenyList(['a', 'b']), ['b', 'c'], set(['c']))

  def test_subtract_filter(self):
    def subtract_check(a, b, ans):
      self.assertEqual(scope.subtract_filters(a, b), ans)

    subtract_check(['a', 'b'], ['b', 'c'], set(['a']))
    subtract_check(True, False, scope.DenyList(False))
    subtract_check(False, False, set())
    subtract_check(True, True, False)
    subtract_check(True, 'a', scope.DenyList('a'))
    subtract_check(
        scope.DenyList(['a', 'b']), scope.DenyList(['b', 'c']), set(['c'])
    )
    subtract_check(
        scope.DenyList(['a', 'b']),
        ['b', 'c'],
        scope.DenyList(set(['a', 'b', 'c'])),
    )

  def test_group_collections(self):
    params = {'dense1': {'x': [10, 20]}}
    batch_stats = {'dense1': {'ema': 5}}
    xs = {'params': params, 'batch_stats': batch_stats}

    # Retrieve all keys only once.
    group = scope.group_collections(xs, ['params', 'params'])
    self.assertEqual(group, ({'params': params}, {}))

    # Ignore non-existing keys.
    self.assertEqual(scope.group_collections(xs, ['vars']), ({},))

    # False gets nothing and True retrieves all keys once.
    self.assertEqual(
        scope.group_collections(xs, [False, True, True]), ({}, xs, {})
    )

  def test_inconsistent_param_shapes(self):
    def f(scope):
      scope.param('test', nn.initializers.ones_init(), (4,))

    msg = (
        r'Initializer expected to generate shape \(2,\) but got shape \(4,\)'
        r' instead for parameter "test" in "/"'
    )
    with self.assertRaisesRegex(errors.ScopeParamShapeError, msg):
      apply(f)(freeze({'params': {'test': np.ones((2,))}}))

  def test_apply_variables_bad_pytree(self):
    def f(scope):
      scope.param('kernel', nn.initializers.ones_init(), (4,))

    params = freeze({
        'params': {
            'kernel': np.ones((4,)),
        },
    })
    apply(f)(params)  # Valid.
    msg = 'but got a dict with an extra params layer'
    with self.assertRaisesRegex(
        errors.ApplyScopeInvalidVariablesStructureError, msg
    ):
      apply(f)({'params': params})

  def test_mutate_undefined_collection(self):
    def f(scope):
      scope.put_variable('state', 'test', 123)

    msg = (
        r'Cannot update variable "test" in "/" because collection "state" is'
        r' immutable.'
    )
    with self.assertRaisesRegex(errors.ModifyScopeVariableError, msg):
      init(f, mutable='params')(random.key(0))

  def test_undefined_param(self):
    def f(scope):
      nn.dense(scope.push('dense'), np.ones((1, 2)), 2)

    msg = r'Could not find parameter named "kernel" in scope "/dense".'
    with self.assertRaisesRegex(errors.ScopeParamNotFoundError, msg):
      apply(f)({'params': {'abc': 1}})

  def test_variable_is_mutable(self):
    def f(scope, should_be_mutable):
      test = scope.variable('state', 'test', lambda: 1)
      self.assertEqual(test.is_mutable(), should_be_mutable)

    _, variables = apply(f, mutable='state')({}, True)
    apply(f, mutable=False)(variables, False)

  def test_rngs_check_w_frozen_dict(self):
    def f(scope, x):
      return x

    _ = apply(f)({}, np.array([0.0]), rngs=freeze({'a': random.key(0)}))

  def test_rng_check_w_old_and_new_keys(self):
    # random.key always returns a new-style typed PRNG key.
    key = random.key(0)
    self.assertTrue(scope._is_valid_rng(key))
    self.assertFalse(scope._is_valid_rng(random.split(key)))

    # random.PRNGKey returns an old-style uint32 key by default.
    old_key = random.PRNGKey(0)
    self.assertTrue(scope._is_valid_rng(old_key))
    self.assertFalse(scope._is_valid_rng(random.split(old_key)))

    # Also explicitly test raw key data, because the jax_enable_custom_prng
    # flag can make PRNGKey return new-style keys.
    raw_key = random.key_data(key)
    self.assertTrue(scope._is_valid_rng(raw_key))
    self.assertFalse(scope._is_valid_rng(random.split(raw_key)))

  def test_jax_leak_detector(self):
    with jax.check_tracer_leaks(True):

      def f(scope):
        def g(scope):
          pass

        scope.child(g)()

      jax.jit(init(f))(random.key(0))

  def test_rng_counter_reuse(self):
    root = Scope({}, {'dropout': random.key(0)})

    def f(scope):
      return scope.make_rng('dropout')

    a = root.child(f)()
    root = root.rewound()
    b = root.child(f)()
    self.assertFalse(jnp.allclose(a, b))

  def test_empty_col_error(self):
    root = Scope({})
    with self.assertRaises(errors.ScopeCollectionNotFound):
      root.param('test', nn.initializers.zeros_init(), ())
    root = Scope({'params': {}})
    with self.assertRaises(errors.ScopeCollectionNotFound):
      root.param('test', nn.initializers.zeros_init(), ())

    root = Scope({'params': {'abc': 1}})
    with self.assertRaises(errors.ScopeCollectionNotFound):
      root.variable('state', 'test', jnp.zeros, ())
    root = Scope({'state': {}})
    with self.assertRaises(errors.ScopeCollectionNotFound):
      root.variable('state', 'test', jnp.zeros, ())

  def test_variable_no_init(self):
    root = Scope({}, mutable='state')
    with self.assertRaises(errors.ScopeCollectionNotFound):
      root.variable('state', 'test')
    root = Scope({'state': {'abc': 1}}, mutable='state')
    abc = root.variable('state', 'abc')
    self.assertEqual(abc.value, 1)
    with self.assertRaises(errors.ScopeVariableNotFoundError):
      root.variable('state', 'test')

  def test_variable_alias(self):
    scope = Scope({}, mutable='state')
    subscope = scope.push(name='a')
    subscope.put_variable('state', 'x', 0.0)
    scope.put_variable('state', 'a', {'x': jnp.array(1.0, jnp.float32)})
    self.assertEqual(
        scope.variables()['state']['a']['x'], subscope.variables()['state']['x']
    )

  def test_lazy_init(self):
    def f(scope, x):
      k = scope.param(
          'kernel', nn.initializers.lecun_normal(), (x.shape[-1], x.shape[-1])
      )
      return x @ k

    init_fn = lazy_init(f)
    # provide a massive input message which would OOM if any compute ops were actually executed
    variables = init_fn(
        random.key(0),
        jax.ShapeDtypeStruct((1024 * 1024 * 1024, 128), jnp.float32),
    )
    self.assertEqual(variables['params']['kernel'].shape, (128, 128))

  def test_lazy_init_fails_on_data_dependence(self):
    def f(scope, x):
      # kernel is initialized with x so params are now dependent on the input
      k = scope.param('kernel', lambda _: x)
      return x * k

    init_fn = lazy_init(f)
    with self.assertRaises(errors.LazyInitError):
      init_fn(random.key(0), jax.ShapeDtypeStruct((8, 4), jnp.float32))

  @temp_flip_flag('fix_rng_separator', True)
  def test_fold_in_static_seperator(self):
    x = LazyRng(random.key(0), ('ab', 'c'))
    y = LazyRng(random.key(0), ('a', 'bc'))
    self.assertFalse(np.all(x.as_jax_rng() == y.as_jax_rng()))


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