function MLmethods_TMG_exp2
% Test Gaussian process regression and ridge regression on TMG data of 
% experiment 2. Six movements are performed for 3 repetetions by 9 
% non-amputated participants. For GPR, RBF kernels with KL divergence,
% SPD, logSPD and Euclidean distance are tested.
% Training data are taken as zero and full activations of all movements,
% while testing data are taken as transition states corresponding to the
% whole range of intermediate activation.
% Results are registered in results/TMGresults%d.txt (RMSE values) and
% predictions are registered in TMGpredictions%d.mat wit d=1:9.
%
%
%
% Writing code takes time. Polishing it and making it available to others takes longer! 
% If some parts of the code were useful for your research of for a better understanding 
% of the algorithms, please cite the related publications.
%
% @article{Jaquier17MDPI,
%   author="Jaquier, N. and Connan, M. and Castellini, C. and Calinon, S.",
%   title="Combining electro- and tactile myography to improve hand and 
%		wrist activity detection in prostheses",
%   year="2017",
%	journal = "Technologies",
%	volume = "",
%	number = "",
%	article number = "",
% }
% 
% Copyright (c) 2017 Idiap Research Institute, http://idiap.ch/
% Written by Noémie Jaquier
% 
% This code is free software: you can redistribute it and/or modify
% it under the terms of the GNU General Public License version 3 as
% published by the Free Software Foundation.
% 
% This code is distributed in the hope that it will be useful,
% but WITHOUT ANY WARRANTY; without even the implied warranty of
% MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
% GNU General Public License for more details.
% 
% You should have received a copy of the GNU General Public License
% along with this code. If not, see <http://www.gnu.org/licenses/>.

for subjectId = 1:9
	%% Parameters
	
	filestr = sprintf('results/TMGresults%d.txt',subjectId);
	predstr = sprintf('results/TMGpredictions%d.mat',subjectId-2);
	
	nbIterCV = 10;
	testProp = 0.6;

	%% Load and preprocess training data
	fprintf('Loading and preprocessing data %d...',subjectId);
	load(sprintf('data/exp2/data_S%d.mat',subjectId));
	
	% Extract data corresponding to activation index 1 (measurement time) 
	% for rest and the six movements -> training data
	activeIdx = find(stim(:,12)==1);
	
	tmgTrainVec = brac(activeIdx,:);
	yTrain = stim(activeIdx,1:9);
	
	% Training input data in matrix form (tactile bracelet of 8 x 40 sensors)
	N = size(tmgTrainVec,1);
	tmgTrain = zeros(N,8,40);
	for i=1:size(tmgTrainVec,1)
		tmgTrain(i,:,:) = reshape(tmgTrainVec(i,:),8,40);
	end
	
	% Extract data corresponding to intermediate results with activation
	% index 0 -> testing data
	activeIdx = find(~mod(stim(:,10),1) & ~mod(stim(:,11),1) & stim(:,12)==0);
	
	tmgTestVec = brac(activeIdx,:);
	yTest = stim(activeIdx,:);
	
	% Testing input data in matrix form
	N = size(tmgTestVec,1);
	tmgTest = zeros(N,8,40);
	for i=1:size(tmgTestVec,1)
		tmgTest(i,:,:) = reshape(tmgTestVec(i,:),8,40);
	end
	
	disp(' -> Done!');
	
	%% Matrix Gaussian process: KL divergence
	disp('Gaussian process with KL divergence');
	
	fprintf('Compute training kernel... ');
	tic
	kerdistTrTr = tensorKLdivGauss(tmgTrain,tmgTrain);
	fprintf('Done in %.2fs.\n',toc);
	
	fprintf('Cross validation to find best parameters...');
	betaParam_KL{1} = logspace(-1,3,15);
	betaParam_KL{2} = logspace(-1,3,15);
	predictor_KL = @(XTr, YTr, XTe, idxTr, idxTe, param) tensorGPR_cv(XTr, YTr, XTe, idxTr, idxTe, param, kerdistTrTr);
	[ ~, teRMSE ] = crossvalidation( tmgTrain, yTrain, nbIterCV, testProp, predictor_KL, betaParam_KL);
	
	meanRMSE_KL = mean(teRMSE,ndims(teRMSE));	
	[minTeRMSE_KL, minId] = min(meanRMSE_KL(:));
	[idTeRMSE_KL(1), idTeRMSE_KL(2)] = ind2sub(size(meanRMSE_KL),minId);
	
	betaKL(1,1,1) = betaParam_KL{1}(idTeRMSE_KL(1));
	betaKL(1,1,2) = betaParam_KL{2}(idTeRMSE_KL(2));
	
	fprintf('Min. training RMSE GP KL: %f for beta1 %.2f and beta2 %.2f\n', minTeRMSE_KL, betaKL(1,1,1), betaKL(1,1,2));
	
	fprintf('Testing on intermediate data... ')
	tic
	kerdistTeTr = tensorKLdivGauss(tmgTest,tmgTrain);
	yhat_tmg_KL = tensorGPR(tmgTrain, yTrain, tmgTest, betaKL, kerdistTrTr, kerdistTeTr);
	fprintf('Done in %.2fs.\n',toc);
	
	%% Matrix Gaussian process SPD distance
	disp('Gaussian process with SPD distance');
	
	fprintf('Compute training kernel... ');
	tic
	kerdistTrTr = tensorCovDist(tmgTrain,tmgTrain,0);
	fprintf('Done in %.2fs.\n',toc);
	
	fprintf('Cross validation to find best parameters...');
	betaParam_SPD{1} = logspace(-1,1.5,15);
	betaParam_SPD{2} = logspace(-1,1.5,15);
	predictor_SPD = @(XTr, YTr, XTe, idxTr, idxTe, param) tensorGPR_cv(XTr, YTr, XTe, idxTr, idxTe, param, kerdistTrTr);
	[ ~, teRMSE ] = crossvalidation( tmgTrain, yTrain, nbIterCV, testProp, predictor_SPD, betaParam_SPD);
	
	meanRMSE_SPD = mean(teRMSE,ndims(teRMSE));
	[minTeRMSE_SPD, minId] = min(meanRMSE_SPD(:));
	[idTeRMSE_SPD(1), idTeRMSE_SPD(2)] = ind2sub(size(meanRMSE_SPD),minId);
	
	betaSPD(1,1,1) = betaParam_SPD{1}(idTeRMSE_SPD(1));
	betaSPD(1,1,2) = betaParam_SPD{2}(idTeRMSE_SPD(2));
	
	fprintf('Min. training RMSE GP SPD: %f for beta1 %.2f and beta2 %.2f\n', minTeRMSE_SPD, betaSPD(1,1,1), betaSPD(1,1,2));
	
	fprintf('Testing on intermediate data... ')
	tic
	kerdistTeTr = tensorCovDist(tmgTest,tmgTrain,0);
	yhat_tmg_SPD = tensorGPR(tmgTrain, yTrain, tmgTest, betaSPD, kerdistTrTr, kerdistTeTr);
	fprintf('Done in %.2fs.\n',toc);
	
	%% Matrix Gaussian process log SPD distance
	disp('Gaussian process with log SPD distance');
	
	fprintf('Compute training kernel... ');
	tic
	kerdistTrTr = tensorCovDist(tmgTrain,tmgTrain,1);
	fprintf('Done in %.2fs.\n',toc);
	
	fprintf('Cross validation to find best parameters...');
	betaParam_logSPD{1} = logspace(-1,1.6,15);
	betaParam_logSPD{2} = logspace(-1,1.6,15);
	predictor_logSPD = @(XTr, YTr, XTe, idxTr, idxTe, param) tensorGPR_cv(XTr, YTr, XTe, idxTr, idxTe, param, kerdistTrTr);
	[ ~, teRMSE ] = crossvalidation( tmgTrain, yTrain, nbIterCV, testProp, predictor_logSPD, betaParam_logSPD);
	
	meanRMSE_logSPD = mean(teRMSE,ndims(teRMSE));
	[minTeRMSE_logSPD, minId] = min(meanRMSE_logSPD(:));
	[idTeRMSE_logSPD(1), idTeRMSE_logSPD(2)] = ind2sub(size(meanRMSE_logSPD),minId);
	
	betaLogSPD(1,1,1) = betaParam_logSPD{1}(idTeRMSE_logSPD(1));
	betaLogSPD(1,1,2) = betaParam_logSPD{2}(idTeRMSE_logSPD(2));
	
	fprintf('Min. training RMSE GP log SPD: %f for beta1 %f and beta2 %f\n', minTeRMSE_logSPD, betaLogSPD(1,1,1), betaLogSPD(1,1,2));
	
	fprintf('Testing on intermediate data... ')
	tic
	kerdistTeTr = tensorCovDist(tmgTest,tmgTrain,1);
	yhat_tmg_logSPD = tensorGPR(tmgTrain, yTrain, tmgTest, betaLogSPD, kerdistTrTr, kerdistTeTr);
	fprintf('Done in %.2fs.\n',toc);
	
	%% Vector Gaussian process Euclidean distance
	disp('Gaussian process with Eucl distance');
	
	fprintf('Compute training kernel... ');
	tic
	kerdistTrTr = compute_Eucldist(tmgTrain,tmgTrain);
	fprintf('Done in %.2fs.\n',toc);
	
	fprintf('Cross validation to find best parameters...');
	betaParam_Eucl = logspace(-1,1,10);
	predictorEucl = @(XTr, YTr, XTe, idxTr, idxTe, param) GPR_cv(XTr, YTr, XTe, idxTr, idxTe, param, kerdistTrTr);
	[ ~, teRMSE ] = crossvalidation( tmgTrain, yTrain, nbIterCV, testProp, predictorEucl, betaParam_Eucl);
	
	meanRMSE_Eucl = mean(teRMSE,ndims(teRMSE));
	[minTeRMSE_Eucl, idTeRMSE_Eucl] = min(meanRMSE_Eucl(:));
	
	betaEucl = betaParam_Eucl(idTeRMSE_Eucl);

	fprintf('Min. training RMSE GP Eucl: %f for beta %.2f\n', minTeRMSE_Eucl, betaEucl);
	
	fprintf('Testing on intermediate data... ')
	tic
	kerdistTeTr = compute_Eucldist(tmgTest,tmgTrain);
	yhat_tmg_Eucl = GPR(tmgTrainVec, yTrain, tmgTestVec, betaEucl, kerdistTrTr, kerdistTeTr);
	fprintf('Done in %.2fs.\n',toc);
	
	%% Ridge regression
	disp('Ridge regression');
	lambdaParam = logspace(-4,-1,10);
	predictorRR = @(XTr, YTr, XTe, idxTr, idxTe, param) RR(XTr, YTr, XTe, idxTr, idxTe, param);
	[ ~, teRMSE ] = crossvalidation( tmgTrainVec, yTrain, nbIterCV, testProp, predictorRR, lambdaParam);
	
	meanRMSE_RR = mean(teRMSE,ndims(teRMSE));
	[minTeRMSE_RR, idTeRMSE_RR] = min(meanRMSE_RR(:));
	
	lambda = lambdaParam(idTeRMSE_RR);
	
	fprintf('Min RMSE RR %f for lambda %f\n', minTeRMSE_RR, lambda);
	
	fprintf('Testing on intermediate data... ')
	tic
	yhat_tmg_rr = RR(tmgTrainVec, yTrain, tmgTestVec, 1, 1, lambda);
	fprintf('Done in %.2fs.\n',toc);
	
	%% Save predictions
	save(predstr,'tmgTrain','tmgTest','yTrain','yTest', 'yhat_tmg_KL','yhat_tmg_SPD',...
		'yhat_tmg_logSPD', 'yhat_tmg_Eucl','yhat_tmg_rr');

	
	%% Compute errors
	err_KL = computeRMSE(yTest(:,1:9),yhat_tmg_KL);

	err_SPD = computeRMSE(yTest(:,1:9),yhat_tmg_SPD);

	err_logSPD = computeRMSE(yTest(:,1:9),yhat_tmg_logSPD);

	err_GP = computeRMSE(yTest(:,1:9),yhat_tmg_Eucl);

	err_RR = computeRMSE(yTest(:,1:9),yhat_tmg_rr);
	%% Write in txt file
	fileID = fopen(filestr,'w');
	
	fprintf(fileID,'KL %2.2f for beta1 %2f and beta2 %2f\n',err_KL, betaKL(1,1,1), betaKL(1,1,2));
	fprintf(fileID,'SPD %2.2f for beta1 %2f and beta2 %2f\n',err_SPD, betaSPD(1,1,1), betaSPD(1,1,2));
	fprintf(fileID,'logSPD %2.2f for beta1 %2f and beta2 %2f\n',err_logSPD, betaLogSPD(1,1,1), betaLogSPD(1,1,2));
	fprintf(fileID,'Eucl %2.2f for beta %2f\n', err_GP,betaEucl);
	fprintf(fileID,'RR %2.2f for lambda %f\n',err_RR,lambda);
	
	fclose(fileID);
end

end

function [ trRMSE, teRMSE ] = crossvalidation( X, Y, iterationsNb, testRatio, predict, param)
%crossvalidation Perform crossvalidation to return the train and test RMSE
%for the method predict and all values of parameters param (up to 2 dim.)
%
% X             Input N x D1 x .. x D_I 
% y             Output N x d
% iterationsNb  Number of random shuffles
% testRatio		Proportion of test data
% predict		Machine learning method used to find Y from X
% param			Values of the parameters of the method predict
% trRMSE		Training RMSE for each parameter and iteration
% teRMSE		Testing RMSE for each parameter and iteration

randomStream = RandStream.getGlobalStream; % using globalStream

trRMSE = [];
teRMSE = [];

getDivision = divideDataSet( X, Y, iterationsNb, testRatio, randomStream );

for iteration = 1:iterationsNb
	if ~mod(iteration,10)
		fprintf('%d out of %d...\n',iteration, iterationsNb);
	end
	% Separe X and y in train and test parts
	[ XTr, XTe, YTr, YTe, idxTr, idxTe ] = getDivision(iteration);
	
	if ~iscell(param)
		for p = 1:length(param)
			% Make prediction for test data
			yTrPred = predict( XTr, YTr, XTr, idxTr, idxTr, param(p) );
			yTePred = predict( XTr, YTr, XTe, idxTr, idxTe, param(p) );
			
			% Compute RMSE
			trRMSE(p,iteration) = computeRMSE(YTr, yTrPred); 
			teRMSE(p,iteration) = computeRMSE(YTe, yTePred);
		end
	else
		for p1 = 1:length(param{1})
			for p2 = 1:length(param{2})
				p(1,1,1) = param{1}(p1);
				p(1,1,2) = param{2}(p2);
				% Make prediction for test data
				yTrPred = predict( XTr, YTr, XTr, idxTr, idxTr, p );
				yTePred = predict( XTr, YTr, XTe, idxTr, idxTe, p );
				
				% Compute RMSE
				trRMSE(p1,p2,iteration) = computeRMSE(YTr, yTrPred);
				teRMSE(p1,p2,iteration) = computeRMSE(YTe, yTePred);
			end
		end
	end
	
end
end

function [ getDivision ] = divideDataSet( X, Y, iterationsNb, testRatio, randomStream )
%divideDataSet Divide the dataset in train and test parts
%
% X             Input N x D1 x .. x D_I 
% y             Output N x d
% iterationsNb  Number of random shuffles
% randomStream  Random stream used to divide the dataset
% getDivision   Function from iteration number to [XTr, XTe, yTr, yTe]

dataSize = size(Y, 1);

trainSize = floor(dataSize*(1-testRatio));
idxPerIt = zeros(iterationsNb, dataSize);

for iteration = 1:iterationsNb
    idxPerIt(iteration, :) = randomStream.randperm(dataSize);
end


    function [XTr, XTe, YTr, YTe, idxTr, idxTe] = i_getDivision(iterationNumber)
        % Get first part of shuffled samples for train, the rest for test
        idxTe = idxPerIt(iterationNumber, trainSize+1:end);
        idxTr = idxPerIt(iterationNumber, 1:trainSize);
        
        dim = ndims(X);
        indsTe = repmat({':'},1,dim);
        indsTe{1} = idxTe;
        YTe = Y(idxTe, :);
        XTe = X(indsTe{:});
        
        indsTr = repmat({':'},1,dim);
        indsTr{1} = idxTr;
        YTr = Y(idxTr, :);
        XTr = X(indsTr{:});
        
    end

getDivision = @i_getDivision;

end

function [ error ] = computeRMSE( Y, Yhat )
%computeRMSE Compute the root mean square error
%
% Y       Expected output N x d
% YHat    Predicted output N x d

N = length(Y);

diff_square = sum((Y-Yhat).^2,2);

error = sum(diff_square);

error = sqrt(error./N);
end

function Yhat = tensorGPR_cv(XTr, YTr, XTe, idxTr, idxTe, beta, kerdist)
%tensorGPR_cv Predict a new output with tensor Gaussian process regression 
%(used for cross-validation only)
%
% XTr		Input N x D1 x .. x D_I
% YTr		Output N x d
% XTe		New input Nte x D1 x .. x D_I
% idxTr		Training indexes
% idxTe		Testing indexes
% beta		TGP parameter
% kerdist	Kernel distance computed for all inputs (N+Nte) x (N+Nte) x I
% Yhat		Predicted output Nte x d

diagRegFact = 1E-4;

M = ndims(kerdist);

kerdistTrTr = kerdist(idxTr,idxTr,:);
kerdistTeTr = kerdist(idxTe,idxTr,:);

kerdistTrTr = kerdistTrTr./repmat((2.*beta.^2).*ones(1,1,M+1-ndims(beta)),size(XTr,1),size(XTr,1),1);
KTrTr = prod(exp(-kerdistTrTr),3);

kerdistTeTr = kerdistTeTr./repmat((2.*beta.^2).*ones(1,1,M+1-ndims(beta)),size(XTe,1),size(XTr,1),1);
KTeTr = prod(exp(-kerdistTeTr),3);

Yhat = KTeTr*((KTrTr+diagRegFact.*eye(size(XTr,1)))\YTr);
end

function Yhat = GPR_cv(XTr, YTr, XTe, idxTr, idxTe, beta, kerdist)
%GPR_cv Predict a new output with Gaussian process regression (used for 
%cross-validation only)
%
% XTr		Input N x D
% YTr		Output N x d
% XTe		New input Nte x D
% idxTr		Training indexes
% idxTe		Testing indexes
% beta		GP parameter
% kerdist	Kernel distance computed for all inputs (N+Nte) x (N+Nte) 
% Yhat		Predicted output Nte x d

diagRegFact = 1E-4;

kerdistTrTr = kerdist(idxTr,idxTr,:);
kerdistTeTr = kerdist(idxTe,idxTr,:);

KTrTr = exp(-kerdistTrTr./(2.*beta.^2));
KTeTr = exp(-kerdistTeTr./(2.*beta.^2));

Yhat = KTeTr*((KTrTr+diagRegFact.*eye(size(XTr,1)))\YTr);
end

function Yhat = tensorGPR(XTr, YTr, XTe, beta, kerdistTrTr, kerdistTeTr)
%tensorGPR Predict a new output with tensor Gaussian process regression 
%
% XTr			Input N x D1 x .. x D_I
% YTr			Output N x d
% XTe			New input Nte x D1 x .. x D_I
% beta			TGP parameter
% kerdistTrTr	Kernel distance between input data N x N x I
% kerdistTeTr	Kernel distance between new input and input data Nte x N x I
% Yhat			Predicted output Nte x d

diagRegFact = 1E-4;

M = ndims(kerdistTrTr);

kerdistTrTr = kerdistTrTr./repmat((2.*beta.^2).*ones(1,1,M+1-ndims(beta)),size(XTr,1),size(XTr,1),1);
KTrTr = prod(exp(-kerdistTrTr),3);

kerdistTeTr = kerdistTeTr./repmat((2.*beta.^2).*ones(1,1,M+1-ndims(beta)),size(XTe,1),size(XTr,1),1);
KTeTr = prod(exp(-kerdistTeTr),3);

Yhat = KTeTr*((KTrTr+diagRegFact.*eye(size(XTr,1)))\YTr);
end

function Yhat = GPR(XTr, YTr, XTe, beta, kerdistTrTr, kerdistTeTr)
%GPR Predict a new output with Gaussian process regression 
%
% XTr			Input N x D
% YTr			Output N x d
% XTe			New input Nte x D
% beta			TGP parameter
% kerdistTrTr	Kernel distance between input data N x N 
% kerdistTeTr	Kernel distance between new input and input data Nte x N
% Yhat			Predicted output Nte x d

diagRegFact = 1E-4;

KTrTr = exp(-kerdistTrTr./(2.*beta.^2));
KTeTr = exp(-kerdistTeTr./(2.*beta.^2));

Yhat = KTeTr*((KTrTr+diagRegFact.*eye(size(XTr,1)))\YTr);
end

function Yhat = RR(XTr, YTr, XTe, idxTr, idxTe, lambda)
%RR Predict a new output with Ridge regression
%
% XTr		Input N x D
% YTr		Output N x d
% XTe		New input Nte x D
% idxTr		Training indexes (for cross-validation)
% idxTe		Testing indexes (for cross-validation)
% lambda	RR parameter
% Yhat		Predicted output Nte x d

% Define lambda matrix
lambdaMatrix = lambda*eye(size(XTr,2));

% Compute beta
beta = (XTr'*XTr + lambdaMatrix)\XTr'*YTr;

% Prediction
Yhat = XTe*beta;
end

function KLdivTens = tensorKLdivGauss(X1,X2)
%tensorKLdivGauss Computes the KL divergence KLdivTens between X1 and X2 by
%treating each matrix data as a generative model

N1 = size(X1,1);
N2 = size(X2,1);

M1 = ndims(X1);
M2 = ndims(X2);
if M1~=M2
	error('Inputs must have the same number of dimensions');
end
M = M1;

% Compute mean and covariance for each data
diagRegFact = 1E-4;
m1 = cell(N1,M-1);
S1 = cell(N1,M-1);
m2 = cell(N2,M-1);
S2 = cell(N2,M-1);

for n = 1:N1
	for m = 2:M
		mat = tens2mat(X1(n,:,:),m,[1:m-1,m+1:M]);
		m1{n,m-1} = mean(mat,2);
		S1{n,m-1} = cov(mat')+diagRegFact.*eye(size(m1{n,m-1},1));
	end
end

for n = 1:N2
	for m = 2:M
		mat = tens2mat(X2(n,:,:),m,[1:m-1,m+1:M]);
		m2{n,m-1} = mean(mat,2);
		S2{n,m-1} = cov(mat')+diagRegFact.*eye(size(m2{n,m-1},1));
	end
end

KLdivTens = zeros(N1,N2,M-1);

symKernel = false;
n2init = 1;
% Check if the kernel is symmetric (if yes, we compute the similar elements
% only once.
if N1 == N2
	if X1 == X2
		symKernel = true;
	end
end

% Compute the kernel
for n1 = 1:N1
	if ~mod(n1,100)
		fprintf('.');
	end
	if symKernel
		n2init = n1;
	end
	for n2 = n2init:N2
		for m = 2:M
			KLdivTens(n1,n2,m-1) = computeKLdivGauss(m1{n1,m-1},S1{n1,m-1},m2{n2,m-1},S2{n2,m-1});
			if symKernel
				KLdivTens(n2,n1,m-1) = KLdivTens(n1,n2,m-1);
			end
		end
	end
end
end

function KLdiv = computeKLdivGauss(m1,S1,m2,S2)
%KLdiv Compute the Kullback-Leibler divergence KLdiv between distributions
%N(m1,S1) and N(m2,S2)

k = size(m1,1);
KLdiv = 0.5*(trace(S1\S2)+((m1-m2)'/S1)*(m1-m2)-k+log(det(S1)/det(S2)));
end

function kerdist = compute_Eucldist(X1,X2)
%kerdist Compute the Euclidean distance kerdist between X1 and X2

N1 = size(X1,1);
N2 = size(X2,1);

symKernel = false;
n2init = 1;
% Check if the kernel is symmetric (if yes, we compute the similar elements
% only once.
if N1 == N2
	if X1 == X2
		symKernel = true;
	end
end

% Compute the kernel
kerdist = zeros(N1,N2);
for n1 = 1:N1
	if ~mod(n1,100)
		fprintf('.');
	end
	if symKernel
		n2init = n1;
	end
	for n2 = n2init:N2
		kerdist(n1,n2) = norm(X1(n1,:)-X2(n2,:),'fro');
		if symKernel
			kerdist(n2,n1) = kerdist(n1,n2);
		end
	end
end

end

function covDistTens = tensorCovDist(X1,X2,logDist)
%tensorCovDist Compute the covariance or log covariance (if logDist = 1)
%distance covDistTens between X1 and X2 (data are not SPD matrices)

N1 = size(X1,1);
N2 = size(X2,1);

M1 = ndims(X1);
M2 = ndims(X2);
if M1~=M2
	error('Inputs must have the same number of dimensions');
end
M = M1;

diagRegFact = 1E-4;

S1 = cell(N1,M-1);
S2 = cell(N2,M-1);

for n = 1:N1
	for m = 2:M
		mat = tens2mat(X1(n,:,:),m,[1:m-1,m+1:M]);
		S1{n,m-1} = cov(mat')+diagRegFact.*eye(size(mat,1));
	end
end

for n = 1:N2
	for m = 2:M
		mat = tens2mat(X2(n,:,:),m,[1:m-1,m+1:M]);
		S2{n,m-1} = cov(mat')+diagRegFact.*eye(size(mat,1));
	end
end

covDistTens = zeros(N1,N2,M-1);

symKernel = false;
n2init = 1;
% Check if the kernel is symmetric (if yes, we compute the similar elements
% only once.
if N1 == N2
	if X1 == X2
		symKernel = true;
	end
end

% Compute the kernel
for n1 = 1:N1
	if ~mod(n1,100)
		fprintf('.');
	end
	if symKernel
		n2init = n1;
	end
	for n2 = n2init:N2
		for m = 2:M
			if(logDist)
				covDistTens(n1,n2,m-1) = computeLogCovDist(S1{n1,m-1},S2{n2,m-1});
			else
				covDistTens(n1,n2,m-1) = computeCovDist(S1{n1,m-1},S2{n2,m-1});
			end
			if symKernel
				covDistTens(n2,n1,m-1) = covDistTens(n1,n2,m-1);
			end
		end
	end
end
end

function covDist = computeCovDist(S1,S2)
%computeCovDist Compute the non-geodesic SPD distance between S1 and S2

covDist = norm(S1-S2,'fro');
end

function covDist = computeLogCovDist(S1,S2)
%computeCovDist Compute the geodesic SPD distance between S1 and S2

covDist = norm(logm(S1)-logm(S2),'fro');
end

function M = tens2mat(T, rows, cols)
%tens2mat Matricisation of a tensor T such that dimensions contained in
%rows and cols become the rows, respectively columns of the new matrix M

if nargin <=2
	cols = [];
end

sizeT = size(T);
N = ndims(T);

if isempty(rows)
	rows = 1:N;
	rows(cols) = [];
end
if isempty(cols)
	cols = 1:N;
	cols(rows) = [];
end

if any(rows(:)' ~= 1:length(rows)) || any(cols(:)' ~= length(rows)+(1:length(cols)))
	T = permute(T,[rows(:)' cols(:)']);
end

M = reshape(T,prod(sizeT(rows)), prod(sizeT(cols)));

end