%####################################################################################################################
%% 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.

%% ##################################################################################################################################
%% GENEARATE THE DATA and REFERENCE MODEL
clc;
close all;
CM = colormap(jet(1024));

NoSubjects = 15;  % play with the number of Subjects (e.g. NoSubjects = 5, NoSubjects = 15...)
NoRepetitions = 3; % Number of contraction repetitions per subject
NoMUs = 15; % play with the number of MUs (e.g. NoMUs = 5, NoMUs = 15...)
portionOfMissedValues = 0.0; % portion of missed MUs (0 for 0%, 0.5 for 50 %) - balanced vs. unbalanced test

% FIXED FACTORS
C = 30;    % maximum contraction level & intercept
mean_recThSlope = -0.3; % mean slope of recruitment threshold factor
mean_recDRSlope = -1; % mean slope of discharge rate factor
recTh_recDR_interactionSlope = 0.0; % interaction between recruitment threshold and recruitment discharge rate

% RANDOM FACTORS
S = 0.2*(rand(1,NoSubjects) - 0.5);  % random effect per subject
R = 0.025*(rand(NoRepetitions,NoSubjects) - 0.5); % random effect per contraction repetition
M = 0.025*(rand(NoRepetitions,NoSubjects,NoMUs) - 0.5); % random effect per MU

figure('Position',get(0,'ScreenSize')/2);
subplot(1,3,1); histogram(S(:),50); title('S'); xlabel('values'); ylabel('No. of elements')
subplot(1,3,2); histogram(R(:),50); title('R'); xlabel('values');
subplot(1,3,3); histogram(M(:),50); title('M'); xlabel('values');

% GENERATE & PLOT DATA per motor unit (MU)
MaxDR = [];  % output variable or observation
RecDR = []; % fixed factor
RecThr = []; % fixed factor    
MUid = []; % grouping variable 
SubID = []; % grouping variable
ContRep = []; % grouping variable

DRintercept = []; % coefficient
recThSlope = []; % coefficient
recDRSlope = []; % coefficient


h=figure('Position',get(0,'ScreenSize'));
for Si = 1:NoSubjects
    for Ri = 1:NoRepetitions
        for Mi = 1:NoMUs  

            if rand >= portionOfMissedValues

                MUid(end+1,1) = Mi;
                SubID(end+1,1) = Si;
                ContRep(end+1,1) = Ri;

                % two factors
                RecThr(end+1,1) = Mi/NoMUs*0.7*C + 1*rand();
                RecDR(end+1,1) = 7+rand();

                % coefficients
                DRintercept(end+1,1) = C; % + S(1,Si) + R(Ri,Si) + M(Ri,Si,Mi); % random effects added to intercept (or not)
                recThSlope(end+1,1) =  mean_recThSlope + S(1,Si) + R(Ri,Si) + M(Ri,Si,Mi);  % random effects added to slope of RecThr (or not)
                recDRSlope(end+1,1) =  mean_recDRSlope;  % slope of RecDR (or not)
             
                % MaxDR generation - the last term ( 0.5*randn(1) ) codes the model residuals
                MaxDR(end+1,1) = DRintercept(end,1) +recDRSlope(end,1)*RecDR(end) + recThSlope(end,1)*RecThr(end) + ...
                    recTh_recDR_interactionSlope*RecDR(end)*RecThr(end) + 0.5*randn(1);

                % plot values per MU/Subject/Repetition
                subplot(NoRepetitions,NoSubjects,(Ri-1)*NoSubjects+Si); hold on; 
                plot([RecThr(end),30],[RecDR(end),MaxDR(end)],'Color',CM(round(Si*1000/NoSubjects+Ri*1),:));
                title(['S ' num2str(Si) ', R ' num2str(Ri)]);
                xlabel('% contraction'); ylabel('discharge rate (Hz)');
                axis([0 30 5 25]);
            end
        end
    end
end

%% plot 3D regression 
[X,Y] = meshgrid(7:0.1:8,0:25);
figure('Position',get(0,'ScreenSize')); 
hold on;
surf(Y,X,mean(DRintercept)+mean(recDRSlope)*X+mean(recThSlope)*Y+recTh_recDR_interactionSlope*X.*Y);
plot3(RecThr,RecDR,MaxDR,'*k','MarkerSize',10,'LineWidth',1); 
xlabel('recruitment threshold (% MVC)'); 
ylabel('recruitment discharge rate (Hz)'); 
zlabel('maximum MU discharge rate (Hz)'); 
colormap(CM); view(325,25); grid on;

%% simple animation
[X,Y] = meshgrid(7:0.1:8,0:25);
figure('Position',get(0,'ScreenSize')); 
hold on;
colormap(CM); 
view(125,25); 
grid on;
plot3(RecThr,RecDR,MaxDR,'*k','MarkerSize',10,'LineWidth',1); 
xlabel('recruitment threshold (% MVC)'); 
ylabel('recruitment discharge rate (Hz)'); 
zlabel('maximum MU discharge rate (Hz)'); 
for dc = 0.1:-0.01:0
    hm = mesh(Y,X,mean(DRintercept)+(dc+mean(recDRSlope))*X+(dc+mean(recThSlope))*Y+recTh_recDR_interactionSlope*X.*Y);
    drawnow;
    pause(0.1);
    delete(hm);
end
surf(Y,X,mean(DRintercept)+mean(recDRSlope)*X+mean(recThSlope)*Y+recTh_recDR_interactionSlope*X.*Y);
for a=125 + (0:5:360)
    view(a,25);
    pause(0.1);
end

%% SUMMARIZE THE STATISTICS OF THE MODEL
disp('##################################################################################################')
disp('');
disp('REGRESSION COEFFICIENTS');
disp(['mean DRintercept ' num2str(mean(DRintercept),5) ' +- '  num2str(std(DRintercept),5)]);
disp(['mean recThSlope ' num2str(mean(recThSlope),5) ' +- '  num2str(std(recThSlope),5)]);
disp(['mean RecDR ' num2str(mean(recDRSlope),5) ' +- '  num2str(std(recDRSlope),5)]);
disp('RANDOM EFFECTS')
disp(['True mean and standard deviation of random effect per Subject = ' num2str(mean(S(:))) ' +- ' num2str(std(S(:)))]);
disp(['True mean and standard deviation of  random effect per Repetition = ' num2str(mean(R(:))) ' +- ' num2str(std(R(:)))]);
disp(['True mean and standard deviation of random effect per motor unit (MU) = ' num2str(mean(M(:))) ' +- ' num2str(std(M(:)))]);
disp('##################################################################################################')
disp('');

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

%% GENERATE THE TABLE WITH MODELED DATA
Table1 = table(SubID,ContRep,MUid,RecThr,RecDR,MaxDR);

% Export data for the examples in Rstudio 
writetable(Table1,'newDATA_SyntheticEMG.csv'); 


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

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% LME model 1 - simpler statistical models, one fixed effect, without random effects 
clc;

lme1 = fitlme(Table1,'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
clc

lme2 = fitlme(Table1,'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');

compare(lme1,lme2) % compare lme2 to lme1

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% LME model 3 - two fixed effects, three random effects (intercepts only), no multilevel grouping of random effects
clc;

lme3 = fitlme(Table1,'MaxDR ~ 1 + RecThr+RecDR+ (1|ContRep) + (1|SubID) + (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');

compare(lme2,lme3) % compare lme3 to lme2

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% LME model 4 - Larger than exact model if you did not change the data generation model in MATLAB
% two fixed effects without interaction, three random effects on the slope and on the intercept of RecThr, 
% no multilevel grouping of random effects clc;

clc;

lme4 = fitlme(Table1,'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');

compare(lme3,lme4) % compare lme4 to lme3

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% LME model 5 - EXACT MODEL if you did not change the data generation model in MATLAB
% 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) clc

clc;

lme5 = fitlme(Table1,'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');

compare(lme4,lme5) % compare lme5 to lme4

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% LME model 6 - Larger than exact model if you did not change the data generation model in MATLAB
% two fixed effects with interaction, three random effects on the slope and but not on the intercept of RecThr, 
% no multilevel grouping of random effects 
clc;

lme6 = fitlme(Table1,'MaxDR ~ 1 + RecThr*RecDR+(-1 + RecThr|ContRep) + (-1 + RecThr|SubID)+ (-1 + RecThr|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');

compare(lme5,lme6)  % compare lme6 to lme5


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

glme1 = fitglme(Table1,'MaxDR ~ 1 + RecThr+RecDR + (-1 + RecThr|ContRep) + (-1 + RecThr|SubID)+ (-1 + RecThr|MUid)',...
    'Distribution','Normal','Link','identity')

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');

% compare(lme5,glme1)  % This creates error, we cannot compare LME and GLME models...

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%  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/predictors/factors (fixed factors from LMM) 
NoPredictors = 2;

% 1) start with diffuse prior model - prior values and distribution are unknown   
PriorMdl1 = bayeslm(NoPredictors,ModelType="diffuse", VarNames=["RecThr" "RecDR"]);
plot(PriorMdl1); set(gcf,'Position',get(0,'ScreenSize'))

X = Table1{:,PriorMdl1.VarNames(2:end)}; % get RecThr, RecDR and ContRep as predictors/factors
y = Table1{:,"MaxDR"}; % make the MaxDR outcome variable
PosteriorMdl1 = estimate(PriorMdl1,X,y);
plot(PosteriorMdl1,PriorMdl1); set(gcf,'Position',get(0,'ScreenSize'))

% 2) start with semiconjugate prior model - factors are independant   
PriorMdl1 = bayeslm(NoPredictors,ModelType="mixsemiconjugate", VarNames=["RecThr" "RecDR"]);
plot(PriorMdl1); set(gcf,'Position',get(0,'ScreenSize'))

X = Table1{:,PriorMdl1.VarNames(2:end)}; % get RecThr, RecDR and ContRep as predictors/factors
y = Table1{:,"MaxDR"}; % make the MaxDR outcome variable
PosteriorMdl1 = estimate(PriorMdl1,X,y);
plot(PosteriorMdl1,PriorMdl1); set(gcf,'Position',get(0,'ScreenSize'))

% 3) correct the means of semiconjugate prior model - factors are independant   
PriorMdl1 = bayeslm(NoPredictors,ModelType="mixsemiconjugate", VarNames=["RecThr" "RecDR"], Mu=[30; -0.3; -1.0;]);
plot(PriorMdl1); set(gcf,'Position',get(0,'ScreenSize'))

X = Table1{:,PriorMdl1.VarNames(2:end)}; % get RecThr, RecDR and ContRep as predictors/factors
y = Table1{:,"MaxDR"}; % make the MaxDR outcome variable
PosteriorMdl1 = estimate(PriorMdl1,X,y);
plot(PosteriorMdl1,PriorMdl1); set(gcf,'Position',get(0,'ScreenSize'))

% 4) empirically estmated means 
PriorMdl1 = bayeslm(NoPredictors,ModelType="mixsemiconjugate", VarNames=["RecThr" "RecDR"], Mu=[mean(Table1.MaxDR); mean(Table1.RecThr); mean(Table1.RecDR)]);
plot(PriorMdl1); set(gcf,'Position',get(0,'ScreenSize'));

X = Table1{:,PriorMdl1.VarNames(2:end)}; % get RecThr, RecDR and ContRep as predictors/factors
y = Table1{:,"MaxDR"}; % make the MaxDR outcome variable
PosteriorMdl1 = estimate(PriorMdl1,X,y);
plot(PosteriorMdl1,PriorMdl1); set(gcf,'Position',get(0,'ScreenSize'))


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% BLR model 2a - 5 effects/predictors/factors (fixed factors + random factors from LMM) 
NoPredictors = 5;

% 1) start with diffuse prior model - prior values and distribution are unknown   
PriorMdl2 = bayeslm(NoPredictors,ModelType="diffuse", VarNames=["RecThr" "RecDR" "ContRep" "SubID" "MUid"]);
plot(PriorMdl2); set(gcf,'Position',get(0,'ScreenSize'))

X = Table1{:,PriorMdl2.VarNames(2:end)}; % get RecThr, RecDR and ContRep as predictors/factors
y = Table1{:,"MaxDR"}; % make the MaxDR outcome variable
PosteriorMdl2 = estimate(PriorMdl2,X,y);
plot(PosteriorMdl2,PriorMdl2); set(gcf,'Position',get(0,'ScreenSize'))

% 2) start with semiconjugate prior model - factors are independant   
PriorMdl2 = bayeslm(NoPredictors,ModelType="conjugate", VarNames=["RecThr" "RecDR"   "ContRep" "SubID" "MUid"]);
plot(PriorMdl2); set(gcf,'Position',get(0,'ScreenSize'))

X = Table1{:,PriorMdl2.VarNames(2:end)}; % get RecThr, RecDR and ContRep as predictors/factors
y = Table1{:,"MaxDR"}; % make the MaxDR outcome variable
PosteriorMdl2 = estimate(PriorMdl2,X,y);
plot(PosteriorMdl2,PriorMdl2); set(gcf,'Position',get(0,'ScreenSize'))

% 3) correct the means of semiconjugate prior model - factors are independant   
PriorMdl2 = bayeslm(NoPredictors,ModelType="conjugate", VarNames=["RecThr" "RecDR"   "ContRep" "SubID" "MUid"], Mu=[30; -0.3; -1.0; 0.0; 0.0; 0.0]);
plot(PriorMdl2); set(gcf,'Position',get(0,'ScreenSize'))

X = Table1{:,PriorMdl2.VarNames(2:end)}; % get RecThr, RecDR and ContRep as predictors/factors
y = Table1{:,"MaxDR"}; % make the MaxDR outcome variable
PosteriorMdl2 = estimate(PriorMdl2,X,y);
plot(PosteriorMdl2,PriorMdl2); set(gcf,'Position',get(0,'ScreenSize'))

%% Plot results
interceptCoeff = PosteriorMdl2.Mu(1);
recThCoeff = PosteriorMdl2.Mu(2);
recDRCoeff = PosteriorMdl2.Mu(3);
recTh_recDR_interactionCoeff = 0;

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

[X,Y] = meshgrid(7:0.1:8,0:25);
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;