# Copyright 2021 The Jax Influence 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.

# Copyright 2021 Google LLC
#
# 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 function_factory."""

from unittest import mock

from absl.testing import absltest
from absl.testing import parameterized
import jax
import jax.numpy as jnp
from jax_influence import function_factory
from jax_influence import selection
from jax_influence import test_utils
import numpy as np


def _per_example_loss_fn(params, batch, batch_stats=None):
  """Returns loss on each example."""
  out = batch['x'] @ params['A']
  out = out + batch['y'] @ params['B']
  out = out + params['c']
  # Batch stats is used as a multiplicative constant.
  if batch_stats is not None:
    out = out * batch_stats
  return out**2 / 2


def _total_loss(params, batch, batch_stats=None):
  """Returns loss on a batch."""

  # Batch stats is used as normalization
  if batch_stats is not None:
    den = batch_stats
  else:
    den = 1
  return jnp.sum(_per_example_loss_fn(params, batch)) / den


class FunctionFactoryTest(parameterized.TestCase):

  def setUp(self):
    super().setUp()
    self.params = {
        'A': jnp.array([1.0, 2.0, -3.0]),
        'B': jnp.array([4.0, 6.0, -3.0, -2.0, -1.0]),
        'c': jnp.array([1.3])
    }
    self.tangents = {
        'A': jnp.array([-0.2, 2.4, 3.22]),
        'B': jnp.array([2.0, 3.0, 1.2, -2.5, -1.44]),
        'c': jnp.array([4.22])
    }

  def get_random_batch(self, seed):
    seed = jax.random.PRNGKey(seed)
    seed1, seed2 = jax.random.split(seed)
    x = jax.random.uniform(seed1, (3,))
    y = jax.random.uniform(seed2, (5,))
    return {'x': x, 'y': y}

  def concat_batches(self, batch0, batch1):
    """Concatenates two batches."""
    return jax.tree_multimap(lambda x, y: jnp.concatenate([x[None], y[None]]),
                             batch0, batch1)

  def get_batch(self):
    """Gets a deterministic batch for testing HVP and JVP."""
    batch0 = {
        'x': jnp.array([1.2, .3, .4]),
        'y': jnp.array([.4, -.5, -2.0, -3.0, -4.2])
    }
    batch1 = {
        'x': jnp.array([3.0, 2.3, 4.4]),
        'y': jnp.array([2.4, -1.5, 2.55, 3.44, -1.0])
    }
    return self.concat_batches(batch0, batch1)

  @parameterized.named_parameters(
      dict(testcase_name='seed_0', seed=0),
      dict(testcase_name='seed_17', seed=17),
  )
  def test_restrict_subset_params(self, seed):
    batch = self.get_random_batch(seed)
    fval = _per_example_loss_fn(self.params, batch)
    fun_rest = function_factory.restrict_subset_params(
        _per_example_loss_fn, self.params, select_fn=lambda x: 'B' in x)
    fvalrest = fun_rest({'B': self.params['B']}, batch)
    test_utils.check_close(fval, fvalrest)

  @parameterized.named_parameters(
      dict(
          testcase_name='no_sel_no_batch_stats',
          select_fn=None,
          batch_stats=None,
          expected={
              'A':
                  np.array([78.24721, 49.799202, 92.40961]),
              'B':
                  np.array(
                      [53.393604, -37.480003, 16.873405, 17.799526, -88.53919]),
              'c':
                  np.array([35.944],)
          }),
      dict(