from scipy.optimize import fminbound

def abspline3_kernel(X, alpha, beta, t1, t2):
    R = np.zeros(X.shape)
    M = np.array([[1, t1, t1**2., t1**3.],
                  [1, t2, t2**2., t2**3.],
                  [0, 1, 2.*t1, 3.*t1**2.],
                  [0, 1, 2.*t2, 3.*t2**2.]])
    v = np.array([1., 1., t1**(-1.0*alpha)*alpha*t1**(alpha-1.), \
                  -1.0*beta*t2**(-1.0*beta-1.)*t2**beta])
    a = np.linalg.solve(M, v)
    r1 = (X >= 0.) & (X < t1)
    r2 = (X >= t1) & (X < t2)
    r3 = X >= t2
       
    R[r1] = np.power(X[r1], alpha)*t1**(-1.0*alpha)
    R[r3] = np.power(X[r3], -1.0*beta)*t2**beta
    X2 = X[r2]
    R[r2] = a[0]+a[1]*X2+a[2]*np.power(X2,2)+a[3]*np.power(X2, 3)
    return R

def abspline3(lmax, nscales, lpfactor=20., a=2., b=2., t1=1., t2=2.):
    lmin = lmax/lpfactor
    t = np.exp(np.linspace(np.log(2./lmin), 
                           np.log(1./lmax), 
                           nscales))
    
    gl = lambda x: np.exp(-np.power(x, 4.))
    gb = lambda x: abspline3_kernel(x, a, b, t1, t2)
    f = lambda x: -1.0*gb(x)
    xstar = fminbound(f, 1., 2.)
    gamma_l = -f(xstar)
    lminfac = 0.6*lmin

    g = [lambda x: gamma_l*gl((1.0/lminfac)*x)] +\
        [lambda x,y=scale: gb(y*x) for scale in t]
    return g,t



def chebyshev_coeff(g, m, n, arange):
    a1 = (arange[1] - arange[0])/2.
    a2 = (arange[1] + arange[0])/2.
    c = np.empty(m+1)
    d = (np.arange(1,n+1)-0.5)*(1./n)
    for j in range(m+1):
        c[j] = np.sum(g(a1*np.cos(np.pi*d) + a2) *\
                      np.cos(np.pi*j*d))*2./n
    return c

def chebyshev_apply(f, L, c, arange):
    nscales = len(c)
    M = map(len, c)
    maxm = np.max(M)
    
    a1 = (arange[1] - arange[0])/2.
    a2 = (arange[1] + arange[0])/2.
    twf_old = f
    twf_cur = (L.dot(f) - a2*f)*(1./a1)
    
    r = [0.5*c[i][0]*twf_old + c[i][1]*twf_cur for i in range(nscales)]
    
    for k in range(1, maxm):
        twf_new = (2./a1)*(L.dot(twf_cur) - a2*twf_cur) - twf_old
        for i in range(nscales):
            if k+1 < M[i]:
                r[i] = r[i] + c[i][k+1]*twf_new
        twf_old = twf_cur
        twf_cur = twf_new
    return r  

def onehot(n, i):
    x = np.zeros(n)
    x[i] = 1.
    return x

from scipy.io import loadmat
L = loadmat('laplacian.mat')['L']

g, t = abspline3(7, 4)
c = [chebyshev_coeff(g[0], 50, 51, [0., 7.])]
d = onehot(500, 32)
wall = chebyshev_apply(d, L, c, [0., 7.])

from mpl_toolkits.mplot3d import Axes3D

X = loadmat('swiss.mat')['x'].T

fig = plt.figure(figsize=(12, 4))
ax = fig.add_subplot(121, projection='3d')
ax.scatter(X[:,0], X[:,1], X[:,2], c=d)
ax.view_init(elev=15, azim=-90)

ax = fig.add_subplot(122, projection='3d')
ax.scatter(X[:,0], X[:,1], X[:,2], c=wall[0])
ax.view_init(elev=15, azim=-90)