clear all; close all;

%addpath ../../Toolboxes/Palamedes

%% Define what to analyse

% Which subjects?
subjects = [1:34];

% Which data?
load('P:/3018029.09/data.mat')
load('P:/3018029.09/data_sd.mat')

%% ======= Removal of invalid trials ======== %%

% Steps:
% 1: Remove all trials on which no direction and/or confidence response was given
% 2: Remove lower-end RT outlier trials

%% STEP 1: Remove no response trials

% Data table (standard): remove no response trials
data_clean1 = data(~isnan(data.responseDir) & ~isnan(data.confidence),:);

% Data table for sd analysis: remove no response trials
data_sd_clean1 = data_sd(~isnan(data_sd.prevResponseDir) & ~isnan(data_sd.currResponseDir) ...
    & ~isnan(data_sd.prevConfidence) & ~isnan(data_sd.currConfidence),:);

% Count how many trials are removed per subject per stimulus level
for iSubject = 1:length(subjects)
    
    subject_data    = data(data.subject == subjects(iSubject), :); % take the original data
    stim_levels     = table2array(unique(subject_data(:,'relativeCoherenceAbs')));
    
    for iStimulus = 1:length(stim_levels)
    
        % rows are participants, columns are stimulus levels
        trials_removed_per_subj_per_stim{1}(iSubject,iStimulus) = ...
            height(subject_data((isnan(subject_data.responseDir) | ...
                isnan(subject_data.confidence)) & ...
                subject_data.relativeCoherenceAbs == stim_levels(iStimulus),:));
            
    end
end

% Count how many trials are removed per stimulus level
trials_removed_per_stim{1} = sum(trials_removed_per_subj_per_stim{1},1);
% Count how many trials are removed per subject
trials_removed_per_subj{1} = sum(trials_removed_per_subj_per_stim{1},2);


%% STEP 2: Remove RT outliers under lower end cut-off

% Criterion (in seconds); trials lower than this rt get removed
RT_low = .3; 

% Data table (standard)
data_clean2 = data_clean1(data_clean1.rt >= RT_low,:);

% Data table (sd)
data_sd_clean2 = data_sd_clean1(data_sd_clean1.prevRt >= RT_low & data_sd_clean1.currRt >= RT_low,:);

% Count how many trials are removed per subject per stimulus level
for iSubject = 1:length(subjects)
    
    subject_data    = data_clean1(data_clean1.subject == subjects(iSubject), :); % take the data as it was cleaned after step 1
    stim_levels     = table2array(unique(subject_data(:,'relativeCoherenceAbs')));
    
    for iStimulus = 1:length(stim_levels)
    
        % rows are participants, columns are stimulus levels
        trials_removed_per_subj_per_stim{2}(iSubject,iStimulus) = ...
            height(subject_data(subject_data.rt < RT_low & ...
                subject_data.relativeCoherenceAbs == stim_levels(iStimulus),:));
            
    end
end

% Count how many trials are removed per stimulus level
trials_removed_per_stim{2} = sum(trials_removed_per_subj_per_stim{2},1);
% Count how many trials are removed per subject
trials_removed_per_subj{2} = sum(trials_removed_per_subj_per_stim{2},2);


%% Count how many trials are removed per step
trials_removed_per_step = ...
    ["total starting trials",   height(data); ...
    "removed in step 1",        height(data)-height(data_clean1); ...
    "removed in step 2",        height(data_clean1)-height(data_clean2); ...
    "total removed",            height(data)-height(data_clean2); ...
    "total remaining trials",   height(data_clean2)];


%% Create overview file
trial_removal_overview = ...
    {"trials removed per step",                         trials_removed_per_step;
    "trials removed in step 1 per pp per stim level",   trials_removed_per_subj_per_stim{1};
    "trials removed in step 2 per pp per stim level",   trials_removed_per_subj_per_stim{2};
    "trials removed in step 1 per pp",                  trials_removed_per_subj{1};
    "trials removed in step 1 per stim level",          trials_removed_per_stim{1};
    "trials removed in step 2 per pp",                  trials_removed_per_subj{2};            
    "trials removed in step 2 per stim level",          trials_removed_per_stim{2};
    "trials removed total per pp",                      trials_removed_per_subj{1} + trials_removed_per_subj{2};
    "trials removed total per stim level",              trials_removed_per_stim{1} + trials_removed_per_stim{2}};

% Save overview file
save('P:/3018029.09/Output variables/Trial removal overview.mat', 'trial_removal_overview')

%% Save cleaned data

% Rename cleaned data to data
data        = data_clean2;
data_sd     = data_sd_clean2;

% Save clean data
save('P:/3018029.09/data_clean.mat', 'data')
save('P:/3018029.09/data_sd_clean.mat', 'data_sd')

