%% Prepare TF data for RSA computation
% apply sliding window
% avg over freqs 
% select same items in encoding and recognition
% sort to same order for ERS computation

% do it once for specified RSA settings
% data can then be used for ERS and BetweenItemSimilarity analysis

function RSA_TFfeatures(subject,rsawidth,rsastep,maxT) 
    tic
    %% set base path
    fs = '\';
    dataDir = '..\Data';
    %ftpath = 'C:\Program Files\MATLAB\R2020b\toolbox\fieldtrip';
    
    %load folder names of subjects
    load('..\subject_list.mat'); %subNames for easier folder accessing
    
    fileName = subNames{subject}; %set folderName of current subject
    disp(fileName)
    
    %% set RSA features
    %define time-features of RSA
    all_wins_start = [-0.5:(rsastep/1000):maxT/1000-(rsawidth)/1000]; %./1000 for s as in TF data 
    all_wins_end = [all_wins_start(1)+(rsawidth/1000):rsastep/1000:maxT/1000];
    numWins = length(all_wins_start);
    
    %define freq-features of RSA
    freqs_start = [1:30 31:5:146];
    freqs_end = [1:30 35:5:150];
    numFreqs = length(freqs_start);
    
    %% load TF data (Morlet-Wavelet, no baseline / z-trans)
    %150freq, sr=100Hz
    
    %load TF-data & extended trialinfo;
    load([dataDir fs fileName fs 'tf_data' fs 'tfData150_enc.mat']); enc = data; clear data;
    load([dataDir fs fileName fs 'encoding_trialinfo.mat']); enc.trialinfo = encoding_trialinfo;
    
    load([dataDir fs fileName fs 'tf_data' fs 'tfData150_ret.mat']); ret = data; clear data;
    load([dataDir fs fileName fs 'retrieval_trialinfo.mat']); ret.trialinfo = retrieval_trialinfo;

    %% Apply baseline normalization / z-transformation across trials
    % perform relchange baseline normalization before any further
    % preparation steps
    cfg = [];
    cfg.baselinetype = 'relchange';
    cfg.baseline = [-0.5 -0.2];
    enc = ft_freqbaseline(cfg,enc);
    ret = ft_freqbaseline(cfg,ret);
        
    %% Prep RSA Time-Frequency features (sliding window & avg over frequencies)
    %init empty datastruct with dims of slidingwindow timesteps and num of
    %freqs
    powspctrm_temp_e = zeros(size(enc.powspctrm,1),size(enc.powspctrm,2),numFreqs,numWins);
    powspctrm_temp_r = zeros(size(ret.powspctrm,1),size(ret.powspctrm,2),numFreqs,numWins);
    
    %loop over encoding and retrieval data, avg over frequencies and cut to
    %respective time-windows
    for freq = 1:numFreqs
        freq_start = freqs_start(freq);
        freq_end = freqs_end(freq);
        %disp([num2str(freq_start) ' - ' num2str(freq_end)])
        for win = 1:numWins
            win_start = all_wins_start(win);
            win_end = all_wins_end(win);
            %disp([num2str(win_start) ' - ' num2str(win_end)])
            
            %write freq+time avg data to template structs
            powspctrm_temp_e(:,:,freq,win) = nanmean(nanmean(enc.powspctrm(:,:,[freq_start : freq_end],[nearest(enc.time,win_start):nearest(enc.time,win_end)]),3),4);
            powspctrm_temp_r(:,:,freq,win) = nanmean(nanmean(ret.powspctrm(:,:,[freq_start : freq_end],[nearest(ret.time,win_start):nearest(ret.time,win_end)]),3),4);
        end
    end
    
    %set new RSA TF features as powspctrm, timepoints of sliding window as
    %time and beginning frequency as frequency
    encRSA = enc; encRSA.powspctrm = powspctrm_temp_e; encRSA.time = all_wins_start; encRSA.freq = freqs_start;
    retRSA = ret; retRSA.powspctrm = powspctrm_temp_r; retRSA.time = all_wins_start; retRSA.freq = freqs_start;
    clearvars enc ret;
    
    %% Prepare TF RSA structures for ERS -> same trials, same order, same sorting (memory+category)

    %drop incorrect encoding trials
    drop_trials = find(encRSA.trialinfo(:,1)==0); 
    encRSA.powspctrm(drop_trials,:,:,:) = [];
    encRSA.trialinfo(drop_trials,:) = [];

    %resort encoding data trials to retrieval order
    ret_trnum = encRSA.trialinfo(:,3);
    [~, sort_idx] = sort(ret_trnum);
    encRSA.powspctrm = encRSA.powspctrm(sort_idx,:,:,:);
    encRSA.trialinfo = encRSA.trialinfo(sort_idx,:);

    %drop "new" and non-response trials from recognition data
    drop_trials = find(ismember(retRSA.trialinfo(:,1),[3 4 5 6]));
    retRSA.powspctrm(drop_trials,:,:,:) = [];
    retRSA.trialinfo(drop_trials,:,:,:) = [];

    %get retnum of trials that exist in both phases
    %disp('looking for identical trials in both phases...')
    enctr = encRSA.trialinfo(:,3);
    rettr = retRSA.trialinfo(:,2);

    tr_in_both = intersect(rettr,enctr);
    enc_trials = find(ismember(encRSA.trialinfo(:,3),tr_in_both));
    ret_trials = find(ismember(retRSA.trialinfo(:,2),tr_in_both));

    %keep selection of ENCODING trials
    encRSA.powspctrm = encRSA.powspctrm(enc_trials,:,:,:);
    encRSA.trialinfo = encRSA.trialinfo(enc_trials,:,:,:);

    %keep selection of RECOGNITION trials
    retRSA.powspctrm = retRSA.powspctrm(ret_trials,:,:,:);
    retRSA.trialinfo = retRSA.trialinfo(ret_trials,:,:,:);

    %% Now sort all trials trials by memory (hit vs miss) and image category (in vs out)
    
    %sort all stimuli bei retNum, ImageCategory, Memory
    [~, memcat_idx] = sortrows(encRSA.trialinfo,[2 5]); %same idx for recognition!
    
    %sort trialinfo and powspctrm fields of encoding and recognition data
    encRSA.trialinfo = encRSA.trialinfo(memcat_idx,:);
    encRSA.powspctrm = encRSA.powspctrm(memcat_idx,:,:,:);
    retRSA.trialinfo = retRSA.trialinfo(memcat_idx,:);
    retRSA.powspctrm = retRSA.powspctrm(memcat_idx,:,:,:);
    
    %% save RSA TF features to disk
    fileTag = ['tfRSA_' num2str(rsawidth) 'width_' num2str(rsastep) 'step_' num2str(all_wins_start(1)) '-' num2str(maxT) 'ms'];
    
    save([dataDir fs fileName fs 'tf_data' fs fileTag '_enc.mat'],'encRSA')
    save([dataDir fs fileName fs 'tf_data' fs fileTag '_ret.mat'],'retRSA')
end