%% Time-delay covariances

clearvars -except EEG

% Signal
EEG.data = double(EEG.data);
nbchans  = size(EEG.data,1);
trials   = size(EEG.data,3);
pnts     = size(EEG.data,2);

%channel sorting
chanlocs                = struct2table(EEG.chanlocs);
[chanlocs_sorted,chidx] = sortrows(chanlocs, 'X','descend');
electrodes              = chanlocs_sorted.labels;
[EEG.chanlocs,chanlocs] = deal(table2struct(chanlocs_sorted));
EEG.data = EEG.data (chidx,:,:);

Rcov_time = [-1500 -1300 ]; % R-covariance time 
[~,Ridx(1)] = min(abs(EEG.times-Rcov_time(1)));
[~,Ridx(2)] = min(abs(EEG.times-Rcov_time(2)));

Scov_time = [ 0 200]; % S-covariance time 
[~,Sidx(1)] = min(abs(EEG.times-Scov_time(1)));
[~,Sidx(2)] = min(abs(EEG.times-Scov_time(2)));


time_idx = dsearchn(EEG.times',[0 200]');
% COVARIANCE
cov_R = zeros(nbchans);
cov_S = zeros(nbchans);
 
for triali=1:trials
    dat_R = squeeze(EEG.data(:,Ridx(1):Ridx(2),triali)); %select time window per trial      
    dat_R = dat_R - mean(dat_R,2);
    cov_R = cov_R + dat_R*dat_R' / (size(dat_R,2)-1);
end


for triali=1:trials
   dat_S = squeeze(EEG.data(:,Sidx(1):Sidx(2),triali));    
   dat_S = dat_S - mean(dat_S,2);
   cov_S = cov_S + dat_S*dat_S' / (size(dat_S,2)-1);
end


figure(3),clf
set(gcf, 'color','w', 'WindowState', 'maximize');

subplot(1,2,1)
imagesc(cov_S)
set(gca,"XTick",1:1:64,"XTickLabel",{chanlocs.labels},"YTick",1:1:64,"YTickLabel",{chanlocs.labels},"LineWidth",2)
title(['Covariance S | t=', num2str(Scov_time(1)), '-', num2str(Scov_time(2))], FontSize=32)
colormap(jet)
axis square

subplot(1,2,2)
imagesc(cov_R)
set(gca,"XTick",1:1:64,"XTickLabel",{chanlocs.labels},"YTick",1:1:64,"YTickLabel",{chanlocs.labels},"LineWidth",2)
title(['Covariance R | t=', num2str(Rcov_time(1)), '-', num2str(Rcov_time(2))], FontSize=32)
colormap(jet)
axis square


%% ==== GED ====

% COVARIANCE
cov_R = zeros(nbchans);
for triali=1:trials
    dat_R = squeeze(EEG.data(:,Ridx(1):Ridx(2),triali)); %select time window per trial      
    dat_R = dat_R - mean(dat_R,2);
    cov_R = cov_R + dat_R*dat_R' / (size(dat_R,2)-1);
end 
   cov_R = cov_R ./ EEG.trials;
    
   % apply shrinkage to R  (for numerical stability)
   shr = .01;
   cov_R = (1-shr)*cov_R + shr*mean(eig(cov_R))*eye(size(cov_R));  
 
cov_S = zeros(nbchans);
for triali=1:trials
   dat_S = squeeze(EEG.data(:,Sidx(1):Sidx(2),triali));    
   dat_S = dat_S - mean(dat_S,2);
   cov_S = cov_S + dat_S*dat_S' / (size(dat_S,2)-1);
end
cov_S = cov_S ./ EEG.trials;

% GED
[evecs,evals] = eig( cov_S,cov_R );
[evals,sidx]  = sort( diag(evals),'descend' );
evecs = evecs(:,sidx);

% normalize eigenvectors to unit length
for v = 1:size(evecs,2)
    evecs(:,v) = evecs(:,v)/norm(evecs(:,v)); 
end

nb_comps = 10;
comp_map = zeros(nb_comps,64); 
comp_ts  = zeros(nb_comps,EEG.pnts,EEG.trials);

for compi = 1:nb_comps       
    % compute map and adjust sign
    comp_map(compi,:) = evecs(:,compi)'*cov_S;
    [~,idx]           = max( abs(comp_map(compi,:)) ); % if the largest is negative, then multiply by -1;
    comp_map(compi,:) = comp_map(compi,:) * sign(comp_map(compi,idx));
           
    % compute time series
    tmp_tmp = evecs(:,compi)'*reshape(EEG.data,nbchans,[]);
    comp_ts(compi,:,:) = reshape( tmp_tmp,[pnts trials] );
end 


%%

% Topographical maps:
figure(4),clf
set(gcf, 'color','w', 'WindowState', 'maximize');
for compi = 1:nb_comps

    subplot(2,5,compi)
    topoplot(comp_map(compi,:),chanlocs,'plotrad',0.537,'maplimits', 'absmax','headrad', 'rim','electrodes', 'on', 'style', 'both','whitebk','on');
    title(['Component | number:', num2str(compi)], FontSize=20)
    colormap jet; axis tight


end

% Time-series of each source:
figure(5),clf
set(gcf, 'color','w', 'WindowState', 'maximize');
for compi = 1:nb_comps

    subplot(2,5,compi)
    plot(EEG.times,squeeze(mean(comp_ts(compi,:,:),3)),'linewidth', 3,'color','k','LineStyle','-')
    box off
    title(['Component | number:', num2str(compi)], FontSize=20)

end

%% T-F on source-level

signal = comp_ts(2,:,:);
% Wavelet parameters
sampRate = 512;       
mtime = -2:1/sampRate:2;
htime = (length(mtime)-1)/2;
num_frex    = 30;
min_freq    = 1;
max_freq    = 30;
frex        = logspace(log10(min_freq),log10(max_freq),num_frex); 
rcycles     = [3 10];
ncycles     = logspace(log10(rcycles(1)),log10(rcycles(end)),num_frex);


% Convolution and FFT parameters
kernel  = length(mtime);
ndata   = EEG.pnts*EEG.trials;
nconv   = kernel+ndata-1;
timevec = EEG.times;
% Create a set of wavelets
cmwFFT  = zeros(length(frex),nconv);
for fi=1:num_frex
    s = ncycles(fi)/(2*pi*frex(fi));
    cmw = exp(2*1i*pi*frex(fi).*mtime) .* exp(-mtime.^2./(2*s^2)); %complex morlet wavelet creation

    tmp = fft(cmw,nconv);
    cmwFFT(fi,:) = tmp./max(tmp); %uV normalization units  
end
signal = reshape(signal(1,:,:),1,[]);
dataFFT     = fft(signal,nconv);
baseidx     = dsearchn(timevec',[-1200 -900]');
[tf_result,bs_result] = deal( zeros(size(frex,2),EEG.pnts) );

for j = 1:length(frex)
    % Convolution across frequencies
    as = ifft(cmwFFT(j,:).*dataFFT,nconv);
    as = as(htime+1:end-htime);
    as = reshape(as,EEG.pnts,EEG.trials);
    % Extract power values
    tf_result(j,:,:) = mean(abs(as).^2,2); 
    % Baseline normalization
    tf_baseline = mean(mean(tf_result(j,baseidx(1):baseidx(2),:),2),3);    
    bs_result(j,:,:) = 100*( (tf_result(j,:,:) - tf_baseline )./tf_baseline);
end

% Plot the result
figure; clf;set(gcf, 'color','w', 'WindowState', 'maximize');
contourf(timevec,frex,bs_result,64,'linecolor','none','fill','on'),hold on 
set(gca,'ylim',ylim,'xlim',xlim,'clim',[-150 150],'FontSize',35,...
    'LineWidth',3) 
    box off;colormap (jet);colorbar;axis square