import scipy
import numpy as np

class GaussianProcess(object):
    def __init__(self,sigma=1.0,l=1.0):
        self.sigma=sigma
        self.l=l
        self.X = []
        self.y = []
        self.deriv = []
        
    def kernel(self,x,x2,deriv,deriv2):
        k = self.sigma**2*np.exp(-0.5*(x-x2)**2/self.l**2)
        if deriv==deriv2==0: return k
        kD =(x-x2)/self.l**2 * k 
        if deriv==0 and deriv2==1: return kD 
        if deriv==1 and deriv2==0: return -kD
        kDD = k/self.l**2 - (x-x2)/self.l**2*kD
        return kDD
        
    def add(self,x,y,deriv):
        self.X.append(x)
        self.y.append(y)
        self.deriv.append(deriv)
        
    def prepare(self):
        Sigma = np.zeros((len(self.y),len(self.y)))
        for idx in range(len(self.y)):
            for idy in range(len(self.y)):
                Sigma[idx,idy] = self.kernel(
                    self.X[idx],self.X[idy],
                    self.deriv[idx],self.deriv[idy])
                
        # Computing the inverse of the covariance matrix Sigma is numerically unstable. 
        # Therefore, we do a Cholesky decomposition instead.
        self.L = scipy.linalg.cholesky(Sigma, lower=True)
        # alpha = Sigma^-1*Y
        self.alpha = scipy.linalg.cho_solve((self.L, True), np.atleast_2d(self.y).T)
        
    def predict(self,xstar):
        
        # vector of [k(xstar,x_i)]_i
        Kstar = np.array([
            self.kernel(xstar,x2,0,deriv2) for x2,deriv2 in zip(self.X,self.deriv)])
            
        # mu = Kstar * Sigma^-1 * Y = Kstar * alpha
        mu = Kstar.dot(self.alpha)[0]
            
        #compute L^-1 K
        v = scipy.linalg.solve_triangular(self.L, Kstar.T, lower=True)
        # K^T * Sigma^-1 * K = v^T * v
        sigma_sq = max(0.0,self.kernel(xstar,xstar,0,0) - v.T.dot(v))
        
        return mu, np.sqrt(sigma_sq)

