function MLmethods_sEMG_exp2
% Test Gaussian process regression and ridge regression on sEMG data of 
% experiment 2. Six movements are performed for 3 repetetions by 9 
% non-amputated participants. For GPR, RBF kernels with 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/EMGresults%d.txt (RMSE values) and
% predictions are registered in EMGpredictions%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/EMGresults%d.txt',subjectId);
	predstr = sprintf('results/EMGpredictions%d.mat',subjectId);
	
	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));
	
	% Covariance features for EMG data
	% Sliding window for spatial covariance computation
	timeWindow = 0.4; % s
	freq = Fs; % Hz
	numberOfSamples = round(timeWindow*freq);
	
	N = size(emg,1);
	D = size(emg,2);
	emgCov = zeros(N,D,D);
	for i = numberOfSamples:size(emg,1) % the first samples are zero, but there are not used
	emgCov(i,:,:) = reshape( cov(emg(i-numberOfSamples+1:i,:))+1E-12*eye(D), 1,D,D);
	end
	
	% Extract data corresponding to activation index 1 (measurement time) 
	% for rest and the six movements -> training data
	activeIdx = find(stim(:,12)==1);

	emgTrain = emg(activeIdx,:);
	emgCovTrain = emgCov(activeIdx,:,:);
	yTrain = stim(activeIdx,1:9);
	
	% 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);
	
	emgTest = emg(activeIdx,:);
	emgCovTest = emgCov(activeIdx,:,:);
	yTest = stim(activeIdx,:);
	
	disp(' -> Done!');
	
	%% Gaussian process SPD distance
	disp('Gaussian process with SPD distance');
	
	fprintf('Compute training kernel... ');
	tic
	kerdistTrTr = covmatCovDist(emgCovTrain,emgCovTrain,0);
	fprintf('Done in %.2fs.\n',toc);
	
	fprintf('Cross validation to find best parameters...');
	betaParam_SPD = [logspace(-2,2,10)];
	predictor_SPD = @(XTr, YTr, XTe, idxTr, idxTe, param) GPR_cv(XTr, YTr, XTe, idxTr, idxTe, param, kerdistTrTr);
	[ ~, teRMSE ] = crossvalidation( emgCovTrain, yTrain, nbIterCV, testProp, predictor_SPD, betaParam_SPD);
	
	meanRMSE_SPD = mean(teRMSE,ndims(teRMSE));
	[minTeRMSE_SPD, idTeRMSE_SPD] = min(meanRMSE_SPD(:));
	
	betaSPD = betaParam_SPD(idTeRMSE_SPD);
	
	fprintf('Min. training RMSE GP SPD: %f for beta1 %.2f\n', minTeRMSE_SPD, betaSPD);
	
	fprintf('Testing on intermediate data... ')
	tic
	kerdistTeTr = covmatCovDist(emgCovTest,emgCovTrain,0);
	yhat_emg_SPD = GPR(emgCovTrain, yTrain, emgCovTest, betaSPD, kerdistTrTr, kerdistTeTr);
	fprintf('Done in %.2fs.\n',toc);
	
	%% Gaussian process log SPD distance
	disp('Gaussian process with log SPD distance');
	
	fprintf('Compute training kernel... ');
	tic
	kerdistTrTr = covmatCovDist(emgCovTrain,emgCovTrain,1);
	fprintf('Done in %.2fs.\n',toc);
	
	fprintf('Cross validation to find best parameters...');
	betaParam_logSPD = [logspace(-2,2,10)];
	predictor_logSPD = @(XTr, YTr, XTe, idxTr, idxTe, param) GPR_cv(XTr, YTr, XTe, idxTr, idxTe, param, kerdistTrTr);
	[ ~, teRMSE ] = crossvalidation( emgCovTrain, yTrain, nbIterCV, testProp, predictor_logSPD, betaParam_logSPD);
	
	meanRMSE_logSPD = mean(teRMSE,ndims(teRMSE));
	[minTeRMSE_logSPD, idTeRMSE_logSPD] = min(meanRMSE_logSPD(:));
	
	betaLogSPD = betaParam_logSPD(idTeRMSE_logSPD);
	
	fprintf('Min. training RMSE GP log SPD: %f for beta1 %f\n', minTeRMSE_logSPD, betaLogSPD);
	
	fprintf('Testing on intermediate data... ')
	tic
	kerdistTeTr = covmatCovDist(emgCovTest,emgCovTrain,1);
	yhat_emg_logSPD = GPR(emgCovTrain, yTrain, emgCovTest, betaLogSPD, kerdistTrTr, kerdistTeTr);
	fprintf('Done in %.2fs.\n',toc);
	
	%% Vector Gaussian process
	disp('Gaussian process with Eucl distance');
	
	fprintf('Compute training kernel... ');
	tic
	kerdistTrTr = compute_Eucldist(emgTrain,emgTrain);
	fprintf('Done in %.2fs.\n',toc);
	
	fprintf('Cross validation to find best parameters...');
	betaParam_Eucl = [0.1:0.2:7];
	predictor_Eucl = @(XTr, YTr, XTe, idxTr, idxTe, param) GPR_cv(XTr, YTr, XTe, idxTr, idxTe, param, kerdistTrTr);
	[ ~, teRMSE ] = crossvalidation( emgTrain, yTrain, nbIterCV, testProp, predictor_Eucl, betaParam_Eucl);
	
	teRMSE_Eucl = mean(teRMSE,ndims(teRMSE));
	[minTeRMSE_Eucl, idTeRMSE_Eucl] = min(teRMSE_Eucl(:));
	
	betaGP = betaParam_Eucl(idTeRMSE_Eucl);

	fprintf('Min. training RMSE GP Eucl: %f for beta %.2f\n', minTeRMSE_Eucl, betaGP);
	
	fprintf('Testing on intermediate data... ')
	tic
	kerdistTeTr = compute_Eucldist(emgCovTest,emgCovTrain);
	yhat_emg_Eucl = GPR(emgTrain, yTrain, emgTest, betaGP, kerdistTrTr, kerdistTeTr);
	fprintf('Done in %.2fs.\n',toc);
	
	%% Ridge Regression
	disp('Ridge regression');
	lambdaParam = logspace(-4,-1,20);
	predictorRR = @(XTr, YTr, XTe, idxTr, idxTe, param) RR(XTr, YTr, XTe, idxTr, idxTe, param);
	[ ~, teRMSE ] = crossvalidation( emgTrain, yTrain, nbIterCV, testProp, predictorRR, lambdaParam);
	
	teRMSE_RR = mean(teRMSE,ndims(teRMSE));
	[minTeRMSE_RR, idTeRMSE_RR] = min(teRMSE_RR(:));
	
	lambda = lambdaParam(idTeRMSE_RR);
	
	fprintf('Min. training RMSE RR %f for lambda %f\n', minTeRMSE_RR, lambda);
	
	fprintf('Testing on intermediate data... ')
	tic
	yhat_emg_rr = RR(emgTrain, yTrain, emgTest, 1, 1, lambda);
	fprintf('Done in %.2fs.\n',toc);
	
	%% Save predictions
	save(predstr,'emgCovTrain','emgCovTest','emgTrain','emgTest','yTrain','yTest', 'yhat_emg_SPD',...
		'yhat_emg_logSPD','yhat_emg_Eucl','yhat_emg_rr');

	
	%% Compute errors
	err_SPD = computeRMSE(yTest(:,1:9),yhat_emg_SPD);

	err_logSPD = computeRMSE(yTest(:,1:9),yhat_emg_logSPD);
	
	err_Eucl = computeRMSE(yTest(:,1:9),yhat_emg_Eucl);

	err_RR = computeRMSE(yTest(:,1:9),yhat_emg_rr);
	%% Write in txt file
	fileID = fopen(filestr,'w');
	
	fprintf(fileID,'SPD %2.2f for beta1 %2f\n',err_SPD, betaSPD);
	fprintf(fileID,'logSPD %2.2f for beta1 %2f\n',err_logSPD, betaLogSPD);
	fprintf(fileID,'Eucl %2.2f for beta %2f\n', err_Eucl,betaGP);
	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 = 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
% YTr		Output
% XTe		New input
% idxTr		Training indexes
% idxTe		Testing indexes
% beta		GP parameter
% kerdist	Kernel distance computed for all inputs
% Yhat		Predicted output

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 = 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 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 = covmatCovDist(X1,X2,logDist)
%tensorCovDist Compute the covariance or log covariance (if logDist = 1)
%distance covDistTens between X1 and X2 (data are SPD matrices)

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

M1 = size(X1,2);
M2 = size(X1,3);

M3 = size(X2,2);
M4 = size(X2,3);

if M1~=M2 || M3~=M4 || M1~=M3
	error('Inputs must have the same size and be covariance matrices.');
end
M = M1;

diagRegFact = 1E-6;

covDistTens = zeros(N1,N2);

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
		S1 = permute(X1(n1,:,:),[2,3,1]) + diagRegFact.*eye(M);
		S2 = permute(X2(n2,:,:),[2,3,1]) + diagRegFact.*eye(M);
			if(logDist)
				covDistTens(n1,n2) = computeLogCovDist(S1,S2);
			else
				covDistTens(n1,n2) = computeCovDist(S1,S2);
			end
			if symKernel
				covDistTens(n2,n1) = covDistTens(n1,n2);
			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 non-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