% Add behavioral information to the EEG data.

% 180123 | written by JQK
% 180327 | including OAs, check for existing data, changed behavioral matrix

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.dynamic_In   = [pn.eeg_root, 'B_data/B_EEG_ET_ByRun/'];
    pn.triggerTiming= [pn.eeg_root, 'C_figures/D_TriggerTiming/'];
    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);
    pn.logs         = [pn.eeg_root, 'Y_logs/H_ArtCorr/'];
    % 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/'];
    pn.triggerTiming= [pn.root, 'C_figures/D_TriggerTiming/'];
    pn.logs         = [pn.root, 'Y_logs/H_ArtCorr/'];
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'};

%% load behavioral data

behavData = load('/Volumes/LNDG/Projects/StateSwitch/dynamic/data/behavior/STSW_dynamic/A_MergeIndividualData/B_data/SS_MergedDynamic_EEG_MRI_YA_09-Mar-2018.mat');

%% extract data by subject

for id = 1:numel(IDs)
    if 1%~exist([pn.EEG, IDs{id}, '_', condEEG, '_EEG_Rlm_Fhl_rdSeg_Art_Cond.mat'])
    for iRun = 1:4

        %% load preprocessed EEG data

        tmp = [];
        tmp.clock = tic; % set tic for processing duration

        load([pn.History, IDs{id}, '_r',num2str(iRun), '_', condEEG, '_config'],'config');
        load([pn.EEG, IDs{id}, '_r',num2str(iRun), '_', condEEG, '_EEG_Rlm_Fhl_rdSeg_Art'],'data_EEG');

        data = data_EEG; clear data_EEG;

        %% remove extraneous channels

        rem             = [];
        rem.channel     = {'all','-IOR','-LHEOG','-RHEOG','-A1'};
        rem.demean      = 'no';
        
        dataNew{iRun} = ft_preprocessing(rem,data); clear rem data;
        
        %% get behavioral data (incl. condition assignment)

        behavIdx = find(strcmp(behavData.IDs_all, IDs{id})); % get index of current subject in behavioral matrix
        SubjBehavData.Atts          = behavData.MergedDataEEG.Atts(:,behavIdx);
        SubjBehavData.StateOrders   = behavData.MergedDataEEG.StateOrders(:,behavIdx);
        SubjBehavData.RTs           = behavData.MergedDataEEG.RTs(:,behavIdx);
        SubjBehavData.Accs          = behavData.MergedDataEEG.Accs(:,behavIdx);
        SubjBehavData.CuedAttributes = reshape(cat(1,behavData.MergedDataEEG.expInfo{behavIdx}.AttCuesRun{:})',256,1);
        
        % create matrix with crucial condition info
        CondData = NaN(256,7);
        for indTrial = 1:256
           tmp_cuedTrial = SubjBehavData.CuedAttributes(indTrial);
           if ismember(1,tmp_cuedTrial{:})
               CondData(indTrial,1) = 1;
           end
           if ismember(2,tmp_cuedTrial{:})
               CondData(indTrial,2) = 1;
           end
           if ismember(3,tmp_cuedTrial{:})
               CondData(indTrial,3) = 1;
           end
           if ismember(4,tmp_cuedTrial{:})
               CondData(indTrial,4) = 1;
           end
        end
        CondData(:,5) = repmat(1:8,1,32)'; % serial position
        CondData(:,6) = SubjBehavData.Atts;
        CondData(:,7) = 1:256;
        CondData(:,8) = SubjBehavData.StateOrders;
        CondData(:,9) = SubjBehavData.RTs;
        CondData(:,10) = SubjBehavData.Accs;
        % select trials from current run
        CurTrials{iRun} = CondData(((iRun-1)*64)+1:iRun*64,:);
        % remove trials that were not retained during preprocessing
        removedTrials = []; removedTrials = setdiff(1:64,config.nartfree_trials);
        CurTrials{iRun}(removedTrials,:) = [];
        
    end; clear CondData;
    %% combine data across runs
    data = ft_appenddata([], dataNew{1}, dataNew{2}, dataNew{3}, dataNew{4});
    data.TrlInfo = cat(1,CurTrials{:});
    data.TrlInfoLabels = {'Att1 cued', 'Att2 cued', 'Att3 cued', 'Att4 cued', ...
        'Serial position', 'WinningAtt', 'Continuous position', 'Cue Dimensionality',...
        'RT', 'Acc'};
    save([pn.EEG, IDs{id}, '_', condEEG, '_EEG_Rlm_Fhl_rdSeg_Art_Cond'],'data');
    else
        disp(['Data for ID ',IDs{id}, ' already exists!']);
    end; % if exist
end