%matplotlib inline

import numpy as np

import matplotlib.pyplot as pl

from revrand.basis_functions import RandomRBF, RandomLaplace, RandomCauchy, RandomMatern32, RandomMatern52, \
    FastFoodRBF, OrthogonalRBF, FastFoodGM, BasisCat
from revrand import Parameter, Positive


# Style
pl.style.use('ggplot')
pl.rc('font', **{'size': 26})

"""
Explanation: Random Basis Functions Test
In this notebook we test our random basis functions against the kernel functions they are designed to approximate. This is a qualitative test in that we just plot the inner product of our bases with the kernel evaluations for an array of inputs, i.e., we compare the shapes of the kernels.
We evaluate these kernels in D > 1, since their Fourier transformation may be a function of D.
End of explanation
"""


N = 1000
D = 10  # Kernel dimension
nbases = 1500
lenscale = 1.5
mean = -1.
inrange = 4

# Data
x = np.zeros((N, D))
y = np.tile(np.linspace(-inrange, inrange, N), (D, 1)).T
    

"""
Explanation: Settings
End of explanation
"""


def dist(power=2):
    return (np.abs((x - y)**power)).sum(axis=1)

# RBF
def kern_rbf():
    return np.exp(- dist() / (2 * lenscale**2))

# Cauchy
def kern_cau():
    return 1. / (1 + dist() / lenscale**2)

# Laplace
def kern_lap():
    return np.exp(- dist(power=1) / lenscale)

# Matern 3/2
def kern_m32():
    dterm = np.sqrt(3) * np.sqrt(dist()) / lenscale
    return (1 + dterm) * np.exp(-dterm)
    
# Matern 5/2
def kern_m52():
    dterm = np.sqrt(5) * np.sqrt(dist()) / lenscale
    return (1 + dterm + dterm**2 / 3.) * np.exp(-dterm)

def kern_combo():
    return 0.5 * kern_lap() + 2 * kern_rbf()


"""
Explanation: Kernel functions
End of explanation
"""


rbf = RandomRBF(Xdim=D, nbases=nbases)
cau = RandomCauchy(Xdim=D, nbases=nbases)
lap = RandomLaplace(Xdim=D, nbases=nbases)
m32 = RandomMatern32(Xdim=D, nbases=nbases)
m52 = RandomMatern52(Xdim=D, nbases=nbases)
ff_rbf = FastFoodRBF(Xdim=D, nbases=nbases)
or_rbf = OrthogonalRBF(Xdim=D, nbases=nbases)

r_lap = Parameter(0.5, Positive())
r_rbf = Parameter(2., Positive())
combo = RandomLaplace(Xdim=D, nbases=nbases, regularizer=r_lap) + \
    RandomRBF(Xdim=D, nbases=nbases, regularizer=r_rbf)

# Get expected kernel evaluations
def radialbasis2kern(basis):
    V, _ = basis.regularizer_diagonal(x)
    l = [lenscale] * len(basis.bases) if isinstance(basis, BasisCat) else [lenscale]
        
    return (basis.transform(x, *l) * basis.transform(y, *l)).dot(V)


"""
Explanation: Basis functions
End of explanation
"""


k_rbf = kern_rbf()
b_rbf = radialbasis2kern(rbf)

k_cau = kern_cau()
b_cau = radialbasis2kern(cau)

k_lap = kern_lap()
b_lap = radialbasis2kern(lap)

k_m32 = kern_m32()
b_m32 = radialbasis2kern(m32)

k_m52 = kern_m52()
b_m52 = radialbasis2kern(m52)

f_rbf = radialbasis2kern(ff_rbf)
o_rbf = radialbasis2kern(or_rbf)

k_combo = kern_combo()
f_combo = radialbasis2kern(combo)

"""
Explanation: Evaluate kernels and bases
End of explanation
"""


distfrom00 = np.sign(y[:, 0]) * np.sqrt(dist(power=2))
def plotkern(k1, k2, k1_label=None, k2_label=None):
    pl.figure(figsize=(15, 10))
    pl.plot(distfrom00, k1, 'b', linewidth=3, alpha=0.5, label=k1_label)
    pl.plot(distfrom00, k2, 'r--', linewidth=3, alpha=0.7, label=k2_label)
    pl.grid(True)
    pl.axis('tight')
    pl.xlabel('$\| x - y \|$')
    pl.ylabel('$k(x - y)$')
    pl.legend()
    pl.show()
    
plotkern(k_rbf, b_rbf, 'RBF kernel', 'RBF basis')
plotkern(k_cau, b_cau, 'Cauchy kernel', 'Cauchy basis')
plotkern(k_lap, b_lap, 'Laplace kernel', 'Laplace basis')
plotkern(k_m32, b_m32, 'Matern32 kernel', 'Matern32 basis')
plotkern(k_m52, b_m52, 'Matern52 kernel', 'Matern52 basis')
plotkern(k_rbf, f_rbf, 'RBF kernel', 'FastFood RBF basis')
plotkern(k_rbf, o_rbf, 'RBF kernel', 'Orthogonal RBF basis')
plotkern(k_combo, f_combo, 'Combo kernel', 'Combo basis')

"""
Explanation: Plot the kernel functions
End of explanation
"""


