# coding=utf-8
# Copyright 2021 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 simplifying network computation graphs."""
from absl.testing import absltest
from absl.testing import parameterized

import jax
from jax import lax
from jax import numpy as jnp
from jax_verify.src import synthetic_primitives


class SyntheticPrimitiveDetectorTest(parameterized.TestCase):

  def _check_correct_impl(self, graph, simplifier, var_is_bound, *inps):
    simplified_graph = synthetic_primitives.simplify_graph(simplifier, graph,
                                                           var_is_bound)

    graph_outs = jax.core.eval_jaxpr(graph, [], *inps)
    simple_graph_outs = jax.core.eval_jaxpr(simplified_graph, [], *inps)

    for graph_out, simple_graph_out in zip(graph_outs, simple_graph_outs):
      self.assertAlmostEqual(jnp.abs(graph_out - simple_graph_out).max(),
                             0., delta=1e-6)

  def _find_eqn_in_simplified_graph(self, graph, simplifier, var_is_bound,
                                    primitive):
    # Check if the primitive is present. This imitates the recursive
    # parsing done in bound_propagation, because depending on the platform,
    # the primitive might be wrapped in a `custom_jvp_call_jaxpr_p`
    # The loop is necessarily terminating because we always remove one level of
    # nesting in the graph, so we will necessarily reach a level with no
    # subgraph.
    simplified_graph = synthetic_primitives.simplify_graph(simplifier, graph,
                                                           var_is_bound)
    for eqn in simplified_graph.eqns:
      sub_graph = synthetic_primitives.jax_primitive_subgraph(eqn)

      if sub_graph is not None:
        subgraph_var_is_bound = {
            sub_invar: var_is_bound[eqn_invar] for sub_invar, eqn_invar
            in zip(sub_graph.invars, eqn.invars)}
        match = self._find_eqn_in_simplified_graph(
            sub_graph, simplifier, subgraph_var_is_bound, primitive)
        if match:
          return match
      elif eqn.primitive == primitive:
        return eqn
    return None


class ActivationDetectorTest(SyntheticPrimitiveDetectorTest):

  @parameterized.named_parameters(('jit', True), ('nojit', False))
  def test_softplus_detected(self, use_jit):

    def softplus_model(inp):
      return jax.nn.softplus(inp)
    inp = jnp.array([[-2., 3.]])

    if use_jit:
      softplus_model = jax.jit(softplus_model)

    jaxpr_maker = jax.make_jaxpr(softplus_model)
    parsed = jaxpr_maker(inp)
    var_is_bound = {parsed.jaxpr.invars[0]: True}

    found_softplus = self._find_eqn_in_simplified_graph(
        parsed.jaxpr,
        synthetic_primitives.activation_simplifier,
        var_is_bound,
        synthetic_primitives.softplus_p)
    self.assertIsNotNone(found_softplus)

    self._check_correct_impl(
        parsed.jaxpr, synthetic_primitives.activation_simplifier,
        var_is_bound, inp)

  @parameterized.named_parameters(('jit', True), ('nojit', False))
  def test_softmax_detected(self, use_jit):

    def softmax_model(x):
      return jax.nn.softmax(x)

    if use_jit:
      softmax_model = jax.jit(softmax_model)

    inp = jnp.array([[-2., 3.]])

    jaxpr_maker = jax.make_jaxpr(softmax_model)
    parsed = jaxpr_maker(inp)
    var_is_bound = {parsed.jaxpr.invars[0]: True}

    match = self._find_eqn_in_simplified_graph(
        parsed.jaxpr,
        synthetic_primitives.activation_simplifier,
        var_is_bound,
        synthetic_primitives.softmax_p)
    self.assertIsNotNone(match)
    self.assertEqual(match.params['axis'], 1)  # pytype: disable=attribute-error

    self._check_correct_impl(
        parsed.jaxpr, synthetic_primitives.activation_simplifier,
        var_is_bound, inp)

  @parameterized.named_parameters(('jit', True), ('nojit', False))
  def test_softmax_expended(self, use_jit):

    def softmax_model(x):
      return jax.nn.softmax(x)

    if use_jit:
      softmax_model = jax.jit(softmax_model)

    inp = jnp.array([[-2., 3.],
                     [3., 3.1],
                     [-2., -2.],
                     [3., 3.]])

    jaxpr_maker = jax.make_jaxpr(softmax_model)
    parsed = jaxpr_maker(inp)
    var_is_bound = {parsed.jaxpr.invars[0]: True}

    expand_softmax_simplifier = synthetic_primitives.simplifier_composition(
        synthetic_primitives.activation_simplifier,
        synthetic_primitives.expand_softmax_simplifier)

    # Check that all of the components of the softmax that we would expect to
    # find are present.
    softmax_primitives = (
        lax.exp_p,
        lax.reduce_sum_p,
        lax.broadcast_in_dim_p,
        synthetic_primitives.posreciprocal_p,
        lax.mul_p)
    for prim_to_match in softmax_primitives:
      match = self._find_eqn_in_simplified_graph(
          parsed.jaxpr,
          expand_softmax_simplifier,
          var_is_bound,
          prim_to_match)
      self.assertIsNotNone(match)

    self._check_correct_impl(
        parsed.jaxpr, expand_softmax_simplifier,
        var_is_bound, inp)

  @parameterized.named_parameters(('jit', True), ('nojit', False))
  def test_relu_detected(self, use_jit):

    def relu_model(x):
      return jax.nn.relu(x)

    if use_jit:
      relu_model = jax.jit(relu_model)

    inp = jnp.array([[-2., 3.]])
    jaxpr_maker = jax.make_jaxpr(relu_model)
    parsed = jaxpr_maker(inp)
    var_is_bound = {parsed.jaxpr.invars[0]: True}

    match = self._find_eqn_in_simplified_graph(
        parsed.jaxpr,
        synthetic_primitives.activation_simplifier,
        var_is_bound,
        synthetic_primitives.relu_p)

    self.assertIsNotNone(match)

    self._check_correct_impl(
        parsed.jaxpr, synthetic_primitives.activation_simplifier,
        var_is_bound, inp)

  @parameterized.named_parameters(('jit', True), ('nojit', False))
  def test_max_with_zero_detected_as_relu(self, use_jit):

    def relu_model(x):
      return jnp.maximum(x, 0.)

    if use_jit:
      relu_model = jax.jit(relu_model)

    inp = jnp.array([[-2., 3.]])
    jaxpr_maker = jax.make_jaxpr(relu_model)
    parsed = jaxpr_maker(inp)
    var_is_bound = {parsed.jaxpr.invars[0]: True}

    match = self._find_eqn_in_simplified_graph(
        parsed.jaxpr,
        synthetic_primitives.activation_simplifier,
        var_is_bound,
        synthetic_primitives.relu_p)

    self.assertIsNotNone(match)

    self._check_correct_impl(
        parsed.jaxpr, synthetic_primitives.activation_simplifier,
        var_is_bound, inp)

  @parameterized.named_parameters(('jit', True), ('nojit', False))
  def test_max_with_one_not_mistaken_for_relu(self, use_jit):

    def notrelu_model(x):
      return jnp.maximum(x, 1.)

    if use_jit:
      notrelu_model = jax.jit(notrelu_model)

    inp = jnp.array([[-2., 3.]])
    jaxpr_maker = jax.make_jaxpr(notrelu_model)
    parsed = jaxpr_maker(inp)
    var_is_bound = {parsed.jaxpr.invars[0]: True}

    match = self._find_eqn_in_simplified_graph(
        parsed.jaxpr,
        synthetic_primitives.activation_simplifier,
        var_is_bound,
        synthetic_primitives.relu_p)

    self.assertIsNone(match)

    self._check_correct_impl(
        parsed.jaxpr, synthetic_primitives.activation_simplifier,
        var_is_bound, inp)

  @parameterized.named_parameters(('jit', True), ('nojit', False))
  def test_leaky_relu_detected(self, use_jit):

    def leaky_relu_model(x):
      return jax.nn.leaky_relu(x)

    if use_jit:
      leaky_relu_model = jax.jit(leaky_relu_model)

    inp = jnp.array([[-2., 3.]])
    jaxpr_maker = jax.make_jaxpr(leaky_relu_model)
    parsed = jaxpr_maker(inp)
    var_is_bound = {parsed.jaxpr.invars[0]: True}

    match = self._find_eqn_in_simplified_graph(
        parsed.jaxpr,
        synthetic_primitives.activation_simplifier,
        var_is_bound,
        synthetic_primitives.leaky_relu_p)

    self.assertIsNotNone(match)

    self._check_correct_impl(
        parsed.jaxpr, synthetic_primitives.activation_simplifier,
        var_is_bound, inp)

  @parameterized.named_parameters(('jit', True), ('nojit', False))
  def test_sigmoid_detected(self, use_jit):

    def sigmoid_model(x):
      return jax.nn.sigmoid(x)

    if use_jit:
      sigmoid_model = jax.jit(sigmoid_model)

    inp = jnp.array([[-2., 3.]])
    jaxpr_maker = jax.make_jaxpr(sigmoid_model)
    parsed = jaxpr_maker(inp)
    var_is_bound = {parsed.jaxpr.invars[0]: True}

    match = self._find_eqn_in_simplified_graph(
        parsed.jaxpr,
        synthetic_primitives.activation_simplifier,
        var_is_bound,
        synthetic_primitives.sigmoid_p)

    self.assertIsNotNone(match)

    self._check_correct_impl(
        parsed.jaxpr, synthetic_primitives.activation_simplifier,
        var_is_bound, inp)

  @parameterized.named_parameters(('jit', True), ('nojit', False))
  def test_linear_detected(self, use_jit):

    inp = jnp.array([[-1., 1.]])
    key = jax.random.PRNGKey(0)
    key_w, key_b = jax.random.split(key, 2)
    w1 = jax.random.uniform(key_w, shape=(2, 5))
    b1 = jax.random.uniform(key_b, shape=(5,))

    def linear_model(inp, w1, b1):
      """Linear function involving several different linear operators."""
      y = inp @ w1 + b1
      centered_y = y - y.mean()
      return centered_y.sum()

    if use_jit:
      linear_model = jax.jit(linear_model)

    jaxpr_maker = jax.make_jaxpr(linear_model)
    parsed = jaxpr_maker(inp, w1, b1)
    var_is_bound = {invar: is_bound for invar, is_bound
                    in zip(parsed.jaxpr.invars, [True, False, False])}

    match = self._find_eqn_in_simplified_graph(
        parsed.jaxpr,
        synthetic_primitives.group_linear_sequence,
        var_is_bound,
        synthetic_primitives.linear_p
    )

    # Check that all the components that we expect are there.
    linear_subgraph = match.params['jax_verify_subgraph']
    subgraph_primitives = [eqn.primitive for eqn in linear_subgraph.eqns]

    self.assertIn(lax.dot_general_p, subgraph_primitives)
    self.assertIn(lax.add_p, subgraph_primitives)
    # There is two reduce_sum, one for the final sum, one for the mean.
    self.assertEqual(2, sum(prim == lax.reduce_sum_p