#!/usr/bin/env python
# -*- coding: utf-8 -*-


#	Written by Chieh Wu, 9/13/2022
#	This function calculates the Gaussian Kernel by approximate it through Random fourier Feature and Orthogonal Random Feature technique.

#	**Random Fourier Feature**
#	Given data $X \in \mathbb{R}^{n \times d}$, RFF says that a kernel function is 
#	$$k(x,y) = E_{p(\omega)}[ cos(w^T x + \theta) cos(w^T y + \theta)] $$
#	and therefore can be approximated with
#	$$k(x,y) = \frac{2}{m} \sum_i^m cos(w_i^T x + \theta) cos(w_i^T y + \theta) $$
#	This implies that the feature map is 
#	$$ \phi(x) \approx \sqrt{\frac{2}{m}} [cos(w_1^T x + \theta), cos(w_2^T x + \theta), ... , cos(w_m^T x + \theta)]$$
#	Given $W \in \mathbb{R}^{d \times m}$ the entire dataset $X$, then
#	$$ \phi(X) \approx \sqrt{\frac{2}{m}} cos[XW + \theta]$$
#	The the element of the W matrix is a normal gaussian distribution divided by σ of the gaussian distribution
#	θ is randomly generated using uniform distribution between 0 to 2π


#	**Structured Orthogonal Random Feature SORF**
#	The SORF replaces the $W$ matrix with
#	$$W = \frac{\sqrt{d}}{\sigma} HD_1HD_2HD_3$$
#	$H$ is a normalized Hadamard matrix, the normalization is such that $I = H^TH$
#	Therefore $H$ is the Hadamard matrix divided by $\frac{1}{\sqrt{d}}$
#	and D is a diagonal matrix with the diagonal elements sampled from Rademacher Distribution

#	Note that since Hadamard matrix size is alway in the power of 2 $2^n$, the dimension of the data is probably not the same.
#	In these cases, pad the data with 0s as the extra dimensions. Also, the width of the SORF is always set to the dimension
#	of the data after padding. Therefore, to make the width even longer, we must generate multiple size d matrices and 
#	concatenate them together. 

import os 
import sys

if os.path.exists('/home/chieh/code/wPlotLib'):
	sys.path.insert(0,'/home/chieh/code/wPlotLib')
if os.path.exists('/home/chieh/code/wuML'):
	sys.path.insert(0,'/home/chieh/code/wuML')

import numpy as np
import sklearn.metrics
import torch
import numpy.matlib
from torch.autograd import Variable
from numpy.random import rand
import torch.nn.functional as F
from scipy.linalg import hadamard
from sklearn.metrics import mean_absolute_error
from sklearn.metrics import mean_squared_error
import wuml
from wuml import jupyter_print

class random_feature:
	# sample_num, the larger the better approximation
	# random_feature_method can be 'orthogonal' (default) or 'rff'
	def __init__(self, kernel='rbk', sigma=1, random_feature_method='orthogonal', sample_num=500):
		self.method = random_feature_method
		self.kernel = kernel
		self.m = sample_num		# the number of samples used to approximate k(x,y)
		self.θ = None			# phase shift
		self.σ = sigma
#
	def sample_rademacher_distribution(self, num_of_samples):
		r = np.random.binomial(1, 0.5, int(num_of_samples))
		r = 2*(r - 0.5)
		D = np.diag(r)
		return D
#
	def initialize_random_features(self, X):
		if self.θ == None:
			X0 = X
			N = X.shape[0]
			d = X.shape[1]
			σˉᑊ = 1/self.σ
			ƻπ = 2*np.pi
			m = self.m
			dᑊᐟᒾ = np.sqrt(d)
	
			if self.method == 'orthogonal':	# perform SORF
				#	Here we will pad the data to ensure data dimension is in power of 2 for Hadamard matrix
				ḿ = np.power(2, np.ceil(np.log2(d))) 	# dimension in power of 2 just greater than d
				Δ = int(ḿ - d)
				if Δ != 0:
					pad0 = np.zeros((N, Δ))
					X0 = np.concatenate((X, pad0), axis=1)
				#	each loop creates ḿ out of m, we need to Figure out how many times we need to regenerate 
				#	Note the sample_num may not be equal to the actual random feature width M
				#		m = the number of random features width suggested by the user
				#		ḿ = the number of random features generated by each hadamard matrix 
				#		M = the acutal number of random features we actually use since power of 2 is required by hadamard
				repeat = int(np.ceil(m/ḿ))
				W = np.empty((int(ḿ),0))
				M = int(repeat*ḿ)
				for ɷ in range(repeat):
					#	Generate H and diagonal matrix D sampled from rademacher distribution
					H = 1/np.sqrt(ḿ)*hadamard(ḿ)		# normalized H to be orthonormal
					Dᑊ = self.sample_rademacher_distribution(ḿ)
					Dᒾ = self.sample_rademacher_distribution(ḿ)
					Dᶾ = self.sample_rademacher_distribution(ḿ)
					#
					HDᑊHDᒾHDᶾ = H.dot(Dᑊ).dot(H).dot(Dᒾ).dot(H).dot(Dᶾ)
					Ŵ = (dᑊᐟᒾ*σˉᑊ*HDᑊHDᒾHDᶾ).T			# we did follow the paper of Wx and use XW format
					W = np.concatenate((W, Ŵ), axis=1)
					#
				θ = ƻπ*rand(1, M)
				return [W, X0, M, θ]
				#
			elif self.method == 'rff':
				θ = ƻπ*rand(1, m)
				W = (σˉᑊ)*np.random.randn(d, m)	# random projection matrix W
				return [W, X, m, θ]
#
	def get_feature_map(self, X):
		#	input data conversion
		input_data_type = wuml.wtype(X)
		X = wuml.ensure_data_type(X)

		[W,Ẋ, m, θ] = self.initialize_random_features(X)
		#
		ẊW = Ẋ.dot(W)
		c = np.sqrt(2.0/m)	
		Φ = c*np.cos(ẊW + θ)	

		#	output data conversion
		Φ = wuml.ensure_data_type(Φ, type_name=input_data_type)

		return Φ
#
	def get_kernel(self, X):
		#	input data conversion
		input_data_type = wuml.wtype(X)
		X = wuml.ensure_data_type(X)

		Φ = self.get_feature_map(X)
		K = Φ.dot(Φ.T)
		K = np.maximum(0,K)	# run a relu on the kernel so no negative values
		if self.kernel == 'rbk': K = np.minimum(1,K) # make sure the kernel values doesn't go beyond 1 for gaussian

		#	output data conversion
		K = wuml.ensure_data_type(K, type_name=input_data_type)
		return K

if __name__ == "__main__":
	np.set_printoptions(precision=3)
	np.set_printoptions(linewidth=300)
	np.set_printoptions(suppress=True)
#
	#	On lower dimension data, it seems that RFF is good enough
	X = np.random.randn(7,5)
	σ = np.median(sklearn.metrics.pairwise.pairwise_distances(X))
	γ = 1.0/(2*σ*σ)
#
	K = sklearn.metrics.pairwise.rbf_kernel(X, gamma=γ)
#
	sorf = random_feature(sigma=σ, random_feature_method='orthogonal')
	Kₒ = sorf.get_kernel(X)			# kernel matrix from orthogonal
#
	rff = random_feature(sigma=σ, random_feature_method='rff')
	Kᵣ = rff.get_kernel(X)			# kernel matrix from orthogonal
#
#
	wuml.print_two_matrices_side_by_side(K, Kₒ, title1='Real', title2='Approx by SORF', auto_print=True)
	wuml.print_two_matrices_side_by_side(K, Kᵣ, title1='Real', title2='Approx by RFF', auto_print=True)
#
	jupyter_print('Notice that RFFperforms better on lower dimension datasets')
	ε = mean_absolute_error(K, Kₒ)
	jupyter_print('Mean absolute Error with SORF %.3f'%ε)
	ε = mean_absolute_error(K, Kᵣ)
	jupyter_print('Mean absolute Error with RFF %.3f\n'%ε)
#
	ε = mean_squared_error(K, Kₒ)
	jupyter_print('Mean Squared Error with SORF %.3f'%ε)
	ε = mean_squared_error(K, Kᵣ)
	jupyter_print('Mean Squared Error with SORF %.3f'%ε)
#
#
	#	SORF starts working better when the dimension gets larger, here we use 16 dimension
	X = wuml.csv_load('../dataset/letters.csv', shuffle_samples=True)
	X = X[0:600,:]
#
	σ = np.median(sklearn.metrics.pairwise.pairwise_distances(X))
	γ = 1.0/(2*σ*σ)
#
	K = sklearn.metrics.pairwise.rbf_kernel(X, gamma=γ)
#
	sorf = random_feature(sigma=σ, random_feature_method='orthogonal')
	Kₒ = sorf.get_kernel(X)			# kernel matrix from orthogonal
#
	rff = random_feature(sigma=σ, random_feature_method='rff')
	Kᵣ = rff.get_kernel(X)			# kernel matrix from orthogonal
#
	jupyter_print('Notice that SORF performs better on higher dimension datasets')
	ε = mean_absolute_error(K, Kₒ)
	jupyter_print('\tMean absolute Error with SORF %.3f'%ε)
	ε = mean_absolute_error(K, Kᵣ)
	jupyter_print('\tMean absolute Error with RFF %.3f\n'%ε)
#
	ε = mean_squared_error(K, Kₒ)
	jupyter_print('\tMean Squared Error with SORF %.3f'%ε)
	ε = mean_squared_error(K, Kᵣ)
	jupyter_print('\tMean Squared Error with SORF %.3f'%ε)

