#
# Copyright (c) 2021 The GPflux Contributors.
#
# 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.
#

import abc

import numpy as np
import pytest
import tensorflow as tf
import tensorflow_probability as tfp

from gpflow.kullback_leiblers import gauss_kl

from gpflux.encoders import DirectlyParameterizedNormalDiag
from gpflux.layers import LatentVariableLayer, LayerWithObservations, TrackableLayer

tf.keras.backend.set_floatx("float64")

############
# Utilities
############


def _zero_one_normal_prior(w_dim):
    """ N(0, I) prior """
    return tfp.distributions.MultivariateNormalDiag(loc=np.zeros(w_dim), scale_diag=np.ones(w_dim))


def get_distributions_with_w_dim():
    distributions = []
    for d in [1, 5]:

        mean = np.zeros(d)
        scale_tri_l = np.eye(d)
        mvn = tfp.distributions.MultivariateNormalTriL(mean, scale_tri_l)

        std = np.ones(d)
        mvn_diag = tfp.distributions.MultivariateNormalDiag(mean, std)

        distributions.append((mvn, d))
        distributions.append((mvn_diag, d))
    return distributions


############
# Tests
############


@pytest.mark.parametrize("distribution, w_dim", get_distributions_with_w_dim())
def test_local_kls(distribution, w_dim):
    lv = LatentVariableLayer(encoder=None, prior=distribution)

    # test kl is 0 when posteriors == priors
    posterior = distribution
    assert lv._local_kls(posterior) == 0

    # test kl > 0 when posteriors != priors
    batch_size = 10
    params = distribution.parameters
    posterior_params = {
        k: [v + 0.5 for _ in range(batch_size)]
        for k, v in params.items()
        if isinstance(v, np.ndarray)
    }
    posterior = lv.distribution_class(**posterior_params)
    local_kls = lv._local_kls(posterior)
    assert np.all(local_kls > 0)
    assert local_kls.shape == (batch_size,)


@pytest.mark.parametrize("w_dim", [1, 5])
def test_local_kl_gpflow_consistency(w_dim):
    num_data = 400
    means = np.random.randn(num_data, w_dim)
    encoder = DirectlyParameterizedNormalDiag(num_data, w_dim, means)

    lv = LatentVariableLayer(encoder=encoder, prior=_zero_one_normal_prior(w_dim))
    posteriors = lv._inference_posteriors(
        [np.random.randn(num_data, 3), np.random.randn(num_data, 2)]
    )

    q_mu = posteriors.parameters["loc"]
    q_sqrt = posteriors.parameters["scale_diag"]

    gpflow_local_kls = gauss_kl(q_mu, q_sqrt)
    tfp_local_kls = tf.reduce_sum(lv._local_kls(posteriors))

    np.testing.assert_allclose(tfp_local_kls, gpflow_local_kls, rtol=1e-10)


class ArrayMatcher:
    def __init__(self, expected):
        self.expected = expected

    def __eq__(self, actual):
        return np.allclose(actual, self.expected, equal_nan=True)


@pytest.mark.parametrize("w_dim", [1, 5])
def test_latent_variable_layer_losses(mocker, w_dim):
    num_data, x_dim, y_dim = 43, 3, 1

    prior_shape = (w_dim,)
    posteriors_shape = (num_data, w_dim)

    prior = tfp.distributions.MultivariateNormalDiag(
        loc=np.random.randn(*prior_shape),
        scale_diag=np.random.randn(*prior_shape) ** 2,
    )
    posteriors = tfp.distributions.MultivariateNormalDiag(
        loc=np.random.randn(*posteriors_shape),
        scale_diag=np.random.randn(*posteriors_shape) ** 2,
    )

    encoder = mocker.Mock(return_value=(posteriors.loc, posteriors.scale.diag))

    lv = LatentVariableLayer(encoder=encoder, prior=prior)

    inputs = np.full((num_data, x_dim), np.nan)
    targets = np.full((num_data, y_dim), np.nan)
    observations = [inputs, targets]
    encoder_inputs = np.concatenate(observations, axis=-1)

    _ = lv(inputs)
    encoder.assert_not_called()
    assert lv.losses == [0.0]

    _ = lv(inputs, observations=observations, training=True)

    # assert_called_once_with uses == for comparison which fails on arrays
    encoder.assert_called_once_with(ArrayMatcher(encoder_inputs), training=True)

    expected_loss = [tf.reduce_mean(posteriors.kl_divergence(prior))]
    np.testing.assert_equal(lv.losses, expected_loss)  # also checks shapes match


@pytest.mark.parametrize("w_dim", [1, 5])
@pytest.mark.parametrize("seed2", [None, 42])
def test_latent_variable_layer_samples(mocker, test_data, w_dim, seed2):
    seed = 123

    inputs, targets = test_data
    num_data, x_dim = inputs.shape

    prior_shape = (w_dim,)
    posteriors_shape = (num_data, w_dim)

    prior = tfp.distributions.MultivariateNormalDiag(
        loc=np.random.randn(*prior_shape),
        scale_diag=np.random.randn(*prior_shape) ** 2,
    )
    posteriors = tfp.distributions.MultivariateNormalDiag(
        loc=np.random.randn(*posteriors_shape),
        scale_diag=np.random.randn(*posteriors_shape) ** 2,
    )

    encoder = mocker.Mock(return_value=(posteriors.loc, posteriors.scale.diag))

    lv = LatentVariableLayer(prior=prior, encoder=encoder)

    tf.random.set_seed(seed)
    sample_prior = lv(inputs, seed=seed2)
    tf.random.set_seed(seed)
    prior_expected = np.concatenate([inputs, prior.sample(num_data, seed=seed2)], axis=-1)
    np.testing.assert_array_equal(sample_prior, prior_expected)

    tf.random.set_seed(seed)
    sample_posterior = lv(inputs, observations=[inputs, targets], training=True, seed=seed2)
    tf.random.set_seed(seed)
    posterior_expected = np.concatenate([inputs, posteriors.sample(seed=seed2)], axis=-1)
    np.testing.assert_array_equal(sample_posterior, posterior_expected)


def test_no_tensorflow_metaclass_overwritten():
    """
    LayerWithObservations is a subclass of tf.keras.layers.Layer (via TrackableLayer);
    this test ensures that TrackableLayer does not have a metaclass, and hence by adding
    the ABCMeta to LayerWithObservations we are not accidentally removing some required
    TensorFlow magic metaclass.
    """
    assert LayerWithObservations.__bases__ == (TrackableLayer,)
    assert type(TrackableLayer) is type
    assert type(LayerWithObservations) is abc.ABCMeta

[((29, 0, 29, 38), 'tensorflow.keras.backend.set_floatx', 'tf.keras.backend.set_floatx', ({(29, 28, 29, 37): '"""float64"""'}, {}), "('float64')", True, 'import tensorflow as tf\n'), ((84, 1, 84, 41), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(84, 25, 84, 32): '"""w_dim"""', (84, 34, 84, 40): '[1, 5]'}, {}), "('w_dim', [1, 5])", False, 'import pytest\n'), ((112, 1, 112, 41), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(112, 25, 112, 32): '"""w_dim"""', (112, 34, 112, 40): '[1, 5]'}, {}), "('w_dim', [1, 5])", False, 'import pytest\n'), ((150, 1, 150, 41), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(150, 25, 150, 32): '"""w_dim"""', (150, 34, 150, 40): '[1, 5]'}, {}), "('w_dim', [1, 5])", False, 'import pytest\n'), ((151, 1, 151, 45), 'pytest.mark.parametrize', 'pytest.mark.parametrize', ({(151, 25, 151, 32): '"""seed2"""', (151, 34, 151, 44): '[None, 42]'}, {}), "('seed2', [None, 42])", False, 'import pytest\n'), ((64, 9, 64, 62), 'gpflux.layers.LatentVariableLayer', 'LatentVariableLayer', (), '', False, 'from gpflux.layers import LatentVariableLayer, LayerWithObservations, TrackableLayer\n'), ((80, 11, 80, 32), 'numpy.all', 'np.all', ({(80, 18, 80, 31): '(local_kls > 0)'}, {}), '(local_kls > 0)', True, 'import numpy as np\n'), ((87, 12, 87, 44), 'numpy.random.randn', 'np.random.randn', ({(87, 28, 87, 36): 'num_data', (87, 38, 87, 43): 'w_dim'}, {}), '(num_data, w_dim)', True, 'import numpy as np\n'), ((88, 14, 88, 69), 'gpflux.encoders.DirectlyParameterizedNormalDiag', 'DirectlyParameterizedNormalDiag', ({(88, 46, 88, 54): 'num_data', (88, 56, 88, 61): 'w_dim', (88, 63, 88, 68): 'means'}, {}), '(num_data, w_dim, means)', False, 'from gpflux.encoders import DirectlyParameterizedNormalDiag\n'), ((98, 23, 98, 45), 'gpflow.kullback_leiblers.gauss_kl', 'gauss_kl', ({(98, 32, 98, 36): 'q_mu', (98, 38, 98, 44): 'q_sqrt'}, {}), '(q_mu, q_sqrt)', False, 'from gpflow.kullback_leiblers import gauss_kl\n'), ((101, 4, 101, 75), 'numpy.testing.assert_allclose', 'np.testing.assert_allclose', (), '', True, 'import numpy as np\n'), ((130, 9, 130, 58), 'gpflux.layers.LatentVariableLayer', 'LatentVariableLayer', (), '', False, 'from gpflux.layers import LatentVariableLayer, LayerWithObservations, TrackableLayer\n'), ((132, 13, 132, 47), 'numpy.full', 'np.full', ({(132, 21, 132, 38): '(num_data, x_dim)', (132, 40, 132, 46): 'np.nan'}, {}), '((num_data, x_dim), np.nan)', True, 'import numpy as np\n'), ((133, 14, 133, 48), 'numpy.full', 'np.full', ({(133, 22, 133, 39): '(num_data, y_dim)', (133, 41, 133, 47): 'np.nan'}, {}), '((num_data, y_dim), np.nan)', True, 'import numpy as np\n'), ((135, 21, 135, 58), 'numpy.concatenate', 'np.concatenate', (), '', True, 'import numpy as np\n'), ((147, 4, 147, 53), 'numpy.testing.assert_equal', 'np.testing.assert_equal', ({(147, 28, 147, 37): 'lv.losses', (147, 39, 147, 52): 'expected_loss'}, {}), '(lv.losses, expected_loss)', True, 'import numpy as np\n'), ((172, 9, 172, 58), 'gpflux.layers.LatentVariableLayer', 'LatentVariableLayer', (), '', False, 'from gpflux.layers import LatentVariableLayer, LayerWithObservations, TrackableLayer\n'), ((174, 4, 174, 28), 'tensorflow.random.set_seed', 'tf.random.set_seed', ({(174, 23, 174, 27): 'seed'}, {}), '(seed)', True, 'import tensorflow as tf\n'), ((176, 4, 176, 28), 'tensorflow.random.set_seed', 'tf.random.set_seed', ({(176, 23, 176, 27): 'seed'}, {}), '(seed)', True, 'import tensorflow as tf\n'), ((178, 4, 178, 63), 'numpy.testing.assert_array_equal', 'np.testing.assert_array_equal', ({(178, 34, 178, 46): 'sample_prior', (178, 48, 178, 62): 'prior_expected'}, {}), '(sample_prior, prior_expected)', True, 'import numpy as np\n'), ((180, 4, 180, 28), 'tensorflow.random.set_seed', 'tf.random.set_seed', ({(180, 23, 180, 27): 'seed'}, {}), '(seed)', True, 'import tensorflow as tf\n'), ((182, 4, 182, 28), 'tensorflow.random.set_seed', 'tf.random.set_seed', ({(182, 23, 182, 27): 'seed'}, {}), '(seed)', True, 'import tensorflow as tf\n'), ((184, 4, 184, 71), 'numpy.testing.assert_array_equal', 'np.testing.assert_array_equal', ({(184, 34, 184, 50): 'sample_posterior', (184, 52, 184, 70): 'posterior_expected'}, {}), '(sample_posterior, posterior_expected)', True, 'import numpy as np\n'), ((45, 15, 45, 26), 'numpy.zeros', 'np.zeros', ({(45, 24, 45, 25): 'd'}, {}), '(d)', True, 'import numpy as np\n'), ((46, 22, 46, 31), 'numpy.eye', 'np.eye', ({(46, 29, 46, 30): 'd'}, {}), '(d)', True, 'import numpy as np\n'), ((47, 14, 47, 73), 'tensorflow_probability.distributions.MultivariateNormalTriL', 'tfp.distributions.MultivariateNormalTriL', ({(47, 55, 47, 59): 'mean', (47, 61, 47, 72): 'scale_tri_l'}, {}), '(mean, scale_tri_l)', True, 'import tensorflow_probability as tfp\n'), ((49, 14, 49, 24), 'numpy.ones', 'np.ones', ({(49, 22, 49, 23): 'd'}, {}), '(d)', True, 'import numpy as np\n'), ((50, 19, 50, 70), 'tensorflow_probability.distributions.MultivariateNormalDiag', 'tfp.distributions.MultivariateNormalDiag', ({(50, 60, 50, 64): 'mean', (50, 66, 50, 69): 'std'}, {}), '(mean, std)', True, 'import tensorflow_probability as tfp\n'), ((109, 15, 109, 65), 'numpy.allclose', 'np.allclose', (), '', True, 'import numpy as np\n'), ((38, 56, 38, 71), 'numpy.zeros', 'np.zeros', ({(38, 65, 38, 70): 'w_dim'}, {}), '(w_dim)', True, 'import numpy as np\n'), ((38, 84, 38, 98), 'numpy.ones', 'np.ones', ({(38, 92, 38, 97): 'w_dim'}, {}), '(w_dim)', True, 'import numpy as np\n'), ((92, 9, 92, 37), 'numpy.random.randn', 'np.random.randn', ({(92, 25, 92, 33): 'num_data', (92, 35, 92, 36): '3'}, {}), '(num_data, 3)', True, 'import numpy as np\n'), ((92, 39, 92, 67), 'numpy.random.randn', 'np.random.randn', ({(92, 55, 92, 63): 'num_data', (92, 65, 92, 66): '2'}, {}), '(num_data, 2)', True, 'import numpy as np\n'), ((120, 12, 120, 41), 'numpy.random.randn', 'np.random.randn', ({(120, 28, 120, 40): '*prior_shape'}, {}), '(*prior_shape)', True, 'import numpy as np\n'), ((124, 12, 124, 46), 'numpy.random.randn', 'np.random.randn', ({(124, 28, 124, 45): '*posteriors_shape'}, {}), '(*posteriors_shape)', True, 'import numpy as np\n'), ((162, 12, 162, 41), 'numpy.random.randn', 'np.random.randn', ({(162, 28, 162, 40): '*prior_shape'}, {}), '(*prior_shape)', True, 'import numpy as np\n'), ((166, 12, 166, 46), 'numpy.random.randn', 'np.random.randn', ({(166, 28, 166, 45): '*posteriors_shape'}, {}), '(*posteriors_shape)', True, 'import numpy as np\n'), ((121, 19, 121, 48), 'numpy.random.randn', 'np.random.randn', ({(121, 35, 121, 47): '*prior_shape'}, {}), '(*prior_shape)', True, 'import numpy as np\n'), ((125, 19, 125, 53), 'numpy.random.randn', 'np.random.randn', ({(125, 35, 125, 52): '*posteriors_shape'}, {}), '(*posteriors_shape)', True, 'import numpy as np\n'), ((163, 19, 163, 48), 'numpy.random.randn', 'np.random.randn', ({(163, 35, 163, 47): '*prior_shape'}, {}), '(*prior_shape)', True, 'import numpy as np\n'), ((167, 19, 167, 53), 'numpy.random.randn', 'np.random.randn', ({(167, 35, 167, 52): '*posteriors_shape'}, {}), '(*posteriors_shape)', True, 'import numpy as np\n')]