% function [] = RunRegression(i,m)
%i = subject, m = model number
%Single trial regression scripts provided by AGF for single-trial regression workshop at MCC 2019. 
%COMPASS ICA analyses by JRW. Thanks to Theo Gründler (fconv_tg) & Gerhard Jocham.
%If you use the scripts, please quote at least Fischer & Ullsperger, Neuron, 2013 & Fischer et al., Scientific Reports, 2016. 
%AGF, 2019

%% 0. Get started
clear; clc; close all

i=2; %index to the subject to be used (can also be run as function)

CF = matlab.desktop.editor.getActive; cd(fileparts(CF.Filename)); %just change to folder of this file
path('STATB/',path);path('STATB/helpers/',path);  %add path of STA-TB to matlab

%define data and output paths
pathEEG = 'data/';                  % path to EEG
pathBeh = 'Part_Behav/';            % path to behavioural data
pathAna = 'AnalysisInfo/';          % path to information about analysis steps
pathOut = 'output/within_subject/'; % output of regression is stored here
pathPic = 'plots/';                 % output of plots

%% 1. Behaviour preparation
%load behaviour
F = dir([pathEEG '*.set']);
load([pathAna F(i).name(1:6) '.mat']);

P_Info = VP_INFO;                   % a structure containing additional information about pre-processing steps
P_name = P_Info.name;               % name of the participants dataset

load([pathBeh F(i).name(1:6) '.mat']);
D = Ulogg;                          % this is a structure containing the behaviour for every trial
disp(['Finished loading analysis information and behavioural data for participant: ' P_name])

%% 1.1 EEG Preparation
Feeg = dir([pathEEG '*' P_name '*.set']); % Find the dataset of this participant in EEG folder

%load EEG
EEG = pop_loadset([pathEEG Feeg.name]);   % load it into EEGlab format 

%Remove artifact ICA components: these have been stored in structure P_Info.IC_All_Rem
display([num2str(length(P_Info.IC_All_Rem)) ' ICs have been removed...'])
EEG = pop_subcomp(EEG, P_Info.IC_All_Rem, 0); %these have been stored in P_Info


%% 2. Re-epoch EEG locked to response and ENSURE that trials are identical in log-file and EEG
event_index = []; epoch_index = []; RT_and_Lat = [];
for e = 1 : length(EEG.epoch) %loop through all epochs
    I0=find([EEG.epoch(e).eventlatency{:}]==0); % Locking event (stimulus)
    if length(EEG.epoch(e).eventtype)-I0>1 %Prevent indexing when no following event is present
        %safety: it is crucial that behavioral file and EEG are exactly the same on every trial, this here just compares stimulus types for identitiy
        if strcmp(EEG.epoch(e).eventtype{I0+1}(6:8), 'Inc') && D(e).congr~=2
            disp(['Incongruent EEG trial but congruent logfile in trial ' num2str(e) '!'])
            AlignmentError = 1;
        elseif strcmp(EEG.epoch(e).eventtype{I0+1}(6:8), 'Con') && D(e).congr~=1
            disp(['Congruent EEG trial but incongruent logfile in trial ' num2str(e) '!']) 
            AlignmentError = 1;
        else
            AlignmentError = 0;
        end
        
        %safety two: the difference in EEG latency and recorded RT should be very small
        RT_and_Lat(e,1) = D(e).rt; %logfile RT just for comparison
        RT_and_Lat(e,2) = EEG.epoch(e).eventlatency{I0+1}; %latency of response relative to stimulus
        
        %if general check passed, include this trial, otherwise, skip this trial
        if strcmp(EEG.epoch(e).eventtype{I0+1}(1), 'R') && ~AlignmentError % First event following stimulus should be a response and there was no alignment mismatch
            event_index = [event_index EEG.epoch(e).event(I0+1)];   % Get the event index
            epoch_index = [epoch_index e];
        end
    end
end
[RESP, indices] = pop_epoch(EEG, {}, [-1.1 1.01], 'eventindices', event_index); %Re-epoch response locked, in case of out of boundary, remember these accepted epochs
D=D(intersect(epoch_index,indices)); %remove trials in behavioral data as well (those selected in the loop, and thos that were additionally out of bounds)

%% make a little plot to compare RT in logfile and EEG
figure;
subplot(211)
plot(RT_and_Lat); legend({'log' 'eeg'}); ylabel({'RT (ms)'}); title('RT comparison')
subplot(212)
plot(RT_and_Lat(:,1)-RT_and_Lat(:,2)); ylabel({'\Delta RT (ms)'}); title(['RT comparison (summed abs-deviation = ' num2str(nansum(abs(diff(RT_and_Lat')))) ' ms).'])
xlabel({'trial no'}); 


%% 3. Remove additional artifacts separately for two categories of trials
S.max_p = 5;            % remove no more than 5% of trials...
S.min_n = 1;            % but find a threshold that removes at least one trials
S.StartSD = 8;          % start at high threshold (conservative, do not remove too many trials!)
S.EpLimit = [-0.5 1];   % sometimes, you may have longer epochs, but artifacts far away from the event of interest do not matter. Just limit the time to look for artifacts here

RemoveErrors  = STA_adaptive_artifact(RESP, [D.error]==1, S); % find trials with artifacts in errors
RemoveCorrect = STA_adaptive_artifact(RESP, [D.error]==0, S); % find trials with artifacts in corrects

%sometimes, there are general trial exclusion criteria. Here, we want only trials in a certain RT range
Behav_Rem_Trials = [D.rt] < 80 | [D.rt] > 900; %trials not in the RT range of 80 to 900 ms

%remove these trials from EEG
RESP = pop_select(RESP, 'trial', setdiff(1:length(RESP.epoch),[RemoveErrors  RemoveCorrect find(Behav_Rem_Trials)]));

%remove these trials from behavior
D([RemoveErrors RemoveCorrect find(Behav_Rem_Trials)])=[];


%% 4. Setup the design matrix

%define regressors
R_error     = [D.error]'*2-1;       % error = 1, correct = -1
R_incongr   = ([D.congr]'-1)*2-1;   % incongruent = 1, congruent = -1
R_logRT     = log([D.rt]');         % log-RT

% R_error     = demean([D.error]'*2-1);       % error = 1, correct = -1
% R_incongr   = demean(([D.congr]'-1)*2-1);   % incongruent = 1, congruent = -1
% R_logRT     = demean(log([D.rt]'));         % log-RT

%define time and electrodes to be included
TimeWin     = [-400 1000]; %the regression is performed in this time window
Electrodes  = structfind(RESP.chanlocs,'labels',{'Fz' 'FCz' 'Cz' 'C3' 'C4' 'F5' 'F6'}); %'Fz' 'FCz'

%just visualize the designmatrix
DM = [R_error R_incongr norm_and_scale(R_logRT, 0, 1)];
clf
imagesc(DM); colormap('gray'); title('Designmatrix'); ylabel('trials'); xlabel('regressors'); set(gca, 'XTick', [1 2 3], 'XTickLabel', {'accuracy' 'congruence' 'RT'})



%% 5. Call the regression
m=1;
switch m
 case 1
    ModelName = 'Simple_Error_Model'; %just a name for the model: will be used to create a folder to save the data in
    
    %here, you can select trials included in this specific model. 
    select_trials = find([D.posterr]==0); %we only want trials that are not preceded by an error, i.e., C-C or C-E trials only. This also removes trials after pauses!

    Predictors = {R_error(select_trials)     R_incongr(select_trials)   R_logRT(select_trials) };   %cell of all predictors
    Reg2Name   = {'Accuracy'                'Congruence'                'RT'                    };  %Cell of predictor names
    RegLables  = [{'cor' 'err'};            {'con' 'inc'};              {'low' 'high'}        ];    %Lables for the names
    %call to the regression script
    [ReggyD, Info] = STA_Fast_Regress(RESP, Electrodes, TimeWin, 'model 1', {{select_trials} Predictors {'EEG'} {'model1'}}, ...
        'PredNames', Reg2Name, 'PredLabels', RegLables, 'bin_size', 4, 'stepsize',1, 'TF', [], 'binEEG', 3, 'RetAll', 3, ...
        'NormaliseOn', 0, 'Downsample', 5, 'ERPimage', 1);
end


%% 6. save the output for this model and this subject
clear VP_MODEL %clear in case we run more models
VP_MODEL.(ModelName) = ReggyD; %combine output
VP_MODEL.([ModelName '_Info']) = Info;
path_reg_m = [pathOut ModelName '/']; %combine path for output and the name of this model
if ~exist(path_reg_m) %if path is nonexist, we create it
    mkdir(path_reg_m);
end
save([path_reg_m P_name], 'VP_MODEL') %save it


%% 7. Plot within subjects effects
ModelName               = 'Simple_Error_Model';
settings                = []; %settings are passed to plot function

settings.plotsteps      = [-200 -100 0 50 100 150 200 300 400 500]; %define time points of topo plots
settings.plotelect      = 'Cz'; %define electrode to plot time course (if empty, uses electrode with maximum effect)
settings.ERP_Im_plot    = 1; %plot an ERP image
settings.PlotPval       = 1; %plot a time course for the p-values
settings.UseValues      = 't'; %Use b or t values?
settings.AddString      = [' ' P_name ' ' settings.UseValues]; %we add the participants name to the output file
% settings.exclude        = setdiff(1:length(F), i); %only relevant when more participants are in the folder, but you want to plot single subjects

STA_Plot_Regression( [pathOut ModelName '/'], [pathPic ModelName], settings )


%% 8. Run with interaction and plot
select_trials = find([D.posterr]==0);
Predictors = {R_error(select_trials)     R_incongr(select_trials)   R_logRT(select_trials)      demean(R_logRT(select_trials)).*demean(R_error(select_trials))}; 
Reg2Name   = {'Accuracy'                'Congruence'                'RT'                        'Acc_X_RT'};
RegLables  = [{'cor' 'err'};            {'con' 'inc'};              {'low' 'high'};             {'low' 'high'}];

ModelName = 'Simple_Error_Model_Interact';

[ReggyD, Info] = STA_Fast_Regress(RESP, Electrodes, TimeWin, 'model 1', {{select_trials} Predictors {'EEG'} {'model1'}}, ...
    'PredNames', Reg2Name, 'PredLabels', RegLables, 'bin_size', 4, 'stepsize',1, 'TF', [], 'binEEG', 3, 'RetAll', 3, ...
    'NormaliseOn', 0, 'Downsample', 5, 'ERPimage', 1);

VP_MODEL2.(ModelName) = ReggyD; VP_MODEL2.([ModelName '_Info']) = Info; path_reg_m = [pathOut ModelName '/']; if ~exist(path_reg_m); mkdir(path_reg_m); end
save([path_reg_m P_name], 'VP_MODEL2')

%% 8.1 Plot the model with the interaction and safe the figure as a PDF
ModelName = 'Simple_Error_Model_Interact';
settings.UseValues = 't';
STA_Plot_Regression( [pathOut ModelName '/'], [pathPic ModelName], settings )

%% 8.2 Orthogonal regressors
select_trials = find([D.posterr]==0);

%orthogonalize RT wrt accuracy --> all shared variance as ascribed to accuracy!
[~,~,RT_resid] = regress(R_logRT(select_trials),[ones(length(select_trials),1) R_error(select_trials)]);

Predictors = {R_error(select_trials)     R_incongr(select_trials)   RT_resid}; 
Reg2Name   = {'Accuracy'                'Congruence'                'RT_Ort'             };
RegLables  = [{'cor' 'err'};            {'con' 'inc'};              {'fast' 'slow'}      ];

ModelName = 'Simple_Error_Model_Orthogonal';

[ReggyD, Info] = STA_Fast_Regress(RESP, Electrodes, TimeWin, 'model 1', {{select_trials} Predictors {'EEG'} {'model1'}}, ...
    'PredNames', Reg2Name, 'PredLabels', RegLables, 'bin_size', 4, 'stepsize',1, 'TF', [], 'binEEG', 3, 'RetAll', 3, ...
    'NormaliseOn', 0, 'Downsample', 5, 'ERPimage', 1);

VP_MODEL2b.(ModelName) = ReggyD; VP_MODEL2b.([ModelName '_Info']) = Info; path_reg_m = [pathOut ModelName '/']; if ~exist(path_reg_m); mkdir(path_reg_m); end
save([path_reg_m P_name], 'VP_MODEL2b')

%% 8.3 Plot the model again and safe the figure as PDF
ModelName = 'Simple_Error_Model_Orthogonal';
settings.UseValues = 't'; settings.exclude = [];
STA_Plot_Regression( [pathOut ModelName '/'], [pathPic ModelName], settings )


%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%LETS HAVE A BREAK!%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% 9. IC based regression
%make a simple plot to confirm the selected ICs
clf; z=2; s=10; 
for c = 1 : 20
    subplot(z,s,c)
    topoplot(RESP.icawinv(:,c), RESP.chanlocs);
    title(['IC no ' num2str(c)])
end

%% Run the regression analysis on ICs instead of electrodes
ModelName = 'Simple_Error_Model_IC_Based';
%use same model as before
select_trials   = find([D.posterr]==0);
Predictors      = {R_error(select_trials)     R_incongr(select_trials)   R_logRT(select_trials)};
Reg2Name        = {'Accuracy'                'Congruence'                'RT'};
RegLables       = [{'cor' 'err'};            {'con' 'inc'};              {'low' 'high'}];

[ReggyD, Info] = STA_Fast_Regress(RESP, [], TimeWin, 'model 1', {{select_trials} Predictors {'EEG'} {'model1'}}, ...
       'PredNames', Reg2Name, 'PredLabels', RegLables,'bin_size', 4, 'stepsize',1,...
       'binEEG', 3, 'RetAll', 3, 'Downsample', 5, 'ERPimage', 1, ...
       'UseIca', {1; 4; [1 4]}, 'icaname', {'ErrC1' 'ErrC2' 'Combo'}); %ICs for VP05 {1; 4; [1 4]}      %ICs for VP60{6; 7; [6 7]}
   
%save data for IC model
VP_MODEL3.(ModelName) = ReggyD; VP_MODEL3.([ModelName '_Info']) = Info; path_reg_m = [pathOut ModelName '/']; if ~exist(path_reg_m); mkdir(path_reg_m); end
save([path_reg_m P_name], 'VP_MODEL3')


%% 10. Plot IC based analysis
ModelName = 'Simple_Error_Model_IC_Based';
settings = [];
settings.plotsteps      = [-200 -100 0 50 100 150 200 300 400 500];
settings.plotelect      = 'Combo';
settings.ERP_Im_plot    = 1; 
settings.UseValues      = 't';
settings.AddString      = [' ' P_name ' ' settings.UseValues ];
settings.exclude        = setdiff(1:length(F), i);

STA_Plot_Regression( [pathOut ModelName '/'], [pathPic ModelName], settings )


%% 11.1 EEG as predictor
%before running, reload (first section) the data and use participant 2!

%define additional regressors
Y_foll_RT        = [D.rt_follow]';           % RT of the following trial
R_foll_congr     = ([D.congr]'-1)*2-1;       % congruent = -1, incongruent = 1
R_foll_dist      = ([D.dist_follow]'-1)*2-1; % Flanker distance of the following trial, close = -1, far = 1
R_foll_RSI       = [D.RSI_follow]';          % RSI (response stimulus interval) between the following and the current trial
R_foll_RSI(R_foll_RSI==250) = -1;            % set to -1 and 1
R_foll_RSI(R_foll_RSI==700) = 1;  

Electrodes  = structfind(RESP.chanlocs,'labels',{'Fz' 'FCz' 'Cz' 'C3' 'C4' 'F5' 'F6' 'CPz' 'Pz' 'Oz'});
TimeWin = [-100 400];

ModelName = 'Following_Trial_Model';
    
select_trials = find([D.error]'==1 & ~isnan(Y_foll_RT) & Y_foll_RT>90 & Y_foll_RT<800); %we only include error trials with reasonable RTs!

Predictors = {R_foll_congr(select_trials)     R_foll_dist(select_trials)   R_foll_RSI(select_trials)};
Reg2Name   = {'Foll congruence'               'Foll distance'              'Foll RSI'};
RegLables  = [{'con' 'inc'};                  {'close' 'far'};             {'250' '700'}];
Observation = Y_foll_RT(select_trials);

[ReggyD, Info] = STA_Fast_Regress(RESP, Electrodes, TimeWin, 'model 1', {{select_trials} Predictors {Observation} {'PESasY'}},...
    'PredNames', Reg2Name, 'PredLabels', RegLables,'bin_size', 4, 'stepsize',1, 'TF', [], ...
      'binEEG', 3, 'RetAll', 3, 'NormaliseOn', 0, 'Downsample', 5); %

%save data for Electrode model
VP_MODEL4.(ModelName) = ReggyD; VP_MODEL4.([ModelName '_Info']) = Info; path_reg_m = [pathOut ModelName '/']; if ~exist(path_reg_m); mkdir(path_reg_m); end
save([path_reg_m P_name], 'VP_MODEL4')

%% 11.1 Plot within subjects effects
ModelName = 'Following_Trial_Model';
settings = [];
%define time points of topo plots
settings.plotsteps = [-100 0 50 80 100 150 200 300 400 ];
settings.plotelect = 'Cz'; %define electrode to plot time course (if empty, uses electrode with maximum effect)
settings.ERP_Im_plot = 0; %Plot ERP images as well
settings.MaskPval = 0.2;
settings.AddString = [' ' P_name];
settings.PlotCorr = 0;
settings.UseValues = 't';
% settings.exclude = 1; %if we have more participants in one folder, but want within-participants plots, we can exclude the other ones here (1 = subject 1 is excluded!)
settings.exclude = setdiff(1:length(F), i); 
STA_Plot_Regression( [pathOut ModelName '/'], [pathPic ModelName], settings )


%% 11.2 EEG as predictor IC based
ModelName = 'Following_Trial_Model_IC_Based';

[ReggyD, Info] = STA_Fast_Regress(RESP, Electrodes, TimeWin, 'model 1', {{select_trials} Predictors {Observation} {'PESasY'}},...
    'PredNames', Reg2Name, 'PredLabels', RegLables,'bin_size', 4, 'stepsize',1, 'TF', [], ...
      'binEEG', 3, 'RetAll', 3, 'NormaliseOn', 0, 'Downsample', 5, 'UseIca', {6; 7; [6 7]}, 'icaname', {'ErrC1' 'ErrC2' 'Combo'}); %

%save data for IC model
VP_MODEL5.(ModelName) = ReggyD; VP_MODEL5.([ModelName '_Info']) = Info; path_reg_m = [pathOut ModelName '/']; if ~exist(path_reg_m); mkdir(path_reg_m); end
save([path_reg_m P_name], 'VP_MODEL5')  

%% Plot within subjects effects
ModelName = 'Following_Trial_Model_IC_Based';
settings = [];
%define time points of topo plots
settings.plotsteps = [-100 0 50 80 100 150 200 300 400 ];
settings.plotelect = 'ErrC1'; %define electrode or IC to plot (e.g., ErrC1, ErrC2 or Combo)
settings.ERP_Im_plot = 0; %Plot ERP images as well
settings.MaskPval = 0.2;
settings.AddString = [' ' P_name];
settings.PlotCorr = 0;
settings.UseValues = 't';
settings.exclude = setdiff(1:length(F), i); %if we have more participants in one folder, but want within-participants plots, we can exclude the other ones here
% settings.exclude = 1; %alternatively, just enter a number
STA_Plot_Regression( [pathOut ModelName '/'], [pathPic ModelName], settings )



































