%matplotlib inline
from IPython.display import display
import numpy as np
import matplotlib.pyplot as plt
import pods
import GPy

kern = GPy.kern.RBF(input_dim=1)
display(kern)

kern = GPy.kern.RBF(input_dim=1, name='signal', variance=4.0, lengthscale=2.0)
display(kern)

kern.plot()

kern.lengthscale = 3.5
display(kern)

kern.lengthscale.name = 'timescale'
display(kern)

kern.timescale = 10.
display(kern)

data = pods.datasets.olympic_marathon_men()
# Load in the times of the olympics
X = data['X']
K=kern.K(X)

def visualize_olympics(K):
    """Helper function for visualizing a covariance computed on the Olympics training data."""
    fig, ax = plt.subplots(figsize=(8,8))
    im = ax.imshow(K, interpolation='None')

    WWI_index = np.argwhere(X==1912)[0][0]
    WWII_index = np.argwhere(X==1936)[0][0]

    ax.axhline(WWI_index+0.5,color='w')
    ax.axvline(WWI_index+0.5,color='w')
    ax.axhline(WWII_index+0.5,color='w')
    ax.axvline(WWII_index+0.5,color='w')
    plt.colorbar(im)
    
visualize_olympics(kern.K(X))

kern.timescale

kern.timescale = 20
visualize_olympics(kern.K(X))

def sample_covariance(kern, X, num_samps=10):
    """Sample a one dimensional function as if its from a Gaussian process with the given covariance function."""
    from IPython.display import HTML
    display(HTML('<h2>Samples from a Gaussian Process with ' + kern.name + ' Covariance</h2>'))
    display(kern)
    K = kern.K(X) 

    # Generate samples paths from a Gaussian with zero mean and covariance K
    F = np.random.multivariate_normal(np.zeros(X.shape[0]), K, num_samps).T

    fig, ax = plt.subplots(figsize=(8,8))
    ax.plot(X,F)



# create an input vector
X = np.linspace(-2, 2, 200)[:, None]

# create a covariance to visualize
kern = GPy.kern.RBF(input_dim=1)

# perform the samples.
sample_covariance(kern, X)

kern = GPy.kern.Poly(input_dim=1, order=6)
X = np.linspace(-.8, .8, 500)[:, None]
sample_covariance(kern, X)

k1 = GPy.kern.RBF(1, variance=4.0, lengthscale=10., name='long term trend')
k2 = GPy.kern.RBF(1, variance=1.0, lengthscale=2., name='short term trend')
kern = k1 + k2
kern.name = 'signal'
kern.long_term_trend.lengthscale.name = 'timescale'
kern.short_term_trend.lengthscale.name = 'timescale'
display(kern)

k1 = GPy.kern.Linear(1)
k2 = GPy.kern.RBF(1)
kern = k1*k2
display(kern)
fig, ax = plt.subplots(figsize=(8,8))
im = ax.imshow(kern.K(X), interpolation='None')
plt.colorbar(im)

# Try plotting sample paths here
# Exercise 4 answer