# Copyright 2019 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 networks."""

# pylint: disable=g-bad-import-order

import chex
import haiku as hk
import jax
from jax.config import config
import jax.numpy as jnp
import numpy as np
import tree

from dqn_zoo import networks
from absl.testing import absltest


def _sample_input(input_shape):
  return jnp.zeros((1,) + input_shape, dtype=jnp.float32)


class SimpleLayersTest(absltest.TestCase):

  def test_linear(self):
    layer = hk.transform(networks.linear(4))
    params = layer.init(jax.random.PRNGKey(1), _sample_input((3,)))
    self.assertCountEqual(['linear'], params)
    lin_params = params['linear']
    self.assertCountEqual(['w', 'b'], lin_params)
    chex.assert_shape(lin_params['w'], (3, 4))
    chex.assert_shape(lin_params['b'], (4,))

  def test_conv(self):
    layer = hk.transform(networks.conv(4, (3, 3), 2))
    params = layer.init(jax.random.PRNGKey(1), _sample_input((7, 7, 3)))
    self.assertCountEqual(['conv2_d'], params)
    conv_params = params['conv2_d']
    self.assertCountEqual(['w', 'b'], conv_params)
    chex.assert_shape(conv_params['w'], (3, 3, 3, 4))
    chex.assert_shape(conv_params['b'], (4,))


class LinearWithSharedBiasTest(absltest.TestCase):

  def setUp(self):
    super().setUp()
    rng_key = jax.random.PRNGKey(1)
    self.init_rng_key, self.apply_rng_key = jax.random.split(rng_key)
    self.input_shape = (4,)
    self.output_shape = (3,)
    self.weights_shape = (self.input_shape[0], self.output_shape[0])
    network_fn = networks.linear_with_shared_bias(self.output_shape[0])
    self.network = hk.transform(network_fn)

  def test_bias_parameter_shape(self):
    params = self.network.init(
        self.init_rng_key, _sample_input(self.input_shape)
    )
    self.assertLen(tree.flatten(params), 2)

    def check_params(path, param):
      if path[-1] == 'b':
        self.assertNotEqual(self.output_shape, param.shape)
        chex.assert_shape(param, (1,))
      elif path[-1] == 'w':
        chex.assert_shape(param, self.weights_shape)
      else:
        self.fail('Unexpected parameter %s.' % path)

    tree.map_structure_with_path(check_params, params)

  def test_output_shares_bias(self):
    bias = 1.23
    params = self.network.init(
        self.init_rng_key, _sample_input(self.input_shape)
    )

    def replace_params(path, param):
      if path[-1] == 'b':
        return jnp.ones_like(param) * bias
      else:
        return jnp.zeros_like(param)

    params = tree.map_structure_with_path(replace_params, params)
    output = self.network.apply(
        params, self.apply_rng_key, jnp.zeros((1,) + self.input_shape)
    )
    chex.assert_shape(output, (1,) + self.output_shape)
    np.testing.assert_allclose([bias] * self.output_shape[0], list(output[0]))


class NoisyLinearTest(absltest.TestCase):

  def setUp(self):
    super().setUp()
    rng_key = jax.random.PRNGKey(1)
    self.init_rng_key, self.apply_rng_key = jax.random.split(rng_key)
    self.input_shape = (4,)
    self.output_shape = (3,)
    self.network_fn = networks.noisy_linear(self.output_shape[0], 0.1)
    self.network = hk.transform(self.network_fn)
    self.params = self.network.init(
        self.init_rng_key, _sample_input(self.input_shape)
    )
    self.inputs = jnp.zeros((2,) + self.input_shape)

  def test_basic(self):
    self.network.apply(self.params, self.apply_rng_key, self.inputs)

  def test_error_raised_if_rng_is_not_passed_in(self):
    with self.assertRaisesRegex(ValueError, 'must be called with an RNG'):
      self.network.apply(self.params, self.inputs)

  def test_error_raised_if_transformed_without_rng_1(self):
    network = hk.without_apply_rng(hk.transform(self.network_fn))
    with self.assertRaisesRegex(ValueError, 'PRNGKey'):
      network.apply(self.params, self.inputs)

  def test_error_raised_if_transformed_without_rng_2(self):
    network = hk.without_apply_rng(hk.transform(self.network_fn))
    with self.assertRaisesRegex(TypeError, 'positional argument'):
      network.apply(self.params, self.apply_rng_key, self.inputs)

  def test_same_rng_produces_same_outputs(self):
    outputs_1 = self.network.apply(self.params, self.apply_rng_key, self.inputs)
    outputs_2 = self.network.apply(self.params, self.apply_rng_key, self.inputs)
    np.testing.assert_allclose(outputs_1, outputs_2)

  def test_different_rngs_produce_different_outputs(self):
    rng_1, rng_2 = jax.random.split(jax.random.PRNGKey(1))
    outputs_1 = self.network.apply(self.params, rng_1, self.inputs)
    outputs_2 = self.network.apply(self.params, rng_2, self.inputs)
    self.assertFalse(np.allclose(outputs_1, outputs_2))

  def test_number_of_params_with_bias_correct(self):
    net_fn = networks.noisy_linear(self.output_shape[0], 0.1, with_bias=True)
    network = hk.transform(net_fn)
    params = network.init(self.init_rng_key, _sample_input(self.input_shape))
    self.assertCountEqual(['mu', 'sigma'], params)
    self.assertCountEqual(['b', 'w'], params['mu'])
    self.assertCountEqual(['b', 'w'], params['sigma'])

  def test_number_of_params_without_bias_correct(self):
    net_fn = networks.noisy_linear(self.output_shape[0], 0.1, with_bias=False)
    network = hk.transform(net_fn)
    params = network.init(self.init_rng_key, _sample_input(self.input_shape))
    self.assertCountEqual(['mu', 'sigma'], params)
    self.assertCountEqual(['w'], params['mu'])
    self.assertCountEqual(['b', 'w'], params['sigma'])

  def test_sigma_params_are_constant(self):
    self.assertCountEqual(['mu', 'sigma'], self.params)
    sigma_params = self.params['sigma']
    sigma_w_values = np.unique(sigma_params['w'])
    sigma_b_values = np.unique(sigma_params['b'])
    self.assertLen(sigma_w_values, 1)
    self.assertLen(sigma_b_values, 1)
    value = 0.1 / np.sqrt(self.input_shape[0])
    self.assertAlmostEqual(value, sigma_w_values)
    self.assertAlmostEqual(value, sigma_b_values)


if __name__ == '__main__':
  config.update('jax_numpy_rank_promotion', 'raise')
  absltest.main()
