% function [] = RunRegressionTF(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=1; %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


%% Preparation is the same as before, just execute this here
for PreParation = 1
%% 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 conatining 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
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], '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)

%% 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)])=[];
end


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

%define time and electrodes to be included
TimeWin     = [-400 950]; %this is just the regression time window, TF is performed with all available points


%% 5. Run TF Decomposition and model
TFElectrodes = structfind(RESP.chanlocs,'labels',{'AF3' 'AF4' 'Fz' 'FCz' 'Cz' 'CPz' 'Pz' 'C3' 'C4'  'O1' 'O2'}); %

TF.frequencies  = [4 20];
TF.stepnumber   = 12;
TF.space        = 'lin';
TF.basetime     = [-500 -300];
TF.bandfreq     = {};
TF.padding      = 0;
TF.cyclenumber  = 4.5;
TF.transform    = 'log';
TF.basetype     = 'none';
TF.bands        = {};

ModelName = 'Simple_Error_Model_TF';

select_trials = find([D.posterr]==0); %we only want trials that are not preceeded 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)};
Reg2Name   = {'Accuracy'                'Congruence'                'RT'};
RegLables  = [{'cor' 'err'};            {'con' 'inc'};              {'low' 'high'}];

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

clear VP_MODEL; VP_MODEL.(ModelName) = ReggyD; VP_MODEL.([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_MODEL')  

%% 5.1 Plot within subjects effects
ModelName = 'Simple_Error_Model_TF';
s1 = [];
%define time points of topo plots
s1.plotsteps = [-200 -100 0 50 100 150 200 300 400 500];
%define electrode to plot time course (if empty, uses electrode with maximum effect)
s1.MaskPval = 1;
s1.printstring = 'pdf';
STA_Plot_Regression( [pathOut ModelName '/'], [pathPic ModelName], s1 )


%% 6. Run TF decomposition in IC space
%get COMPASS output components
load(['COMPASS/ERN/Subject-' num2str(i) '-Components_Final.mat']);
ModelName = 'Simple_Error_Model_TF_IC';

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


clear VP_MODEL; VP_MODEL.(ModelName) = ReggyD; VP_MODEL.([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_MODEL') 

%% 6.1 Plot within subjects effects IC based
ModelName = 'Simple_Error_Model_TF_IC';
s1 = [];
%define time points of topo plots
s1.plotsteps = [-200 -100 0 50 100 150 200 300 400 500];
%define electrode to plot time course (if empty, uses electrode with maximum effect)
s1.MaskPval = 1;
s1.printstring = 'pdf';
s1.plotelect   = 'ErrC1';
s1.MaskPval    = 1;

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

%% 6.2 Use FDR correction for this plot
ModelName = 'Simple_Error_Model_TF_IC';

s1.MaskPval       = 'fdr';
s1.fdr_q          = 0.05 / 3;
s1.UseValues      = 'b';
s1.AddString      = ' fdr';

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


%% 7. Run a theta specific model in IC space
%get COMPASS output components
TF.frequencies  = [4 8];    %outline of frequencies
TF.stepnumber   = 8;        %8 steps = steps in 0.5 Hz
TF.space        = 'lin';    %use linear steps (4 4.5 4 5.5 etc)
TF.bandfreq     = {[4 8]};  %Average the data we have across these frequencies (this is what we entered, but we could well define low theta ([4 6]) and high theta ([6.1 8]) if we have specific hypotheses)
TF.bands        = {'theta'}; %theta = 4 to 8 Hz
load(['COMPASS/ERN/Subject-' num2str(i) '-Components_Final.mat']); %these are the ICs selected by COMPASS, we want the one for our subject i

ModelName = 'Simple_Error_Model_TF_IC_Theta';
[ReggyD, Info] = STA_Fast_Regress(RESP, TFElectrodes, TimeWin, 'model 1', {{select_trials} Predictors {'EEG'} {'TF1'}},...
    'PredNames', Reg2Name, 'PredLabels', RegLables, 'TF', TF, 'bin_size', 4, 'stepsize',1,'binEEG', 3,...
     'binEEG', 3, 'RetAll', 3, 'NormaliseOn', 0, 'Downsample', 5, 'ERPimage', 0, 'UseIca', {components'}, 'icaname',{'ErrC1'}, 'IcaCSpace',1);


clear VP_MODEL; VP_MODEL.(ModelName) = ReggyD; VP_MODEL.([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_MODEL') 

%% 7.1 Plot within subjects effects IC based
ModelName = 'Simple_Error_Model_TF_IC_Theta';
s1 = [];
s1.plotsteps = [-200 -100 0 50 100 150 200 300 400 500]; %define time points of topo plots
s1.MaskPval = 0.05; %we use conventional masking here
s1.printstring = 'pdf'; %print this as PDF
s1.plotelect = 'ErrC1'; %define electrode to plot time course (if empty, uses electrode with maximum effect)
s1.st.IndiviPlot = 1;
STA_Plot_Regression( [pathOut ModelName '/'], [pathPic ModelName], s1 )




























