import numpy as np
from scipy.spatial.distance import cdist
import scipy.linalg

import matplotlib.pyplot as plt

class SEKernel:
    def k(self, x, xPrime):
        x = np.reshape(x, [-1, self.d])
        xPrime = np.reshape(xPrime, [-1, self.d])
        distances = cdist(x, xPrime, 'sqeuclidean')
        k = self.a*np.exp(-distances / self.lR / self.lR)
        return k

    def kDiag(self, x):
        return self.a*np.ones(np.reshape(x, [-1, self.d]).shape[0])

    def setDim(self, d):
        self.d = d

    def __init__(self, lR=0.2, a=1.0):
        self.d = None
        self.lR = lR
        self.a = a

class GP:
    def trainGP(self):
        X = self.X
        y = self.C
        K = self.kernel.k(X, X)
        G = K + self.sigma0*self.sigma0*np.eye(X.shape[0])
        self.GInv = np.linalg.inv(G)
        self.b = self.GInv@y

    def mu(self, x):
        kappa = self.kernel.k(x, self.X)
        cPred = np.matmul(kappa, self.b)
        return cPred

    def sigma(self, x):
        kxx = self.kernel.kDiag(x)
        kXx = self.kernel.k(self.X, x)
        sig = kxx - np.diag(kXx.T@self.GInv@kXx) # this is horribly inefficient, there is a much better way utilizing Cholesky decompositions
        return np.reshape(sig, [-1,1])

    def addDataPoint(self, x, c):
        x = np.reshape(x, (-1, self.d))
        c = np.reshape(c, (-1, 1))

        self.X = np.vstack((self.X, x))
        self.C = np.vstack((self.C, c))


    def clearData(self):
        self.X = np.empty((0, self.d))
        self.C = np.empty((0, 1))
        self.L = None
        self.b = None

    def __init__(self, kernel, d=3, sigma0 = 0.001):
        self.kernel = kernel
        self.d = d
        self.sigma0 = sigma0

        self.kernel.setDim(self.d)

        self.L = None
        self.b = None

        self.X = np.empty((0, self.d))
        self.C = np.empty((0, 1))



if __name__ == "__main__":
    kernel = SEKernel(lR=0.8, a=2.0)
    gp = GP(d=1, kernel=kernel, sigma0=0.1)
    X = np.array([[0.5], [4], [7]])
    C = np.array([1, 3, -2])
    gp.addDataPoint(X, C)
    gp.trainGP()


    xTest = np.linspace(-1,10,100)
    yTest = gp.mu(xTest)
    sigTest = gp.sigma(xTest)

    plt.plot(xTest, yTest)
    plt.plot(xTest, yTest+sigTest)
    plt.plot(xTest, yTest-sigTest)
    plt.show()