import numpy as np

from scipy.stats import multivariate_normal as mvn

import matplotlib.pyplot as plt

import seaborn as sns

from sklearn.gaussian_process.kernels import RBF

import matplotlib

font = {"size": 30}

matplotlib.rc("font", **font)

matplotlib.rcParams["text.usetex"] = True

lengthscale = 1.0

amplitude = 1.0

noise_stddev = 1e-6

xlims = [-5, 5]

n = 100

X = np.linspace(xlims[0], xlims[1], n)

X = np.expand_dims(X, 1)

## Draw function

K_XX = amplitude * RBF(length_scale=lengthscale)(X, X) + noise_stddev * np.eye(n)

mean = X.squeeze()

# mean = np.zeros(n)

Y = mvn(mean, K_XX).rvs()

# import ipdb; ipdb.set_trace()

plt.figure(figsize=(7, 6))

plt.plot(X, Y, linewidth=5)

plt.xlabel("Observed spatial coordinate")

plt.ylabel("Warped spatial coordinate")

plt.title(r"$\sigma^2 = {}, \ell = {}$".format(amplitude, lengthscale))

plt.tight_layout()

plt.savefig("../../plots/mean_function_example.png")

plt.show()