import sklearn.metrics
from sklearn.gaussian_process import GaussianProcessRegressor
import numpy as np
from sklearn.gaussian_process.kernels import RBF, GenericKernelMixin, Kernel
from sklearn.metrics.pairwise import pairwise_distances, pairwise_kernels

np.random.seed(0)


class MinL2Kernel(GenericKernelMixin, Kernel):
    """
    A minimal (but valid) convolutional kernel for sequences of variable
    lengths."""

    def __init__(self):
        pass

    def _f(self, sample1, sample2):
        """
        kernel value between a pair of sequences
        """
        sample1 = sample1.reshape(-1, 3)
        sample2 = sample2.reshape(-1, 3)
        dist = pairwise_distances(sample1, sample2)
        mean_dist = dist.mean()
        closenest = np.exp(-mean_dist)
        return closenest

    def __call__(self, X, Y=None, eval_gradient=False):
        if Y is None:
            Y = X

        if eval_gradient:
            raise NotImplementedError()
        else:
            return np.array([[self._f(x, y) for y in Y] for x in X])

    def diag(self, X):
        return np.array([self._f(x, x) for x in X])

    def is_stationary(self):
        return True


class RJSDkernel(GenericKernelMixin, Kernel):
    """
    A minimal (but valid) convolutional kernel for sequences of variable
    lengths."""

    def __init__(self):
        pass

    def _f(self, sample1, sample2):
        """
        kernel value between a pair of sequences
        """
        div = RJSDk(sample1, sample2)
        closenest = np.exp(-div)
        print(f'{closenest:.4f}')
        return closenest

    def __call__(self, X, Y=None, eval_gradient=False):
        if Y is None:
            Y = X

        if eval_gradient:
            raise NotImplementedError()
        else:
            return np.array([[self._f(x, y) for y in Y] for x in X])

    def diag(self, X):
        return np.array([self._f(x, x) for x in X])

    def is_stationary(self):
        return True


def RJSDk(sample_1, sample_2):
    sample_1 = sample_1.reshape(-1, 3)
    sample_2 = sample_2.reshape(-1, 3)
    n1 = sample_1.shape[0]
    n2 = sample_2.shape[0]
    pi1 = n1 / (n1 + n2)
    pi2 = n2 / (n1 + n2)
    Z = np.concatenate([sample_1, sample_2])
    # Kz = pairwise_kernels(Z, metric='rbf', n_jobs=-1)
    Kz = pairwise_kernels(Z, metric='cosine', n_jobs=-1)
    Kx = Kz[:n1, :n1]
    Ky = Kz[n1:, n1:]

    SKz = S(Kz)
    SKx = S(Kx)
    SKy = S(Ky)
    return SKz - (pi1 * SKx + pi2 * SKy)

def S(K):
    K = K/np.trace(K)
    M = K @ np.log(K)
    s = -np.trace(M)
    return s
    # eigval, _ = np.linalg.eig(K)
    # accum = 0
    # for lamda_i in eigval:
    #     accum += (lamda_i * np.log(lamda_i))
    # return -accum


def target_function(X):
    X = X.reshape(-1,3)
    return X[:,0]**3 + 2.1*X[:,1]**2 + X[:,0] + 0.1


# X = np.random.rand(10,3)
# X /= X.sum(axis=1, keepdims=True)
# Y = np.random.rand(10,3)
# Y /= Y.sum(axis=1, keepdims=True)
#
# X = X.flatten()
# Y = Y.flatten()
#
# d = RJSDk(X, Y)
#
# print(d)
#
# import sys ; sys.exit(0)

X_train = [np.random.rand(10*3) for _ in range(15)]
y_train = [target_function(X).mean() for X in X_train]

X_test = [np.random.rand(10*3) for _ in range(11)]
y_test = [target_function(X).mean() for X in X_test]


print('fit')
#kernel = 1 * RBF(length_scale=1.0, length_scale_bounds=(1e-2, 1e2))
# kernel = MinL2Kernel()
kernel = RJSDkernel()
gaussian_process = GaussianProcessRegressor(kernel=kernel, n_restarts_optimizer=9)
gaussian_process.fit(X_train, y_train)
print('[done]')

print(gaussian_process.kernel_)

y_pred = gaussian_process.predict(X_test)

mse = np.mean((y_test - y_pred)**2)

print(mse)