%% 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 pat h + 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'));

% visualize the electrodes, volume conductor model and dipole positions
figure;
ft_plot_mesh(headmodel.bnd(3), 'edgecolor', 'none', 'facecolor', 'r');
ft_plot_mesh(headmodel.bnd(2), 'edgecolor', 'none', 'facecolor', 'g');
ft_plot_mesh(headmodel.bnd(1), 'edgecolor', 'none', 'facecolor', 'b');
ft_plot_mesh(leadfield.pos(leadfield.inside,:),'vertexcolor','m');
ft_plot_sens(elec_aligned);
view([0 -90 0])
alpha 0.3

%% SOURCE ANALYSIS USING LCMV TO LOCALIZE THE N1
tic
for i = 1:nsubj
    datafile = fullfile(path_data, 'data', subjects{i}, 'data.mat');
    
    % load the single epoch data
    load(datafile)

    % remove bad channels
    cfg         = [];
    cfg.channel = {'all','-E30','-Cz','-E192'};
    data        = ft_selectdata(cfg,data);
    
    % filter data
    cfg            = [];
    cfg.bpfilter   = 'yes';
    cfg.bpfilttype = 'firws';
    cfg.bpfreq     = [1 40];
    datafilt       = ft_preprocessing(cfg,data);
    
    % compute covariance
    cfg                  = [];
    cfg.covariance       = 'yes';
    cfg.vartrllength     = 2;
    cfg.covariancewindow = [-.5 .5];
    avgall               = ft_timelockanalysis(cfg,datafilt);
    
    % the data consists of fewer channels than the precomputed
    % leadfields, the following chunk of code takes care of this
    [a,b] = match_str(datafilt.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 lcmv method
    cfg                    = [];
    cfg.channel            = datafilt.label;
    cfg.elec               = elec_aligned;
    cfg.method             = 'lcmv';
    cfg.grid               = leadfield;
    cfg.headmodel          = headmodel;
    cfg.lcmv.keepfilter    = 'yes';
    cfg.lcmv.fixedori      = 'yes';
    cfg.lcmv.projectnoise  = 'yes';
    cfg.lcmv.weightnorm    = 'nai';
    cfg.lcmv.lambda        = '5%';
    sourceall              = ft_sourceanalysis(cfg, avgall);
    
    % average across time the dipole moments within the N1 latency range
    ind    = find(data.time{1}>=0.17 & data.time{1}<=0.23);
    tmpmom = sourceall.avg.mom(sourceall.inside);
    mom    = sourceall.avg.pow(sourceall.inside);
    for ii = 1:length(tmpmom)
        mom(ii) = mean(abs(tmpmom{ii}(ind)));
    end
    
    % insert the N1 amplitude in the 'pow' field and save to disk, the
    % original pow contains the mean amplitude-squared across the
    % time-window used for the channel-level covariance computation
    source = sourceall;
    source.avg.pow(source.inside) = abs(mom);
    source.cfg = rmfield(source.cfg, {'headmodel' 'callinfo'}); % this is removed because it takes up a lot of memory
    save(fullfile(path_data, 'data', subjects{i}, 'source_lcmv.mat'), 'source');
end

toc

% load source reconstructed data
keep subjects path_data nsubj path_ft
all = cell(1,nsubj);
for i = 1:nsubj
    datafile = fullfile(path_data, 'data', subjects{i}, 'source_lcmv.mat');
    load(datafile)
    
    source.cfg = []; % to save memory
    all{i}     = source;
end

% grand average
cfg           = [];
cfg.parameter = 'avg.pow';
gaall         = ft_sourcegrandaverage(cfg, all{:});

% interpolate onto MRI
load(fullfile(path_data,'templates','mri152.mat'));
cfg           = [];
cfg.parameter = 'pow';
sourceNAIint  = ft_sourceinterpolate(cfg, gaall, mri152);

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