% Matlab code included as Supplement A in Smith EE, Bel-Bahar T, Kayser J
%                         (2022, in revision). Psychopysiology.
%
% URL: http://psychophysiology.cpmc.columbia.edu/CSDfcPCA2022supplA.html
%
% Copyright (C) 2022 by Jürgen Kayser (Email: kayser@columbia.edu) and
%                       Ezra E. Smith (Email: ezraesmith@gmail.com)
% GNU General Public License (http://www.gnu.org/licenses/gpl.txt)
% Updated: $Date: 2022/02/02 11:38:00 $ $Author: jk $
%

clear all;
clc;

% MyPath = 'D:\CSD-fcPCA\Matlab\'; % add path with CSDfcPCA functions/files
% addpath(MyPath);
% cd(MyPath);                      % make <MyPath> current working directory
% cd                               % confirm directory  

%%
% EEG = [];
% load('EEG1.mat'); % load EEGlab EEG structure for single subject (session 1)
EEG
            
EEG.times = [EEG.xmin:1/EEG.srate:EEG.xmax]*1000; % timeline for each epoch [ms]

T = tabulate(cell2mat({EEG.event.epochtype}));
table(T(T(:,2)>0,:))

ec_idx = find( cell2mat({EEG.event.epochtype})  ==  20 | ...
               cell2mat({EEG.event.epochtype})  ==  30);
eo_idx = find( cell2mat({EEG.event.epochtype})  == 110 | ...
               cell2mat({EEG.event.epochtype})  == 140);

E = {EEG.chanlocs.labels}';                % get EEG montage and predefined ...
M = ExtractMontage('10-5-System_Mastoids_EGI129.csd',E);  % ... scalp locations
[G,H] = GetGH(M,4);   % create transformation matrices (spline flexibility = 4)
C = CSD(EEG.data,G,H); 

size(EEG.data)    % matrix dimension (channels-by-samples-by-epochs)
size(C)           % matrix dimension (channels-by-samples-by-epochs)

EEG.data = C; % assign CSD-transformed epochs to EEG structure

% save('EEGcsd.mat','EEG'); % save CSD-transformed EEG structure

%%
% load('EEGcsd.mat');

fx_range = [2 50]; % set frequency range …
num_frex = 40;     % … and number of frequencies
frex = logspace(log10(fx_range(1)),log10(fx_range(2)),num_frex); % wavelet ...
s_all = logspace(log10(3),log10(10),num_frex)./(2*pi*frex);      % ... family
dwPLI  = zeros(EEG.nbchan,EEG.nbchan,EEG.trials,num_frex); % initialize FC matrix
time2analyze = [-500 500]; % specify time-frequency parameters [ms]
tidx = dsearchn(EEG.times', time2analyze'); % time to analyze in indices
                                            % setup wavelets
time          = -1:1/EEG.srate:1;           % time [s]
half_wavelet  = (length(time)-1)/2;         % ½-wavelet time points
n_wavelet     = length(time);               % number of wavelet time points
n_data        = EEG.pnts*EEG.trials;        % total number of data points
n_convol      = n_wavelet+n_data-1;         % number of convolution points
WaveletFFT_ = @(nfx,time,n_convol) ...      % box function for wavelet FFT:
    fft(exp(2*1i*pi*frex(nfx).*time) .* ... % nfx = frequency index for ...
    exp(-time.^2./(2*(s_all(nfx)^2))), ...  %       ... center frequency
    n_convol);
sig_ = @(i,A) squeeze(A(i,tidx(1):tidx(2),:)); % box function for data extraction
xsd_ = @(sg1,sg2) sg1 .* conj(sg2);         % box function for complex conjugate

for nfx = 1:num_frex % loop on center frequencies
    tic
    disp(sprintf('%3d/%3d looping. Calculating FC for: %.4f Hz', ...
         nfx,num_frex,frex(nfx))); 
    center_freq = frex(nfx); % center frequency [Hz]
    dwpli = zeros(EEG.nbchan,EEG.nbchan,EEG.trials); % initialize fc matrix
    wavelet_fft = WaveletFFT_(nfx,time,n_convol); % create wavelet for FFT
    ASig = zeros(EEG.nbchan,EEG.pnts,EEG.trials); % initialize signal matrix
    for chani = 1:EEG.nbchan % compute analytic signal for all channels
        data_fft = fft(reshape(EEG.data(chani,:,:),1,n_data), ... % FFT of data
                       n_convol);
        conv_result = ifft(wavelet_fft.*data_fft,n_convol); % convolution
        conv_result = conv_result(half_wavelet+1:end-half_wavelet);
        ASig(chani,:,:) = reshape(conv_result,EEG.pnts,EEG.trials);
    end
    SIG = zeros(half_wavelet+1,size(ASig,3),EEG.nbchan); % initialize SIG matrix
    for ii = 1:EEG.nbchan; SIG(:,:,ii) = sig_(ii,ASig); end; % assign signals
    for chani=1:EEG.nbchan-1 % compute dwPLI for all unique channel pairs
        sig1 = SIG(:,:,chani); % extract data for first channel
        for chanj=chani+1:EEG.nbchan 
            sig2 = SIG(:,:,chanj); % extract data for second channel
            xsd = xsd_(sig1,sig2); % complex conjugate
            cdi = imag(xsd);       % take imaginary part of signal only
            imagsum      = sum(cdi);
            imagsumW     = sum(abs(cdi));
            debiasfactor = sum(cdi.^2);
            dwpli(chani,chanj,:) = ... % fc matrix
                 (imagsum.^2 - debiasfactor)./(imagsumW.^2 - debiasfactor);
        end
    end
    dwPLI(:,:,:,nfx) = dwpli; % assign data for this center frequency to FC matrix
    toc
end

                                           % compute means for each condition
EO = squeeze(mean(dwPLI(:,:,eo_idx,:),3)); % mean FC for eyes open ...
EC = squeeze(mean(dwPLI(:,:,ec_idx,:),3)); % ... and eyes closed ...
EO_odd = squeeze(mean(dwPLI(:,:,eo_idx([1:2:end]),:),3));  % ... EO odd epochs ...
EO_even = squeeze(mean(dwPLI(:,:,eo_idx([2:2:end]),:),3)); % ... EO even epochs ...
EC_odd = squeeze(mean(dwPLI(:,:,ec_idx([1:2:end]),:),3));  % ... EC odd epochs ...
EC_even = squeeze(mean(dwPLI(:,:,ec_idx([2:2:end]),:),3)); % ... EC even epochs

% visualize FC matrix for first frequency (eyes closed, odd)
for nfx = 1
    idx=triu(ones(EEG.nbchan),+1);
    ec_odd = EC_odd(:,:,nfx);
    ec_odd(~idx)=nan;
    figure('name','FC Matrix');
    ax = imagesc(ec_odd,[min(ec_odd,[],'all') max(ec_odd,[],'all')]);
    xlabel('Site','FontWeight','bold');
    ylabel('Site','FontWeight','bold');
    h = colorbar;
    ylabel(h,'dwPLI');
    colormap([1 1 1; jet]);
    title(sprintf('Subject #%d, %s, %.1f Hz',1,'EC odd',frex(nfx)));
end

%%
% re-arrange FC matrix (eyes closed, odd) to edges-by-frequency bins
n_sites = EEG.nbchan; % number of electrode sites
n_edges = (n_sites * (n_sites-1))/2; % number of unique channel pairs
A = zeros(n_edges,num_frex); % initialize matrix
for nfx = 1:num_frex % loop through all frequencies
    ec_odd = EC_odd(:,:,nfx); % get FC matrix for one condition and mark ...
    ec_odd(~idx)=nan;         % ... lower triangle and diagonale as NaN
    d = reshape(ec_odd,n_sites^2,1); % reshape to column vector and ...
    A(:,nfx) = d(~isnan(d));         % ... remove all NaNs
end

B = zeros(n_edges,2*num_frex); % initialize matrix for frequency interpolation
InterpFreqs = logspace(log10(fx_range(1)),log10(fx_range(2)),2*num_frex);
for c = 1:n_edges                              % loop through all edges
    B(c,:) = interp1(frex,A(c,:),InterpFreqs); % interpolate frequency bins
end

fx_range2 = [3 16]; % set new frequency range [Hz]
R = InterpFreqs >= fx_range2(1) & InterpFreqs <= fx_range2(2);
frex2 = InterpFreqs(R); % interpolated frequency bin vector 
B2 = B(:,R); % matrix for interpolated new frequency range

frex2([1:4 end-3:end]) % first and last four interpolated frequency bins
size(B2)               % matrix dimension (edges-by-frequency bins)

%%
% YY = load('PCA solution (Y).mat'); 
% D1 = YY.Mat1;

size(D1)    % dimensions (edges-by-frequency bins) of dwPLI data matrix
disp(sprintf('dwPLI value range: %.3f .. %.3f',min(min(D1)),max(max(D1))));
Y = erpPCA_Varimax4M_progress(D1) % PCA solution for step 1

%%
figure('name','Spectral FC Loadings');
n_factors = size(Y.LR,2); % number of extracted PCA factors
[m,idx] = max(Y.LR); % get maxima of factor loadings with index
clear flab; % get labels: peak loadings frequency (variance explained)
for i = 1:n_factors
    fspec = ['%4.1f Hz (%4.1f%%)'];
    flab(i,:) = {sprintf(fspec,frex2(idx(i)),Y.VT(i,4))};
end;
R = Y.VT(:,4) >= 1;       % plot factor loadings >= 1% explained variance
plot(frex2,Y.LR(:,R'));
xlabel('Frequency [Hz]');
ylabel('dwPLI');
ax = gca;
ax.XLim = [frex2(1)-0.2 frex2(end)+0.2];
title('PCA factor loadings');
legend(flab(R),'location','northwest');

%%
theFact = 1;                     % spectral factor for back-projection ...
disp(sprintf('Factor %d: %s',theFact,flab{theFact}));
U = Y.FSr(:,theFact) * Y.LR(:,theFact)';         % ... to dwPLI unit space
size(U) % matrix dimension (all edges-by-frequency bins)

%%
% YY2 = load('Edge mat, des, U, and Y2.mat');
% U = YY2.U;

n_freqs = size(U,2); % number of frequency bins
n_all = size(U,1)/n_edges; % number of all FC matrices
D2 = zeros(n_all*n_freqs,n_edges); % initialize re-arranged matrix <U>
for i = 1:n_all
    D2((i-1)*n_freqs+1:i*n_freqs,:) = ... % transpose each FC matrix and ...
      U((i-1)*n_edges+1:i*n_edges,:)';    % ... assign to new matrix <D2>
end
Y2 = erpPCA_Varimax4M_progress(D2) % PCA solution for step 2
% Y2 = erpPCA_Varimax4M_progress(D2,'rank',50) % syntax for data subsets

%%
figure('name','Spatial FC Loadings'); 
theFact2 = 1; % spatial factor selected for visualization
subplot(1,2,1); % plot edges
axis square; 
load('chanlocs.mat'); % load EEGlab structure with EEG montage locations 
pct2display = 10; % display top 10% of loadings
load('mod_parula.mat');      % load and assign a color ... 
cmap = colormap(mod_parula); % ... map for plotting
ConnPCA_Plot_Pct_Top_Links(Y2.LR,theFact,n_sites,chanlocs,pct2display,cmap);
cbe = colorbar;
topidx = 1:round([(pct2display/100)*size(Y2.LR,1)]);
set(cbe,'Position',[.46 .46 .04 .2], ... % include scale
        'Ticks',[0 1],'TickLabels',[1 length(topidx)],'FontWeight','bold');
Step1Var = Y.VT(theFact,4);   % variance of step-1 PCA solution 
Step2Var = Y2.VT(theFact2,4); % variance of step-2 PCA solution
Step1x2Var = (Step1Var/100)*(Step2Var/100)*100; % total explained variance
title(sprintf('Edges (top %.0f%%)',pct2display));

subplot(1,2,2); % plot node degree
[B,I] = sort(Y2.LR(:,theFact2),1,'descend'); % identify suprathreshold connections
top_loads = I(topidx);
axis square;
nLR = nan(n_edges,1);
nLR(top_loads,1) = Y2.LR(top_loads,theFact2);
binarizedLR = nLR(:,1);
binarizedLR(~isnan(binarizedLR)) = 1;
adjmat = nan(n_sites,n_sites);
c = 0;
for i = 1:n_sites-1
    for j = i+1:n_sites
        c = c + 1;
        adjmat(i,j) = binarizedLR(c,1);
    end
end
symmat = triu(adjmat) + triu(adjmat)';
node_degree = nansum(symmat);
topoplot_MG(node_degree,chanlocs, 'electrodes','off', 'style','map', ...
            'shading','interp', 'whitebk','on', ...
            'maplimits',[0 max(node_degree)],'colormap',cmap);
cbn = colorbar;
set(cbn,'Position',[.92 .46 .04 .2], ... % include scale
        'Ticks',[0 max(node_degree)],'TickLabels',[1 max(node_degree)], ...
        'FontWeight','bold');
title(sprintf('Nodes (top %.0f%%)',pct2display));
set(gcf,'color','w'); % set figure background to white

MainTitle = sprintf('Step-1 %s × Step-2 #%d (%.1f%%) = %.1f%%', ...
                    flab{theFact},theFact2,Step2Var,Step1x2Var);
annotation('textbox',[.13 .85 .8 .1], 'string',MainTitle, ...
           'HorizontalAlignment','center', 'FontWeight','bold', ...
           'LineSTyle','none', 'FontSize',12);
