#
# 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

[((897, 935), 'tensorflow.keras.backend.set_floatx', 'tf.keras.backend.set_floatx', (['"""float64"""'], {}), "('float64')\n", (924, 935), True, 'import tensorflow as tf\n'), ((2330, 2370), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""w_dim"""', '[1, 5]'], {}), "('w_dim', [1, 5])\n", (2353, 2370), False, 'import pytest\n'), ((3216, 3256), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""w_dim"""', '[1, 5]'], {}), "('w_dim', [1, 5])\n", (3239, 3256), False, 'import pytest\n'), ((4542, 4582), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""w_dim"""', '[1, 5]'], {}), "('w_dim', [1, 5])\n", (4565, 4582), False, 'import pytest\n'), ((4584, 4628), 'pytest.mark.parametrize', 'pytest.mark.parametrize', (['"""seed2"""', '[None, 42]'], {}), "('seed2', [None, 42])\n", (4607, 4628), False, 'import pytest\n'), ((1728, 1781), 'gpflux.layers.LatentVariableLayer', 'LatentVariableLayer', ([], {'encoder': 'None', 'prior': 'distribution'}), '(encoder=None, prior=distribution)\n', (1747, 1781), False, 'from gpflux.layers import LatentVariableLayer, LayerWithObservations, TrackableLayer\n'), ((2261, 2282), 'numpy.all', 'np.all', (['(local_kls > 0)'], {}), '(local_kls > 0)\n', (2267, 2282), True, 'import numpy as np\n'), ((2447, 2479), 'numpy.random.randn', 'np.random.randn', (['num_data', 'w_dim'], {}), '(num_data, w_dim)\n', (2462, 2479), True, 'import numpy as np\n'), ((2494, 2549), 'gpflux.encoders.DirectlyParameterizedNormalDiag', 'DirectlyParameterizedNormalDiag', (['num_data', 'w_dim', 'means'], {}), '(num_data, w_dim, means)\n', (2525, 2549), False, 'from gpflux.encoders import DirectlyParameterizedNormalDiag\n'), ((2866, 2888), 'gpflow.kullback_leiblers.gauss_kl', 'gauss_kl', (['q_mu', 'q_sqrt'], {}), '(q_mu, q_sqrt)\n', (2874, 2888), False, 'from gpflow.kullback_leiblers import gauss_kl\n'), ((2955, 3026), 'numpy.testing.assert_allclose', 'np.testing.assert_allclose', (['tfp_local_kls', 'gpflow_local_kls'], {'rtol': '(1e-10)'}), '(tfp_local_kls, gpflow_local_kls, rtol=1e-10)\n', (2981, 3026), True, 'import numpy as np\n'), ((3841, 3890), 'gpflux.layers.LatentVariableLayer', 'LatentVariableLayer', ([], {'encoder': 'encoder', 'prior': 'prior'}), '(encoder=encoder, prior=prior)\n', (3860, 3890), False, 'from gpflux.layers import LatentVariableLayer, LayerWithObservations, TrackableLayer\n'), ((3905, 3939), 'numpy.full', 'np.full', (['(num_data, x_dim)', 'np.nan'], {}), '((num_data, x_dim), np.nan)\n', (3912, 3939), True, 'import numpy as np\n'), ((3954, 3988), 'numpy.full', 'np.full', (['(num_data, y_dim)', 'np.nan'], {}), '((num_data, y_dim), np.nan)\n', (3961, 3988), True, 'import numpy as np\n'), ((4047, 4084), 'numpy.concatenate', 'np.concatenate', (['observations'], {'axis': '(-1)'}), '(observations, axis=-1)\n', (4061, 4084), True, 'import numpy as np\n'), ((4461, 4510), 'numpy.testing.assert_equal', 'np.testing.assert_equal', (['lv.losses', 'expected_loss'], {}), '(lv.losses, expected_loss)\n', (4484, 4510), True, 'import numpy as np\n'), ((5277, 5326), 'gpflux.layers.LatentVariableLayer', 'LatentVariableLayer', ([], {'prior': 'prior', 'encoder': 'encoder'}), '(prior=prior, encoder=encoder)\n', (5296, 5326), False, 'from gpflux.layers import LatentVariableLayer, LayerWithObservations, TrackableLayer\n'), ((5332, 5356), 'tensorflow.random.set_seed', 'tf.random.set_seed', (['seed'], {}), '(seed)\n', (5350, 5356), True, 'import tensorflow as tf\n'), ((5403, 5427), 'tensorflow.random.set_seed', 'tf.random.set_seed', (['seed'], {}), '(seed)\n', (5421, 5427), True, 'import tensorflow as tf\n'), ((5523, 5582), 'numpy.testing.assert_array_equal', 'np.testing.assert_array_equal', (['sample_prior', 'prior_expected'], {}), '(sample_prior, prior_expected)\n', (5552, 5582), True, 'import numpy as np\n'), ((5588, 5612), 'tensorflow.random.set_seed', 'tf.random.set_seed', (['seed'], {}), '(seed)\n', (5606, 5612), True, 'import tensorflow as tf\n'), ((5710, 5734), 'tensorflow.random.set_seed', 'tf.random.set_seed', (['seed'], {}), '(seed)\n', (5728, 5734), True, 'import tensorflow as tf\n'), ((5829, 5896), 'numpy.testing.assert_array_equal', 'np.testing.assert_array_equal', (['sample_posterior', 'posterior_expected'], {}), '(sample_posterior, posterior_expected)\n', (5858, 5896), True, 'import numpy as np\n'), ((1236, 1247), 'numpy.zeros', 'np.zeros', (['d'], {}), '(d)\n', (1244, 1247), True, 'import numpy as np\n'), ((1270, 1279), 'numpy.eye', 'np.eye', (['d'], {}), '(d)\n', (1276, 1279), True, 'import numpy as np\n'), ((1294, 1353), 'tensorflow_probability.distributions.MultivariateNormalTriL', 'tfp.distributions.MultivariateNormalTriL', (['mean', 'scale_tri_l'], {}), '(mean, scale_tri_l)\n', (1334, 1353), True, 'import tensorflow_probability as tfp\n'), ((1369, 1379), 'numpy.ones', 'np.ones', (['d'], {}), '(d)\n', (1376, 1379), True, 'import numpy as np\n'), ((1399, 1450), 'tensorflow_probability.distributions.MultivariateNormalDiag', 'tfp.distributions.MultivariateNormalDiag', (['mean', 'std'], {}), '(mean, std)\n', (1439, 1450), True, 'import tensorflow_probability as tfp\n'), ((3162, 3212), 'numpy.allclose', 'np.allclose', (['actual', 'self.expected'], {'equal_nan': '(True)'}), '(actual, self.expected, equal_nan=True)\n', (3173, 3212), True, 'import numpy as np\n'), ((1094, 1109), 'numpy.zeros', 'np.zeros', (['w_dim'], {}), '(w_dim)\n', (1102, 1109), True, 'import numpy as np\n'), ((1122, 1136), 'numpy.ones', 'np.ones', (['w_dim'], {}), '(w_dim)\n', (1129, 1136), True, 'import numpy as np\n'), ((2686, 2714), 'numpy.random.randn', 'np.random.randn', (['num_data', '(3)'], {}), '(num_data, 3)\n', (2701, 2714), True, 'import numpy as np\n'), ((2716, 2744), 'numpy.random.randn', 'np.random.randn', (['num_data', '(2)'], {}), '(num_data, 2)\n', (2731, 2744), True, 'import numpy as np\n'), ((3485, 3514), 'numpy.random.randn', 'np.random.randn', (['*prior_shape'], {}), '(*prior_shape)\n', (3500, 3514), True, 'import numpy as np\n'), ((3648, 3682), 'numpy.random.randn', 'np.random.randn', (['*posteriors_shape'], {}), '(*posteriors_shape)\n', (3663, 3682), True, 'import numpy as np\n'), ((4921, 4950), 'numpy.random.randn', 'np.random.randn', (['*prior_shape'], {}), '(*prior_shape)\n', (4936, 4950), True, 'import numpy as np\n'), ((5084, 5118), 'numpy.random.randn', 'np.random.randn', (['*posteriors_shape'], {}), '(*posteriors_shape)\n', (5099, 5118), True, 'import numpy as np\n'), ((3535, 3564), 'numpy.random.randn', 'np.random.randn', (['*prior_shape'], {}), '(*prior_shape)\n', (3550, 3564), True, 'import numpy as np\n'), ((3703, 3737), 'numpy.random.randn', 'np.random.randn', (['*posteriors_shape'], {}), '(*posteriors_shape)\n', (3718, 3737), True, 'import numpy as np\n'), ((4971, 5000), 'numpy.random.randn', 'np.random.randn', (['*prior_shape'], {}), '(*prior_shape)\n', (4986, 5000), True, 'import numpy as np\n'), ((5139, 5173), 'numpy.random.randn', 'np.random.randn', (['*posteriors_shape'], {}), '(*posteriors_shape)\n', (5154, 5173), True, 'import numpy as np\n')]