#@title ##### Licensed under the Apache License, Version 2.0 (the "License"); { display-mode: "form" }
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://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.

"""
Explanation: Copyright 2020 The TensorFlow Probability Authors.
End of explanation
"""


# Importing the TFP with Jax backend
!pip3 install -q 'tfp-nightly[jax]' tf-nightly-cpu  # We (currently) still require TF, but TF's smaller CPU build will work.
import tensorflow_probability as tfp
tfp = tfp.experimental.substrates.jax
tf = tfp.tf2jax

# Standard TFP Imports
tfd = tfp.distributions
tfb = tfp.bijectors
tfpk = tfp.math.psd_kernels

# Jax imports
import jax
import jax.numpy as np
from jax import random

# Other imports
import matplotlib.pyplot as plt
import seaborn as sns
sns.set(style='white')

"""
Explanation: TFP, backed by Jax
Jax-backed TFP is a work in progress, but many distributions and bijectors are currently working! How do you use the alternative backend?
<table class="tfo-notebook-buttons" align="left">
  <td>
    <a target="_blank" href="https://colab.research.google.com/github/tensorflow/probability/blob/main/discussion/examples/TFP_and_Jax.ipynb"><img src="https://www.tensorflow.org/images/colab_logo_32px.png" />Run in Google Colab</a>
  </td>
  <td>
    <a target="_blank" href="https://github.com/tensorflow/probability/blob/main/discussion/examples/TFP_and_Jax.ipynb"><img src="https://www.tensorflow.org/images/GitHub-Mark-32px.png" />View source on GitHub</a>
  </td>
</table>

Importing
End of explanation
"""


tf.ones(5)

tf.matmul(tf.ones([1, 2]), tf.ones([2, 4]))

"""
Explanation: TF Interface to Jax
We've reimplemented the TF API, but with Jax functions instead of TF functions and DeviceArrays instead of TF Tensors.
End of explanation
"""


tf.ones(5).shape

"""
Explanation: Some differences:
Shapes are tuples, not TensorShapes
End of explanation
"""


tf.random.stateless_uniform([1, 2], seed=random.PRNGKey(0))

"""
Explanation: Randomness is stateless, like in Jax and requires Jax PRNGKeys to operate.
End of explanation
"""


tf.compat.v1.placeholder_with_default(tf.ones(5), (5,))

"""
Explanation: Placeholders don't exist.
End of explanation
"""


bij = tfb.Shift(1.)(tfb.Scale(3.))
print(bij.forward(np.ones(5)))
print(bij.inverse(np.ones(5)))

"""
Explanation: Math libraries
TFP's math libraries are now largely working, i.e. tfp.math
Bijectors
Most bijectors have tests passing!
Unary bijectors
End of explanation
"""


b = tfb.FillScaleTriL(diag_bijector=tfb.Exp(), diag_shift=None)
print(b.forward(x=[0., 0., 0.]))
print(b.inverse(y=[[1., 0], [.5, 2]]))

b = tfb.Chain([tfb.Exp(), tfb.Softplus()])
# or: 
# b = tfb.Exp()(tfb.Softplus())
print(b.forward(-np.ones(5)))

"""
Explanation: Meta bijectors
End of explanation
"""


dist = tfd.Normal(loc=0., scale=1.)
print(dist.sample(seed=random.PRNGKey(0)))

"""
Explanation: MCMC coming soon
We are migrating TFP's random samplers to be internally-stateless, then will update MCMC to support JAX.
Some don't work yet

For example: FFJORD, MAF (WIP), Real NVP

Distributions
When sampling, we need to pass in a seed.
End of explanation
"""


dist =  tfd.Normal(np.zeros(5), np.ones(5))
s = dist.sample(sample_shape=(10, 2), seed=random.PRNGKey(0))
print(dist.log_prob(s).shape)

dist =  tfd.Independent(tfd.Normal(np.zeros(5), np.ones(5)), 1)
s = dist.sample(sample_shape=(10, 2), seed=random.PRNGKey(0))
print(dist.log_prob(s).shape)

"""
Explanation: Jax distributions obey the same batching semantics as their TensorFlow counterparts.
End of explanation
"""


dist = tfd.TransformedDistribution(
    tfd.MultivariateNormalDiag(tf.zeros(5), tf.ones(5)),
    tfb.Exp())
# or:
# dist = tfb.Exp()(tfd.MultivariateNormalDiag(tf.zeros(5), tf.ones(5)))
s = dist.sample(sample_shape=2, seed=random.PRNGKey(0))
print(s)
print(dist.log_prob(s).shape)

"""
Explanation: Most meta distributions are working!
End of explanation
"""


k1, k2, k3 = random.split(random.PRNGKey(0), 3)
observation_noise_variance = 0.01
f = lambda x: np.sin(10*x[..., 0]) * np.exp(-x[..., 0]**2)
observation_index_points = tf.random.stateless_uniform(
    [50], minval=-1.,maxval= 1., seed=k1)[..., np.newaxis]
observations = f(observation_index_points) + tfd.Normal(loc=0., scale=np.sqrt(observation_noise_variance)).sample(seed=k2)

index_points = np.linspace(-1., 1., 100)[..., np.newaxis]

kernel = tfpk.ExponentiatedQuadratic(length_scale=0.1)

gprm = tfd.GaussianProcessRegressionModel(
    kernel=kernel,
    index_points=index_points,
    observation_index_points=observation_index_points,
    observations=observations,
    observation_noise_variance=observation_noise_variance)

samples = gprm.sample(10, seed=k3)
for i in range(10):
  plt.plot(index_points, samples[i])
plt.show()

"""
Explanation: Gaussian processes and PSD kernels also work.
End of explanation
"""


