import numpy as np
import matplotlib.pyplot as plt
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF, WhiteKernel

np.random.seed(1)

def f(x):
    return x * np.sin(x)


def generate_observations(n=10):
    X = np.random.uniform(0.1, 9.9, n).reshape(-1, 1)
    y = f(X).ravel()
    dy = 0.5   1.0 * np.random.random(y.shape)
    noise = np.random.normal(0, dy)
    y  = noise
    return X, y, dy


def train_gp(X, y, dy):
    kernel = 1.0 * RBF(length_scale=1.0)   WhiteKernel(noise_level=1.0)
    if dy is None:
        alpha = 1e-10
    else:
        alpha = (dy / y) ** 2
    gp = GaussianProcessRegressor(kernel=kernel, alpha=alpha, n_restarts_optimizer=100)
    gp.fit(X, y)
    return gp



def main():
    # Generate observation   train GP
    X_obs, y_obs, dy_obs = generate_observations()
    gp = train_gp(X_obs, y_obs, dy_obs)
    gp_no_unc = train_gp(X_obs, y_obs, None)

    # Generate plotting data
    x = np.linspace(0, 10, 1000).reshape(-1, 1)
    y_true = f(x)
    y_pred, sigma = gp.predict(x, return_std=True)
    y_pred_no_unc, sigma_no_unc = gp_no_unc.predict(x, return_std=True)

    # Plot
    plt.figure(figsize=(10, 6))
    plt.plot(x, y_true, ls=':', color='k')
    plt.errorbar(X_obs.ravel(), y_obs, dy_obs, fmt='.', color='k', markersize=10, label=u'Observations')


    plt.plot(x, y_pred_no_unc, color="tab:orange")
    plt.fill_between(x.ravel(), y_pred_no_unc - 1.96 * sigma_no_unc, y_pred_no_unc   1.960 * sigma_no_unc, alpha=0.1, color="tab:orange",
                     label='GP no input uncertainty')

    plt.plot(x, y_pred, color="tab:blue")
    plt.fill_between(x.ravel(), y_pred - 1.96 * sigma, y_pred   1.960 * sigma, alpha=0.3, color="tab:blue",
                     label='GP with input uncertainty')

    plt.xlabel('$x$')
    plt.ylabel('$f(x)$')
    plt.ylim(-10, 12)
    plt.xlim(0,10)
    plt.legend(loc='upper left')
    plt.show()


if __name__ == '__main__':
    main()