import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
from gpflow.kernels import White, RBF
from gpflow.likelihoods import Gaussian

from deep_gp import DeepGP

np.random.seed(0)
tf.random.set_seed(0)


def get_data():
    Ns = 300
    Xs = np.linspace(-0.5, 1.5, Ns)[:, None]

    N, M = 50, 25
    X = np.random.uniform(0, 1, N)[:, None]
    Z = np.random.uniform(0, 1, M)[:, None]
    f_step = lambda x: 0. if x < 0.5 else 1.
    Y = np.reshape([f_step(x) for x in X], X.shape) + np.random.randn(
        *X.shape) * 1e-2
    return Xs, X, Y, Z


def make_deep_GP(num_layers, X, Y, Z):
    kernels = []
    layer_sizes = []
    for l in range(num_layers):
        kernel = RBF(lengthscales=0.2, variance=1.0) + White(variance=1e-5)
        kernels.append(kernel)
        layer_sizes.append(1)

    dgp = DeepGP(X, Y, Z, kernels, layer_sizes, Gaussian(), num_samples=100)

    # init hidden layers to be near deterministic
    for layer in dgp.layers[:-1]:
        layer.q_sqrt.assign(layer.q_sqrt * 1e-5)
    return dgp


if __name__ == '__main__':
    Xs, X_train, Y_train, Z = get_data()
    dgp = make_deep_GP(3, X_train, Y_train, Z)
    optimizer = tf.optimizers.Adam(learning_rate=0.01, epsilon=1e-08)

    for _ in range(1500):
        with tf.GradientTape(watch_accessed_variables=False) as tape:
            tape.watch(dgp.trainable_variables)
            objective = -dgp.elbo((X_train, Y_train))
            gradients = tape.gradient(objective, dgp.trainable_variables)
        optimizer.apply_gradients(zip(gradients, dgp.trainable_variables))
        print(f"ELBO: {-objective.numpy()}")

    samples, _, _ = dgp.predict_all_layers(Xs, num_samples=50, full_cov=True)
    plt.plot(Xs, samples[-1].numpy()[:, :, 0].T, color='r', alpha=0.3)

    plt.title('Deep Gaussian Process')
    plt.scatter(X_train, Y_train)
    plt.show()
[((9, 0, 9, 17), 'numpy.random.seed', 'np.random.seed', ({(9, 15, 9, 16): '(0)'}, {}), '(0)', True, 'import numpy as np\n'), ((10, 0, 10, 21), 'tensorflow.random.set_seed', 'tf.random.set_seed', ({(10, 19, 10, 20): '(0)'}, {}), '(0)', True, 'import tensorflow as tf\n'), ((45, 16, 45, 69), 'tensorflow.optimizers.Adam', 'tf.optimizers.Adam', (), '', True, 'import tensorflow as tf\n'), ((58, 4, 58, 38), 'matplotlib.pyplot.title', 'plt.title', ({(58, 14, 58, 37): '"""Deep Gaussian Process"""'}, {}), "('Deep Gaussian Process')", True, 'import matplotlib.pyplot as plt\n'), ((59, 4, 59, 33), 'matplotlib.pyplot.scatter', 'plt.scatter', ({(59, 16, 59, 23): 'X_train', (59, 25, 59, 32): 'Y_train'}, {}), '(X_train, Y_train)', True, 'import matplotlib.pyplot as plt\n'), ((60, 4, 60, 14), 'matplotlib.pyplot.show', 'plt.show', ({}, {}), '()', True, 'import matplotlib.pyplot as plt\n'), ((15, 9, 15, 35), 'numpy.linspace', 'np.linspace', ({(15, 21, 15, 25): '(-0.5)', (15, 27, 15, 30): '(1.5)', (15, 32, 15, 34): 'Ns'}, {}), '(-0.5, 1.5, Ns)', True, 'import numpy as np\n'), ((18, 8, 18, 34), 'numpy.random.uniform', 'np.random.uniform', ({(18, 26, 18, 27): '(0)', (18, 29, 18, 30): '(1)', (18, 32, 18, 33): 'N'}, {}), '(0, 1, N)', True, 'import numpy as np\n'), ((19, 8, 19, 34), 'numpy.random.uniform', 'np.random.uniform', ({(19, 26, 19, 27): '(0)', (19, 29, 19, 30): '(1)', (19, 32, 19, 33): 'M'}, {}), '(0, 1, M)', True, 'import numpy as np\n'), ((34, 48, 34, 58), 'gpflow.likelihoods.Gaussian', 'Gaussian', ({}, {}), '()', False, 'from gpflow.likelihoods import Gaussian\n'), ((21, 54, 22, 17), 'numpy.random.randn', 'np.random.randn', ({(22, 8, 22, 16): '*X.shape'}, {}), '(*X.shape)', True, 'import numpy as np\n'), ((30, 17, 30, 52), 'gpflow.kernels.RBF', 'RBF', (), '', False, 'from gpflow.kernels import White, RBF\n'), ((30, 55, 30, 75), 'gpflow.kernels.White', 'White', (), '', False, 'from gpflow.kernels import White, RBF\n'), ((48, 13, 48, 60), 'tensorflow.GradientTape', 'tf.GradientTape', (), '', True, 'import tensorflow as tf\n')]