# -*- coding: utf-8 -*-
"""
Created on Fri Dec  8 08:19:48 2017

@author: Carl Steyn
"""

from __future__ import division
import numpy as np
import matplotlib.pyplot as pl
#%%
# Unknown function we're approximating
f = lambda x: np.sin(0.9*x).flatten()

def kernel(a,b):
    kernelParameter = 1
    sqdist = np.sum(a**2,1).reshape(-1,1) + np.sum(b**2,1) - 2*np.dot(a,b.T)
    return np.exp(-.5*(1/kernelParameter)*sqdist)
N = 10 #training points
n = 50 #test points
s = 0.00005 #noise variance
#Sample input points and noisy versions of the function evaluated at these points
X = np.random.uniform(-5,5,size=(N,1))
y = f(X) + s*np.random.randn(N)

K = kernel(X,X)
L = np.linalg.cholesky(K + s*np.eye(N))

# points we're going to make predictions at
Xtest = np.linspace(-5,5,n).reshape(-1,1)

# compute the posterior mean at our test points
Lk = np.linalg.solve(L,kernel(X,Xtest))
mu = np.dot(Lk.T, np.linalg.solve(L,y))

# compute variance at test points
K_ = kernel(Xtest,Xtest)
s2 = np.diag(K_) - np.sum(Lk**2,axis=0)
s = np.sqrt(s2)

#%%
# PLOTS:
pl.figure(1)
pl.clf()
pl.plot(X, y, 'r+', ms=20)
pl.plot(Xtest, f(Xtest), 'b-')
pl.gca().fill_between(Xtest.flat, mu-3*s, mu+3*s, color="#dddddd")
pl.plot(Xtest, mu, 'r--', lw=2)
pl.savefig('predictive.png', bbox_inches='tight')
pl.title('Mean predictions plus 3 st.deviations')
pl.axis([-5, 5, -3, 3])
#%%
# draw samples from the prior at our test points.
L = np.linalg.cholesky(K_ + 1e-6*np.eye(n))
f_prior = np.dot(L, np.random.normal(size=(n,10)))
pl.figure(2)
pl.clf()
pl.plot(Xtest, f_prior)
pl.title('Ten samples from the GP prior')
pl.axis([-5, 5, -3, 3])
pl.savefig('prior.png', bbox_inches='tight')
#%%
# draw samples from the posterior at our test points.
L = np.linalg.cholesky(K_ + 1e-6*np.eye(n) - np.dot(Lk.T, Lk))
f_post = mu.reshape(-1,1) + np.dot(L, np.random.normal(size=(n,10)))
pl.figure(3)
pl.clf()
pl.plot(Xtest, f_post)
pl.title('Ten samples from the GP posterior')
pl.axis([-5, 5, -3, 3])
pl.savefig('post.png', bbox_inches='tight')

pl.show()