%% I_SS_prep_data_for_analysis_20170922

% 180124 | now also exclude ART and EMG artifacts
% 180327 | include OAs, deal wih dropped trials of subjects with missing onsets
%        | interpolate data to requested times vs. choosing shortest
%        duration (zero interpolation)

% This script fixes issues with missing onsets in the following way:

% IMPORTANT: For 2160, the first two trials were not available,
% hence we have to add +2 to the index to identify the correct
% behavioral trials that belong to the EEG data. For 2203 the
% first trial is missing. The third and second trial should
% also be removed, regardless of whether they are indicated to
% be dropped due to automatic correction.

%% initialize

restoredefaultpath;
clear all; close all; pack; clc;

%% pathdef

if ismac
    pn.study        = '/Volumes/LNDG/Projects/StateSwitch/';
    pn.eeg_root     = [pn.study, 'dynamic/data/eeg/task/A_preproc/SA_preproc_study/'];
    pn.EEG          = [pn.eeg_root, 'B_data/C_EEG_FT/']; %mkdir(pn.EEG);
    pn.History      = [pn.eeg_root, 'B_data/D_History/']; %mkdir(pn.History);
    % add ConMemEEG tools
    pn.MWBtools     = [pn.eeg_root, 'T_tools/fnct_MWB/'];           addpath(genpath(pn.MWBtools));
    pn.THGtools     = [pn.eeg_root, 'T_tools/fnct_THG/'];           addpath(genpath(pn.THGtools));
    pn.commontools  = [pn.eeg_root, 'T_tools/fnct_common/'];        addpath(genpath(pn.commontools));
    pn.fnct_JQK     = [pn.eeg_root, 'T_tools/fnct_JQK/'];           addpath(genpath(pn.fnct_JQK));
    pn.FT           = [pn.eeg_root, 'T_tools/fieldtrip-20170904/']; addpath(pn.FT); ft_defaults;
    pn.helper       = [pn.eeg_root, 'A_scripts/helper/'];           addpath(pn.helper);
else
    pn.root         = '/home/mpib/LNDG/StateSwitch/WIP_eeg/SA_preproc_study/';
    pn.EEG          = [pn.root, 'B_data/C_EEG_FT/'];
    pn.History      = [pn.root, 'B_data/D_History/'];
end

%% define Condition & IDs for preprocessing

condEEG = 'dynamic';

%% define IDs for segmentation

% N = 47 YAs + 53 OAs;
IDs = {'1117';'1118';'1120';'1124';'1126';'1131';'1132';'1135';'1136';'1138';...
    '1144';'1151';'1158';'1160';'1163';'1164';'1167';'1169';'1172';'1173';...
    '1178';'1182';'1215';'1216';'1219';'1221';'1223';'1227';'1228';'1233';...
    '1234';'1237';'1239';'1240';'1243';'1245';'1247';'1250';'1252';'1257';...
    '1261';'1265';'1266';'1268';'1270';'1276';'1281';...
    '2104';'2107';'2108';'2112';'2118';'2120';'2121';'2123';'2125';'2129';...
    '2130';'2131';'2132';'2133';'2134';'2135';'2139';'2140';'2145';'2147';...
    '2149';'2157';'2160';'2201';'2202';'2203';'2205';'2206';'2209';'2210';...
    '2211';'2213';'2214';'2215';'2216';'2217';'2219';'2222';'2224';'2226';...
    '2227';'2236';'2237';'2238';'2241';'2244';'2246';'2248';'2250';'2251';...
    '2252';'2258';'2261'};

%%  loop IDs

setting.plot = 0; % plot imagesc of resulting timeseries

% prior to loop create matrixes for overviews over artifact-free trials and total trial numbers
overview_nartfree_trials = {};
overview_trial_numbers = {};
ID_unavailable = cell(length(IDs),4);
for id = 1:length(IDs)
    for iRun = 1:4
        
%         if ~exist([pn.EEG, IDs{id}, '_r',num2str(iRun), '_', condEEG, '_EEG_Rlm_Fhl_rdSeg_Art.mat'])
%             display(['processing ID ' IDs{id}, ' Run: ', num2str(iRun)]);
%         else
%             display(['ID ' IDs{id}, ' Run: ', num2str(iRun), ' has already been processed.']);
%             continue;
%         end
        
        try
        %%  load files
            
            % load data
            load([pn.EEG, IDs{id}, '_r',num2str(iRun), '_', condEEG, '_EEG_Rlm_Fhl_rdSeg'],'data_EEG');
            dataICA = load([pn.EEG, IDs{id}, '_', condEEG, '_EEG_Rlm_Fhl_Ica'],'data');
            data = data_EEG; clear data_EEG;
            data.elec = dataICA.data.elec;
            data.chanlocs = dataICA.data.chanlocs; clear dataICA;
            % load config
            load([pn.History, IDs{id}, '_r',num2str(iRun), '_', condEEG, '_config'],'config')

        %% cut data to requested time (necessary due to acquisition jitter)
        % Note that this does not affect the data, but we simply cut some
        % sample points from the ITI which is not of interest anyways.

            trig.eventDurInS = [1.5, 2, 1, 3, 2, 1.5]; % total time to extract is 11s
            trig.eventTriggers = {'S 17'; 'S 72'; 'S 20'; 'S 24'; 'S 64'};
        
            % 2160 has only 62 trials, i.e. the first two trials were
            % not recorded.

            cfg = [];
            cfg.toilim = [-1.5 9.5];
            data = ft_redefinetrial(cfg, data);
            
        %%  define parameter

            % ICs to reject
            iclabels    = config.ica1.iclabels.manual;
            %ics2reject = [iclabels.bli(:); iclabels.mov(:); iclabels.hrt(:); iclabels.ref(:)];
            ics2reject = [iclabels.bli(:); iclabels.mov(:); iclabels.hrt(:); iclabels.ref(:); iclabels.art(:); iclabels.emg(:)];
            clear iclabels

            % trials to remove
            trls2remove = config.ArtDect.trials;

            % channels to interpolate
            chns2interp = config.ArtDect.channels;

        %% for select IDs, interpolate channels prior to ICA
        
        if strcmp(IDs{id}, '2112')
            cfg = [];
            cfg.channel = {'all','-ECG','-A2'};
            dataSelect = ft_preprocessing(cfg,data);
            cfg = [];
            cfg.method     = 'spline';
            if strcmp(IDs{id}, '2112')
                cfg.badchannel = {'Fz'; 'FCz'};
            end
            cfg.trials     = 'all';
            cfg.lambda     = 1e-5; 
            cfg.order      = 4; 
            data = ft_channelrepair(cfg,data);
            clear cfg;
        end
            
        %% ICA (from weights)

            % ica config
            cfg.method           = 'runica';
            cfg.channel          = {'all','-ECG','-A2'};
            cfg.trials           = 'all';
            cfg.numcomponent     = 'all';
            cfg.demean           = 'yes';
            cfg.runica.extended  = 1;

            % ICA solution
            cfg.unmixing     = config.ica1.unmixing;
            cfg.topolabel    = config.ica1.topolabel;

            % components
            comp = ft_componentanalysis(cfg,data);

            % clear cfg
            clear cfg data

        %% remove components

            % cfg for rejecting components (reject: blinks, eye movements, ecg, ref)
            cfg.component = sortrows(ics2reject)';
            cfg.demean    = 'yes';

            % reject components
            data = ft_rejectcomponent(cfg,comp);

            % clear cfg
            clear cfg comp ics2reject

        %%  remove trials

            % IMPORTANT: For 2160, the first two trials were not available,
            % hence we have to add +2 to the index to identify the correct
            % behavioral trials that belong to the EEG data. For 2203 the
            % first trial is missing. The third and second trial should
            % also be removed, regardless of whether they are indicated to
            % be dropped due to automatic correction.
        
            if strcmp(IDs{id}, '2160') && iRun==1
                trls2remove=[trls2remove;1];
            elseif strcmp(IDs{id}, '2203') && iRun==1
                trls2remove=[trls2remove;1];
            end
            
            % define trials to keep
            trials               = 1:length(data.trial);
            trials(trls2remove)  = [];

            if strcmp(IDs{id}, '2160') && iRun==1
                trials=trials+2;
                data.time = cat(2,data.time{end},data.time{end},data.time);
                data.trial = cat(2,data.trial{end},data.trial{end},data.trial);
            elseif strcmp(IDs{id}, '2203') && iRun==1
                trials=trials+1;
                data.time = cat(2,data.time{end},data.time);
                data.trial = cat(2,data.trial{end},data.trial);
            end
            
            % config for deleting trials
            cfg.trials   = trials;

            % remove trials
            data   = ft_preprocessing(cfg,data);
            nartfree_trials = trials;
            overview_nartfree_trials{id,iRun} = nartfree_trials;

            config.nartfree_trials = nartfree_trials;
            
            % clear variables
            clear cfg trials trls2remove

        %% remove eye & reference channels (before interpolation)

            cfg.channel     = {'all','-IOR','-LHEOG','-RHEOG','-A1'};
            cfg.demean      = 'yes';

            % remove channels
            tmpdat = ft_preprocessing(cfg,data);

            % clear cfg
            clear cfg

            % interpolation cfg
            cfg.method     = 'spline';
            cfg.badchannel = tmpdat.label(chns2interp);
            cfg.trials     = 'all';
            cfg.lambda     = 1e-5; 
            cfg.order      = 4; 
            cfg.elec       = config.elec;

            % interpolate
            tmpdat = ft_channelrepair(cfg,tmpdat);

            % clear cfg
            clear cfg chns2interp

        %%  replace interpolated data

            for j = 1:length(data.trial)

                data.trial{j}(1:60,:) = tmpdat.trial{j};

            end; clear j

            % clear temporary data
            clear tmpdat

        %%  keep channel x trial artifacts  

            data.ArtDect.channels_x_trials = config.ArtDect.channels_x_trials;            

        %%  save data
            data_EEG = data; clear data;
            save([pn.EEG, IDs{id}, '_r',num2str(iRun), '_', condEEG, '_EEG_Rlm_Fhl_rdSeg_Art'],'data_EEG')

%         check data
%         data_EEG = THG_label_artefacts_20110916(data_EEG);
%         figure; imagesc(config.ArtDect.channels_x_trials)

            % plot resulting trials

            if setting.plot == 1
                h = figure;
                data = NaN(64,size(data_EEG.trial,1),5540);
                for indTrial = 1:size(data_EEG.trial,2)
                    data(:,indTrial,1:size(data_EEG.trial{indTrial},2))= data_EEG.trial{indTrial};
                    subplot(8,8,indTrial);
                    imagesc(squeeze(data_EEG.trial{indTrial}));
                end; suptitle(['processing ID ' IDs{id}, ' Run: ', num2str(iRun)]);
                close(h);
            end

        %% create overview of percentage of excluded trials, update config.trl

            tmp_trialNumbers(1,1) = length(data_EEG.trial);
            tmp_trialNumbers(2,1) = 64; % maximum amount of trials
            tmp_trialNumbers(3,1) = (100-((tmp_trialNumbers(1,1)/tmp_trialNumbers(2,1))*100)); % percentage of excluded trials
            overview_trial_numbers{id, iRun} = tmp_trialNumbers;
            config.overview_trial_numbers = tmp_trialNumbers;

        %%  update config.trl with trigger codes

            % add info whether trial is included in final set; 1 = trial containing no artifact
            config.trl(:,4)                 = 0;
            config.trl(nartfree_trials, 4)  = 1;
            % add trial number
            config.trl(:,5) = 0;
            config.trl(nartfree_trials, 5)  = nartfree_trials;
            % load marker file
            mrk = config.mrk;
            for i = 1:size(mrk,2)
               mrk_val{i,1} = mrk(1,i).value;
            end
            % add onset and offset triggers
            config.trl(:,6) = 17;
            config.trl(:,7) = 64;
            % save config
            save([pn.History, IDs{id}, '_r',num2str(iRun), '_', condEEG, '_config'],'config')
            
        catch ME
            warning('Error occured. Please check.');
            rethrow(ME)
            fprintf('\n ID not availble! Skip! \n')
            ID_unavailable{id,iRun} = (IDs{id});
        end % try...catch
    end % run
end % id