%% Kernel Regression for signals over graphs 
% Arun Venkitaraman 2018-01-01

close all
clear all
clc
tic;
n=8;  % No of training samples to be used for KRG


R=5; % No of folds for R-fold crossvalidation
SNR=0; % Signal to noise ratio of the additive noise

dataset='erdos';

[D,L,alpvec,betvec, Ntrain,Ntest,m,offset,city_ip,city_op]=get_dataset(dataset);
% D: Data matrix
%L: graph-Laplacian
%Ntrain: subset of indices of D corresponding to training set
%Ntest: data indices corresponding to test data
%m: size of graph
%offset: offset of days to be used in the case of temperature data
% city_ip: is the portion of the entrie data used for input (for example
% some cities in the ETEX data)
% city_op: is the portion of the entrie data used for output that lies over associated graph with Laplacian L 

% In order to run synthesized small-world graphs, gsbox toolbox needs to be compiled. 
% The information regarding compiling the associated files are provided in folder 'gspbox'

%% Mask to simulate random large perturbations in training data
Mask=ones(n,m);
for i=1:n
Mask(i,randperm(m,5))=5;
end

%% Mask to simulate random missing samples in training data
Mask=ones(n,m);
for i=1:n
Mask(i,randperm(m,5))=0;
end


%% Cross validation to find parameters
% Choosing subset of n samples from the full training set
ns=length(Ntrain);
ntrain=Ntrain(randperm(ns,n));
ntest=Ntest;


ntest=ntest;
X_train=(D((ntrain)+offset,city_ip));
Y_train=(D((ntrain),city_op))*pinv(eye(m)+0*L);
X_test=(D((ntest)+offset,city_ip));
Y_test=(D((ntest),city_op))*pinv(eye(m)+0*L);

% Generating noisy data
sig_train=1*sqrt((norm(Y_train,'fro')^2/(length(Y_train(:))))*10^(-SNR/10));  % computing the variance for additive noise of given SNR

%% Use for large perturbations
%T_train=Mask.*Y_train;  

%% Use for additive noise
T_train=(Y_train+1*sig_train*randn(size(Y_train))); % 


%% Final computation on test data using the alpha, beta parameters obtained from the crossvalidation step

[all_alpha,all_beta]=return_KRG_xvalparams(X_train, Y_train,T_train,n,L,R,alpvec,betvec);
a1=all_alpha(1);
a2=all_alpha(2);
a3=all_alpha(3);
a4=all_alpha(4);
b2=all_beta(1);
b4=all_beta(2);


%% KRG on test data

ytrain0=Y_train(1:n,:); 
ytrain=T_train(1:n,:);
ytest=Y_test;

for r=1:1
    
    clear train;
 
    Phi_train=X_train;
    Phi_test=X_test;

    ytrain=T_train;
    
    K1=Phi_train*Phi_train';
    k1=(Phi_train*Phi_test')';
    
    K2=pdist2(Phi_train,Phi_train).^2;
    sig_rbf=mean(K2(:));
    K2=exp(-K2/sig_rbf);
    k2=pdist2(Phi_test,Phi_train).^2;
    k2=exp(-k2/sig_rbf);
    
    
    %% Computing the eigenvectors of kernel and graph Laplacian matrices for fast computation of KRG
    [u1,d_thet1]=eig(K1);
    d_thet1=diag(d_thet1);
    [u2,d_thet2]=eig(K2);
    d_thet2=diag(d_thet2);
    [v,d_lam]=eig(L);
    d_lam=diag(d_lam);
    
    
    Z1=kron(v,u1);
    Z2=kron(v,u2);
  
    %% Kernel
    Psi_lin=KerRegGraph_fast((a1(1)),0,ytrain,d_lam,d_thet1,Z1);
    Psi_lin_g=KerRegGraph_fast((a2),(b2(1)),ytrain,d_lam,d_thet1,Z1);
    
    Psi_ker=KerRegGraph_fast((a3(1)),0,ytrain,d_lam,d_thet2,Z2);
    Psi_ker_g=KerRegGraph_fast((a4(1)),(b4(1)),ytrain,d_lam,d_thet2,Z2);
    
    
    % Training prediction
    y_lin_train=K1*Psi_lin;
    y_lin_g_train=K1*Psi_lin_g;
    y_ker_train=K2*Psi_ker;
    y_ker_g_train=K2*Psi_ker_g;
    
    % Test prediction
    y_lin_test=k1*Psi_lin;
    y_lin_g_test=k1*Psi_lin_g;
    y_ker_test=k2*Psi_ker;
    y_ker_g_test=k2*Psi_ker_g;
    
    % MSE
    mse_lin_train_f(r)=(norm(ytrain0(:)-y_lin_train(:),2)^2);
    mse_lin_g_train_f(r)=(norm(ytrain0(:)-y_lin_g_train(:),2)^2);
    e_train_f(r)=norm(ytrain0(:),2)^2;
    mse_lin_test_f(r)=(norm(ytest(:)-y_lin_test(:),2)^2);
    mse_lin_g_test_f(r)=(norm(ytest(:)-y_lin_g_test(:),2)^2);
    e_test_f(r)=norm(ytest(:),2)^2;
     
    mse_ker_train_f(r)=(norm(ytrain0(:)-y_ker_train(:),2)^2);
    mse_ker_g_train_f(r)=(norm(ytrain0(:)-y_ker_g_train(:),2)^2);
    mse_ker_test_f(r)=(norm(ytest(:)-y_ker_test(:),2)^2);
    mse_ker_g_test_f(r)=(norm(ytest(:)-y_ker_g_test(:),2)^2);
end

%% Final MSE values of four approches:
% 1. LR, 2. LRG, 3. KR, and 4. KRG
allmse=10*log10([mean(mse_lin_test_f)/mean(e_test_f) mean(mse_lin_g_test_f)/mean(e_test_f) mean(mse_ker_test_f)/mean(e_test_f) mean(mse_ker_g_test_f)/mean(e_test_f) ])

 
%% Plotting an example realization
figure, plot(ytrain0(:),'k','Linewidth',2), hold on,plot(ytrain(:),'g'), hold on
hold on, plot(y_lin_train(:),'b'),plot(y_lin_g_train(:),'r'),
plot(y_ker_train(:),'bO-'),plot(y_ker_g_train(:),'rO-')

 figure, plot(ytest(:),'k','Linewidth',2),  hold on,
plot(y_lin_test(:),'b'),plot(y_lin_g_test(:),'r'),
plot(ytest(:),'k','Linewidth',2),
plot(y_ker_test(:),'bO-'),plot(y_ker_g_test(:),'rO-')
% 
