%% SOME preLIMINARIES
clear 
close all
clc

% initialize fieldtrip
path_ft   = '/path/to/fieldtrip';
path_data = '/path/to/data/ASSR';

% this adds fieldtrip to the matlab path + the relevant subdirectories
addpath(path_ft);
ft_defaults;

%% the variable subjects contains the list of subject-specific directories
subjects = {'45048'	'47723'	'49449'	'48288'	'45670'...
            '49353'	'47522'	'44340'	'45880'	'45611'...
            '46658'	'48993'	'45859'	'45644'	'47790'...
            '45549'	'48411'	'47371'	'48135'	'46637'...
            '45631'	'47498'	'44296'	'49121'	'46659'...
            '49152'	'44212'	'49004'	'45249'};
nsubj = numel(subjects);

% this loads in the precomputed geometric data that is needed for source
% reconstruction: the forward model (leadfields), as well as the volume
% conductor model of the head, and the electrodes. Strictly speaking, once
% the forward model is available, the other objects are not needed anymore,
% but FieldTrip expects them to be present in the cfg. 
load(fullfile(path_data,'templates','headmodel.mat'));
load(fullfile(path_data,'templates','leadfield.mat'));
load(fullfile(path_data,'templates','elec_aligned.mat'));

% the ASSR is a very high SNR signal, leading to (near) simultaneous (i.e.
% zero-lag correlated) activation in bilateral auditory areas. Therefore,
% the traditional beamformer will fail to yield good results. Instead, we
% will use a 'double-dipole' as sourcemodel, using left/right symmetric
% dipole pairs. In this example here, we have already computed leadfields,
% so we will create for each original dipole location a leadfield
% consisting of a concatenation of this dipole's leadfield with its
% counterpart in the opposite hemisphere. Note that the x-coordinate here
% is the left/right direction.
lf1 = reshape(leadfield.leadfield, leadfield.dim);
lf2 = flip(lf1,1);
for k = 1:numel(lf1)
    if ~isempty(lf1{k})&&~isempty(lf2{k})
        leadfield.leadfield{k} = [lf1{k} lf2{k}];
    else
        leadfield.leadfield{k} = [];
        leadfield.inside(k)    = false;
    end
end
clear lf1 lf2

%% SOURCE ANALYSIS USING DICS
tic
for i = 1:nsubj
    datafile = fullfile(path_data, 'data', subjects{i}, 'data.mat');
    load(datafile)
   
    % remove bad channels
    cfg         = [];
    cfg.channel = {'all' '-E30' '-E3' '-E4' '-E11' '-E12' '-Cz' '-E192'};
    data        = ft_selectdata(cfg, data);
    
    % this chunk of code creates a 'dummy' reference channel to be used for
    % the coherence analysis
    trial = cell(size(data.trial));
    for k = 1:numel(trial)
        trial{k} = sin(2.*pi.*40.*data.time{k});
    end
    refdata.trial = trial;
    refdata.time  = data.time;
    refdata.label = {'refchan'};
    data          = ft_appenddata([], data, refdata);
    data.fsample  = 1000;
    
    % re-segment the data into pre and post stimulus onset intervals
    cfg        = [];                                              
    cfg.toilim = [-1 0-1./data.fsample];    
    datapre    = ft_redefinetrial(cfg, data);
      
    cfg.toilim = [1 2-1./data.fsample]; 
    datapost   = ft_redefinetrial(cfg, data);

    % append data to facilitate the computation of a 'common' spatial filter 
    dataall = ft_appenddata([], datapre, datapost);
    
    % calculate the cross-spectral density matrics
    cfg        = [];
    cfg.method = 'mtmfft';
    cfg.output = 'fourier';
    cfg.taper  = 'hanning';
    cfg.foilim = [40 40];
    freqpre    = ft_freqanalysis(cfg, datapre);
    freqpost   = ft_freqanalysis(cfg, datapost);
    freqall    = ft_freqanalysis(cfg, dataall);
    
    % the data consists of fewer channels than the precomputed
    % leadfields, the following chunk of code takes care of this
    [a,b] = match_str(data.label, leadfield.label);
    for k = 1:numel(leadfield.leadfield)
        if ~isempty(leadfield.leadfield{k})
            tmp = leadfield.leadfield{k};
            tmp = tmp(b,:);
            tmp = tmp-repmat(mean(tmp,1),[size(tmp,1) 1]); % average re-ref
            leadfield.leadfield{k} = tmp;
        end
    end
    leadfield.label = leadfield.label(b);
    
    % perform source reconstruction using the dics method
    cfg                   = []; 
    cfg.method            = 'dics';
    cfg.frequency         = 40;  
    cfg.channel           = data.label;
    cfg.elec              = elec_aligned;
    cfg.grid              = leadfield;
    cfg.headmodel         = headmodel;
    cfg.dics.keepfilter   = 'yes';
    cfg.dics.fixedori     = 'no';
    cfg.dics.projectnoise = 'yes';
    cfg.dics.lambda       = '5%';
    cfg.dics.realfilter   = 'yes';
    cfg.dics.keepcsd      = 'yes';
    cfg.refchan           = {'refchan'};
    sourceall             = ft_sourceanalysis(cfg, freqall);
    
    % apply common filters to pre and post stimulus data
    cfg.grid.filter = sourceall.avg.filter;
    
    % now we need to extract the dipole pairs' full csd matrix with respect
    % to the reference channel, which is not possible in the traditional
    % DICS implementation, but can be achieved with pcc
    cfg.method   = 'pcc';
    cfg          = rmfield(cfg, 'dics');
    sourcepre    = ft_sourceanalysis(cfg, freqpre);
    sourcepst    = ft_sourceanalysis(cfg, freqpost);
    
    % now we need to adjust the 'csdlabel', in order to compute the single
    % dipole-of-interest coherence with the reference.
    % source.avg.csd contains for each dipole
    % pair the cross-spectral density as a 7x7 matrix, two triplets of
    % rows/columns for each of the dipoles in the pair, and one for the
    % reference signal. ft_sourcedescriptives will return the coherence 
    % between the 'scandip' and the 'refchan', so we need to relabel the
    % second dipole's label
    for k = 1:size(sourcepre.pos,1)
      if ~isempty(sourcepre.avg.csdlabel{k})
        sourcepre.avg.csdlabel{k}(4:6) = {'supdip'};
        sourcepst.avg.csdlabel{k}(4:6) = {'supdip'};
      end
    end
    sourcepre = ft_sourcedescriptives([], sourcepre);
    sourcepst = ft_sourcedescriptives([], sourcepst);
    
    % contrast post stimulus onset activity with respect to baseline
    cfg           = [];
    cfg.operation = 'x2-x1';
    cfg.parameter = 'avg.coh';
    source_dics   = ft_math(cfg,sourcepre,sourcepst);
    

    save(fullfile(path_data, 'data/', subjects{i}, 'source_dics.mat'), 'source_dics');
end
toc

tic
allsource = cell(1,nsubj);
for i = 1:nsubj
    datafile = fullfile(path_data,'data', subjects{i}, 'source_dics.mat');
    load(datafile)
    allsource{i} = source_dics;
end

% grand average
cfg           = [];
cfg.parameter = 'coh';
gaall         = ft_sourcegrandaverage(cfg, allsource{:});

% interpolate onto MRI
load(fullfile(path_data,'templates','mri152.mat'));

cfg           = [];
cfg.parameter = 'coh';
sourceCOHint  = ft_sourceinterpolate(cfg, gaall, mri152);

% plot 
aal = ft_read_atlas(fullfile(path_ft, 'template/atlas/aal/ROI_MNI_V4.nii'));% load the AAL atlas
sourceCOHint.mask =( sourceCOHint.coh > max(sourceCOHint.coh(:))*.95 |...
                     sourceCOHint.coh < min(sourceCOHint.coh(:))*.9);
cfg               = [];
cfg.method        = 'ortho';
cfg.funparameter  = 'coh';
cfg.funcolormap   = 'jet';
cfg.maskparameter = 'mask';
cfg.funcolormap   = parula;
cfg.atlas         = aal;
cfg.location      = [55.5 -22.5 10.5];
ft_sourceplot(cfg,sourceCOHint);
toc
