%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% MATLAB examples of HybridNeuro WEBINAR on Validation of results: statistical models and MU identification accuracy
%  
% Topics covered:
% - Accuracy of MU identification from HDEMG
% - Regression analysis with Linear Mixed Models 
% - Bayesian linear regression models
%
% EXPERIMENTAL HDEMG SIGNALS
%
% Authors: Aleš Holobar, Nina Murks,
% University of Maribor, Faculty of Electrical Engineering and Computer Science, Slovenia
%
% Website: 
% -  https://www.hybridneuro.feri.um.si/
% -  https://www.hybridneuro.feri.um.si/webinars/webinar_validation_of_results.html
% 
% HybridNeuro project project has received funding from the European Union's Horizon Europe Research and 
% Innovation Programme under grant agreement no. 101079392 and from the UK Research and Innovation (UKRI) 
% government’s Horizon Europe funding guarantee scheme under grant agreement no. 10052152.

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Read the data from experimental HDEMG signals. 
% HDEMG signals were recorded from Tibialis Anterior muscle of 19 healthy young subjects 
% during isometric ramp-up contraction from 0% to 30% of MVC. Recorded signals were decomposed
% by CKC decomposition method and edited by two experts. Unreliably identified motor units (MUs)
% were discarded and excluded from further analysis
Table1 = readtable('DATA_ExperimentalEMG.csv'); 
Table1.SubID = categorical(Table1.SubID);
Table1.MUid = categorical(Table1.MUid);
Table1.ContRep = categorical(Table1.ContRep);


%% PARAMETERS
PNRthreshold = 20; % analyse only motor units (MUs) with Pulse-to-Noise Ratio above the PNRthreshold (in dB)
% for details on PNR, see Holobar A, Minetto MA, Farina D. Accurate identification of motor unit discharge patterns from high-density
% surface EMG and validation with a novel signal-based performance metric. Journal of neural engineering. 2014 Jan 8;11(1):016008
CM = colormap(jet(1024)); % colormap

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
pTable = Table1(Table1.PNR>PNRthreshold,:); % filter the values for plots 
AllSubjects = unique(pTable.SubID);
AllRepetitions = unique(pTable.ContRep);
figure('Position',get(0,'ScreenSize')); 
for r = 1:size(pTable,1)
    Si = find(AllSubjects == pTable.SubID(r));
    Ri = find(AllRepetitions == pTable.ContRep(r));
    subplot(length(AllRepetitions),length(AllSubjects),(Ri-1)*length(AllSubjects)+Si);
    hold on;
    plot([pTable.RecThr(r),30],[pTable.RecDR(r),pTable.MaxDR(r)],'Color',CM(Si*50+Ri*1,:));
    title(['S ' num2str(Si) ',R ' num2str(Ri)]);
    if Ri == 3
        xlabel('% MVC');
    end
    if Si == 1
        ylabel('Discharge rate (Hz)');
    end
    axis([0 30 0 30])
end

%% PLOT DISTRIBURIONS
figure('Position',get(0,'ScreenSize')); 
subplot(2,2,1); histogram(Table1.MaxDR,50); title('MaxDR'); xlabel('values'); ylabel('No. of elements')
subplot(2,2,2); histogram(Table1.RecDR,50); title('RecDR'); xlabel('values'); ylabel('No. of elements')
subplot(2,2,3); histogram(Table1.RecThr,50); title('RecThr'); xlabel('values'); ylabel('No. of elements')
subplot(2,2,4); histogram(Table1.PNR,50); title('PNR'); xlabel('values (dB)'); ylabel('No. of elements')


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Linear Mixed Effect (LME) models 
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% LME model 1 - one fixed effect, without random effects 
lme1 = fitlme(Table1(Table1.PNR>PNRthreshold,:),'MaxDR ~ 1 + RecThr')
figure; subplot(2,2,1); plotResiduals(lme1,'probability');
subplot(2,2,2); plotResiduals(lme1,'lagged');
subplot(2,2,3); plotResiduals(lme1,'histogram');
subplot(2,2,4); plotResiduals(lme1,'caseorder');

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% LME model 2 - two fixed effects, without random effects
lme2 = fitlme(Table1(Table1.PNR>PNRthreshold,:),'MaxDR ~ 1 + RecThr + RecDR')
figure; subplot(2,2,1); plotResiduals(lme2,'probability');
subplot(2,2,2); plotResiduals(lme2,'lagged');
subplot(2,2,3); plotResiduals(lme2,'histogram');
subplot(2,2,4); plotResiduals(lme2,'caseorder');

results = compare(lme1,lme2)   % compare lme2 to lme1

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% LME model 3 - two fixed effects, two random effects (intercepts only), no multilevel grouping of random effects
lme3 = fitlme(Table1(Table1.PNR>PNRthreshold,:),'MaxDR ~ 1 + RecThr+RecDR + (1|SubID) + (1|ContRep) + (1|MUid)')
figure; subplot(2,2,1); plotResiduals(lme3,'probability');
subplot(2,2,2); plotResiduals(lme3,'lagged');
subplot(2,2,3); plotResiduals(lme3,'histogram');
subplot(2,2,4); plotResiduals(lme3,'caseorder');

results = compare(lme2,lme3)  % compare lme3 to lme2

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% LME model 4 - two fixed effects, three random effects (intercepts and slopes), no multilevel grouping of random effects 
lme4 = fitlme(Table1(Table1.PNR>PNRthreshold,:),['MaxDR ~ 1 +RecThr+RecDR + ' ...
    '(1+RecThr|ContRep) + (1+RecThr|SubID) + (1+RecThr|MUid)']) 
figure; subplot(2,2,1); plotResiduals(lme4,'probability');
subplot(2,2,2); plotResiduals(lme4,'lagged');
subplot(2,2,3); plotResiduals(lme4,'histogram');
subplot(2,2,4); plotResiduals(lme4,'caseorder');

results = compare(lme3,lme4)  % compare lme4 to lme3

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% LME model 5 - two fixed effects (no interaction), three random effects on the slope and but not on the intercept of RecThr, 
% multilevel grouping of random effects (MUid are per subject & are tracked across contraction repetitions ContRep) 
lme5 = fitlme(Table1(Table1.PNR>PNRthreshold,:),['MaxDR ~ 1 + RecThr+RecDR+ (-1 + RecThr|ContRep) + (-1 + RecThr|SubID) + ' ...
    '(-1 + RecThr|SubID:MUid)'])
figure; subplot(2,2,1); plotResiduals(lme5,'probability');
subplot(2,2,2); plotResiduals(lme5,'lagged');
subplot(2,2,3); plotResiduals(lme5,'histogram');
subplot(2,2,4); plotResiduals(lme5,'caseorder');

results = compare(lme4,lme5)    % compare lme5 to lme4
results = compare(lme3,lme5)    % compare lme5 to lme3

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% LME model 6 - two fixed effects + their interaction, three random effects on the slope and but not on the intercept of RecThr 
% + three random effects on the slope but not on the intercept of RecDR, 
% multilevel grouping of random effects (MUid are per subject & are tracked across contraction repetitions ContRep) 
lme6 = fitlme(Table1(Table1.PNR>PNRthreshold,:),['MaxDR ~ 1 + RecThr*RecDR+(-1 + RecThr|ContRep) + (-1 + RecThr|SubID)+ ' ...
    '(-1 + RecThr|SubID:MUid) +(-1 + RecDR|ContRep) + (-1 + RecDR|SubID)+ (-1 + RecDR|SubID:MUid)'])
figure; subplot(2,2,1); plotResiduals(lme6,'probability');
subplot(2,2,2); plotResiduals(lme6,'lagged');
subplot(2,2,3); plotResiduals(lme6,'histogram');
subplot(2,2,4); plotResiduals(lme6,'caseorder');

results = compare(lme5,lme6)   % compare lme6 to lme5

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% LME model 7 - simplified model LME 6 (without (-1 + RecThr|ContRep)  + (-1 + RecDR|ContRep) ) 
lme7 = fitlme(Table1(Table1.PNR>PNRthreshold,:),['MaxDR ~ 1 + RecThr*RecDR + (-1 + RecThr|SubID) + (-1 + RecThr|ContRep) +' ...
    ' (-1 + RecThr|SubID:MUid) + (-1 + RecDR|SubID)+ (-1 + RecDR|SubID:MUid)'])
figure; subplot(2,2,1); plotResiduals(lme7,'probability');
subplot(2,2,2); plotResiduals(lme7,'lagged');
subplot(2,2,3); plotResiduals(lme7,'histogram');
subplot(2,2,4); plotResiduals(lme7,'caseorder');

results = compare(lme6,lme7)    % compare lme7 to lme6
results = compare(lme7,lme6)    % compare lme6 to lme7 - sometimes comparison in both directions is required...


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%  Generalzed Linear Mixed Effect (GLME) models
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% GLME model 1 - two fixed effects + their interaction, three random effects on the slope and but not on the intercept of RecThr
% + three random effects on the slope but not on the intercept of RecDR, 
% multilevel grouping of random effects (MUid are per subject & are tracked across contraction repetitions ContRep) 

% OptimizerOptions = statset('fitglme'); % use Optimizer options to control the efficiency of model fitting.
% OptimizerOptions.Display = 'iter'; % display optimization iterations, default value: 'off'
% OptimizerOptions.UseParallel = true; % use Parallel toolbox, dafault value: false
glme1 = fitglme(Table1(Table1.PNR>PNRthreshold,:),['MaxDR ~ 1 + RecThr*RecDR+(-1 + RecThr|ContRep) + (-1 + RecThr|SubID)+' ...
    '(-1 + RecThr|SubID:MUid) + (-1 + RecDR|ContRep) + (-1 + RecDR|SubID)+ (-1 + RecDR|SubID:MUid)'], ...
    'Distribution','Gamma') %'OptimizerOptions',OptimizerOptions
    
figure; subplot(2,2,1); plotResiduals(glme1,'probability');
subplot(2,2,2); plotResiduals(glme1,'lagged');
subplot(2,2,3); plotResiduals(glme1,'histogram');
subplot(2,2,4); plotResiduals(glme1,'caseorder');

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Plot results of LME fitting
interceptCoeff = lme6.Coefficients{1,2};
recThCoeff = lme6.Coefficients{2,2};
recDRCoeff = lme6.Coefficients{3,2};
recTh_recDR_interactionCoeff = lme6.Coefficients{4,2};

pRecThr= Table1.RecThr(Table1.PNR>PNRthreshold);
pRecDR = Table1.RecDR(Table1.PNR>PNRthreshold);
pmaxDR = Table1.MaxDR(Table1.PNR>PNRthreshold);
pSubjectID = Table1.SubID(Table1.PNR>PNRthreshold);

[X,Y] = meshgrid(5:0.5:35,0:30);
figure('Position',get(0,'ScreenSize')); 
hold on;
surf(Y,X,interceptCoeff+recDRCoeff*X+recThCoeff*Y+recTh_recDR_interactionCoeff*X.*Y);
for p=1:length(pRecThr)
    plot3(pRecThr(p),pRecDR(p),pmaxDR(p),'k*','MarkerSize',10,'LineWidth',1) %'Color',CM(pSubjectID(p)*50,:); 
end
xlabel('recruitment threshold (% MVC)'); 
ylabel('recruitment discharge rate (Hz)'); 
zlabel('maximum MU discharge rate (Hz)'); 
colormap(CM); view(325,25); grid on;



%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%  Bayesian Linear Regression (BLR) models - requires Econometrics Toolbox
%   relatively limited models, for more complex models, STAN (https://mc-stan.org/) or
%   High-Dimensional Bayesian Regularised Regression with the BayesReg
%   Package, arXiv:1611.06649 (https://www.mathworks.com/matlabcentral/fileexchange/60823-flexible-bayesian-penalized-regression-modelling?status=SUCCESS#version_history_tab)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% BLR model 1a - two effects
NoPredictors=2;
PriorMdl1 = bayeslm(NoPredictors,ModelType="diffuse", VarNames=["RecThr" "RecDR"]);
PriorMdl1 = bayeslm(NoPredictors,ModelType="mixconjugateblm", VarNames=["RecThr" "RecDR"]);
plot(PriorMdl1); set(gcf,'Position',get(0,'ScreenSize'))

X = Table1{Table1.ContLev==30 & Table1.PNR>PNRthreshold,PriorMdl1.VarNames(2:end)};
y = Table1{Table1.ContLev==30 & Table1.PNR>PNRthreshold,"MaxDR"};
PosteriorMdl1 = estimate(PriorMdl1,X,y);
plot(PosteriorMdl1,PriorMdl1); set(gcf,'Position',get(0,'ScreenSize'))

%% Only in the case of empirical distributions
PosteriorMdl1 = estimate(PriorMdl1,X,y,'Sampler','hmc','NumDraws',1e4,'Burnin',1e3);
plot(PosteriorMdl1,PriorMdl1); set(gcf,'Position',get(0,'ScreenSize'))

% check the quality of Markov chain Monte Carlo (MCMC) traces &
% their autocorrelations
figure;
K = size(PosteriorMdl1.BetaDraws,1);
for k=1:K
    subplot(K,2,(k-1)*2+1)
    plot(PosteriorMdl1.BetaDraws(k,1001:end));
    title(['Trace Plot - \beta_' num2str(k)]); xlabel('MCMC Draw'); ylabel('Simulation Index');
    subplot(K,2,(k-1)*2+2)
    autocorr(PosteriorMdl1.BetaDraws(2,1001:end))
end

figure;
subplot(2,1,1)
plot(PosteriorMdl1.Sigma2Draws(1001:end));
title(['Trace Plot - Disturbance Variance']);
xlabel('MCMC Draw')
ylabel('Simulation Index')
subplot(2,1,2)
autocorr(PosteriorMdl1.Sigma2Draws(1001:end))


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% BLR model 2a - five effects
NoPredictors=5;
PriorMdl1 = bayeslm(NoPredictors,ModelType="diffuse", VarNames=["RecThr" "RecDR" "PNR" "SubID" "ContRep"]);
PriorMdl1 = bayeslm(NoPredictors,ModelType="mixconjugateblm", VarNames=["RecThr" "RecDR" "PNR" "SubID" "ContRep"]);
plot(PriorMdl1); set(gcf,'Position',get(0,'ScreenSize'))

X = Table1{Table1.ContLev==30 & Table1.PNR>PNRthreshold,PriorMdl1.VarNames(2:end)};
y = Table1{Table1.ContLev==30 & Table1.PNR>PNRthreshold,"MaxDR"};
PosteriorMdl1 = estimate(PriorMdl1,X,y);
plot(PosteriorMdl1,PriorMdl1); set(gcf,'Position',get(0,'ScreenSize'))

%% Only in the case of empirical distributions
PosteriorMdl1 = estimate(PriorMdl1,X,y,'Sampler','hmc','NumDraws',1e4,'Burnin',1e3);
plot(PosteriorMdl1,PriorMdl1); set(gcf,'Position',get(0,'ScreenSize'))

% check the quality of Markov chain Monte Carlo (MCMC) traces &
% their autocorrelations
figure;
K = size(PosteriorMdl1.BetaDraws,1);
for k=1:K
    subplot(K,2,(k-1)*2+1)
    plot(PosteriorMdl1.BetaDraws(k,1001:end));
    title(['Trace Plot - \beta_' num2str(k)]); xlabel('MCMC Draw'); ylabel('Simulation Index');
    subplot(K,2,(k-1)*2+2)
    autocorr(PosteriorMdl1.BetaDraws(2,1001:end))
end

figure;
subplot(2,1,1)
plot(PosteriorMdl1.Sigma2Draws(1001:end));
title(['Trace Plot - Disturbance Variance']);
xlabel('MCMC Draw')
ylabel('Simulation Index')
subplot(2,1,2)
autocorr(PosteriorMdl1.Sigma2Draws(1001:end))


%% Plot results
interceptCoeff = mean(PosteriorMdl1.BetaDraws(1,1001:end)); % random draws by MCMC to estimate the statistics...
recThCoeff = mean(PosteriorMdl1.BetaDraws(2,1001:end)); % random draws by MCMC to estimate the statistics...
recDRCoeff = mean(PosteriorMdl1.BetaDraws(3,1001:end)); % random draws by MCMC to estimate the statistics...
recTh_recDR_interactionCoeff = 0; % the model in Matlab does not support interaction assessment (Rstudio supports it). 

pRecThr= Table1.RecThr;
pRecDR = Table1.RecDR;
pmaxDR = Table1.MaxDR;
pSubjectID = Table1.SubID;

[X,Y] = meshgrid(5:0.5:22,0:30);
figure('Position',get(0,'ScreenSize')); 
hold on;
surf(Y,X,interceptCoeff+recDRCoeff*X+recThCoeff*Y+recTh_recDR_interactionCoeff*X.*Y);
for p=1:length(pRecThr)
    plot3(pRecThr(p),pRecDR(p),pmaxDR(p),'k*','MarkerSize',10,'LineWidth',1); 
end
xlabel('recruitment threshold (% MVC)'); 
ylabel('recruitment discharge rate (Hz)'); 
zlabel('maximum MU discharge rate (Hz)'); 
colormap(CM); view(325,25); grid on;