# -*- coding: utf-8 -*-

import numpy as np
import matplotlib.pyplot as plt
import matplotlib
#%%  Kernel matrix for given kernel and data
def kernel_matrix(K,paras, X):
    n = X.shape[0]
    return np.array([[K(X[i],X[j],paras) for i in range(n)] for j in range(n)])

def polynomial_kernel(x1,x2,c_d = [.1,2]):
    inner = np.dot(x1.T,x2)
    return np.power(inner+c_d[0] , c_d[1])

def sigmoid_kernel(x1,x2,alpha_c = [1,0]):
    inner = np.dot(x1.T,x2)
    return np.tanh(alpha_c[0]*inner+alpha_c[1])

def gaussian_kernel(x1,x2,sigmasq = .1):
    dist = np.sum(np.power(x1 - x2,2))
    return np.exp(-dist/(2*sigmasq))

def laplace_kernel(x1,x2,sigma = .1):
    dist = np.sqrt(np.sum(np.power(x1 - x2,2)))
    return np.exp(-dist/sigma)
    
#%% GD and SGD algorithm
def SGD(K,y, init, GD = True, batch_size = 1,stepsize1 = 10e-3,stepsize2 = 10e-4,iters1 = 20,iters2 = 50):
    if len(init) != K.shape[0]:
        print("The shape of initialization not correct, expect {}, get {}.".format(K.shape[0],len(init)))
    solution_path = [init]
    n = K.shape[0]
    init.reshape(-1,1)
    y.reshape(-1,1)
    if GD:
        for _ in range(iters1 + iters2):
            init.reshape(1,-1)
            temp = init - stepsize1*np.dot(K,np.dot(K,init) - y)/n
            init = temp
            solution_path.append(temp.reshape(-1))
    else:
        for _ in range(iters1):
            init.reshape(1,-1)
            samples = np.random.choice(n,batch_size,replace = False)
            temp = init - stepsize1*np.dot(K[samples].T,(np.dot(K[samples],init) - y[samples]))/batch_size
            solution_path.append(temp.reshape(-1))
            init = temp
        for _ in range(iters2):
            init.reshape(1,-1)
            samples = np.random.choice(n,batch_size,replace = False)
            temp = init - stepsize2*np.dot(K[samples].T,(np.dot(K[samples],init) - y[samples]))/batch_size
            solution_path.append(temp.reshape(-1))
            init = temp
                    
    return solution_path
        
#%%
### simulate the X matrix
np.random.seed(122)
n = 10
p = 100
X = np.random.normal(size = (n,p))
### normalize the rows of X such that l2 norms are in [.49, 1]
X_normalized = np.array([X[i]*np.random.uniform(0.7,1)/np.sqrt(np.sum(np.power(X[i],2))) for i in range(n)])
X = X_normalized
### Kernel matrix
K = kernel_matrix(polynomial_kernel,[10e-2,2],X)
### Simulate y
def f(x):
    return np.sum(np.sin(x))

y = np.array([f(X[i]) + np.random.normal(scale = 0.01) for i in range(n)])

#%% Run SGD and GD algo
matplotlib.rcParams.update({'font.size':25})

solution_GD = SGD(K,y, np.zeros(n), GD = True,stepsize1 = .1,stepsize2 = .1,iters2 = 400)
solution_SGD = SGD(K,y, np.zeros(n), GD = False,stepsize1 = .01,stepsize2 = .001,iters2 = 400)
#% Plot Rayleigh Quotient
RQ_GD = [np.sum(np.power(np.dot(K,i.reshape(-1,1)),2))/np.sum(np.power(i,2)) for i in solution_GD[1:]]
RQ_SGD = [np.sum(np.power(np.dot(K,i.reshape(-1,1)),2))/np.sum(np.power(i,2)) for i in solution_SGD[1:]]

plt.plot(RQ_GD,'-.',label = "GD")
plt.plot(RQ_SGD,label = "SGD")
plt.legend()
plt.ylim(0.5,1.4)
plt.ylabel('Rayleigh Quotient')
plt.xlabel('# iterations')

plt.savefig(r"RQ_sim.png",bbox_inches='tight')