%% MAIN SCRIPT TO COMPUTE ERS MATRICES

%load preprocessed TF data (bipolar ref, bpf, MW-transformation,
%z-normalized)

%prepare TF features for RSA (sliding window, avg across high freqs (1:30
%31:5:150)

%reduce to OLD trials in encoding and recognition and sort

%compute RSA over channels and freqs

%save ROI-specific channels to disk

%INPUTS:
%subject num
%type of Atlas ('BA' or 'HO')
%idx of areas (Brodmann Area Number or HO label number(s))
    %BA IDX is same as BA numbers 
    %HO IDX is idx of all labels from 'labels_both.mat' in
    %'F:\iEEG_PFC\ECoG_MemDev\ATLAS\HarvardOxford' --> cortical labels +
    %Left and Right Hippocampus
%RSA sliding window parameters
%rsatag - folder naming according to rsa settings
%roitag - folder naming according to roi selection

function EncodingRecognitionSimilarity(subject,channels,atlasType,atlasIdx,rsawidth,rsastep,maxT,rsa_tag,roi_tag) %subject, atlasType ('BA','HO') , winsize, stepsize
    tic
    %% set base path
    fs = '\';
    dataDir = '..\Data';
    ftpath = 'C:\Program Files\MATLAB\R2020b\toolbox\fieldtrip';
    save_to_disk = 1;
    
    %load folder names of subjects
    load('..\subject_list.mat'); %subNames for easier folder accessing
    
    fileName = subNames{subject}; %set folderName of current subject
    disp(fileName)
        
    %% load RSA-prepped tf data for specified RSA settings
    fileTag = ['tfRSA_' num2str(rsawidth) 'width_' num2str(rsastep) 'step_' num2str(maxT) 'ms'];
    load([dataDir fs fileName fs 'tf_data' fs fileTag '_enc.mat']); %encRSA
    load([dataDir fs fileName fs 'tf_data' fs fileTag '_ret.mat']); %retRSA
    
    %% Fetch indices of channels over which to calculate the RSA
    load([dataDir fs fileName fs 'coordinates' fs 'electrodes_bp_atlas_extended.mat'])

    chanIdx = channels;
    chanLabels = electrodes.BA(chanIdx);
    
    %% COMPUTE ENCODING-RECOGNITION SIMILARITY
    %temporal generalization -> all combinations of timepoints
    ERS = zeros(length(encRSA.time),length(encRSA.time),size(encRSA.powspctrm,1),size(encRSA.powspctrm,1));
        
    % prepare trial-wise correlation
    % prepare arrays for linear indices of ALL cells
    %get linear indices of the matrix
    confMat = ones(size(encRSA.powspctrm,1),size(encRSA.powspctrm,1)); %STIM X STIM
    linIdx = find(confMat); %stores linear indices for stimXstim combinations

    [r,c] = ind2sub(size(confMat,1),linIdx);
    linIdx(:,2) = r;
    linIdx(:,3) = c;
    
    %vectorize data and compute similarity analysis
    for stimpair = 1:length(linIdx)
        %reshape array to n Chan*freq by timesteps
        encoding_stim = reshape(squeeze(encRSA.powspctrm(linIdx(stimpair,2),chanIdx,:,:)),[],length(encRSA.time));
        retrieval_stim = reshape(squeeze(retRSA.powspctrm(linIdx(stimpair,3),chanIdx,:,:)),[],length(retRSA.time));

        %vectorizes data with dimension (n chans x freq) X timesteps
        %computes timeXtime correlation matrix for individual stimulus pair!
        ERS(:,:,linIdx(stimpair,2),linIdx(stimpair,3)) = corr(encoding_stim, retrieval_stim,'rows','complete','type','spearman');
    end

    %Fisher Z Transform correlations
    ERS_Z = 0.5.*log((1+ERS)./(1-ERS)); clear ERS;
    ERS = ERS_Z;
    %disp(toc)

    %% create infofile
    info = {};
    info.atlasType = atlasType;
    info.atlasIdx = atlasIdx;
    info.channels = chanIdx;
    info.channel_labels = chanLabels;
    info.trialinfo = encRSA.trialinfo;

    %% save to disk
    %set directory
    statsDir = ['stats' fs 'ERS' fs rsa_tag fs roi_tag];
    
    if save_to_disk
        % save to disk
        % set path
        status = mkdir([dataDir fs fileName fs statsDir]);

        save([dataDir fs fileName fs statsDir fs 'ERS_tempGen.mat'],'ERS')
        save([dataDir fs fileName fs statsDir fs 'ERSinfo.mat'],'info')
    end
    disp(['ERS computations took ' num2str(toc) ' seconds'])
end