function [signal, state] = flt_sourceLocalize(varargin)
% Localize sources for a given head model.
% Estimated sources and inverse operator will be stored in 
% signal.CSD and signal.invOperator
% 
% Author: Tim Mullen, Jan 2013, SCCN/INC/UCSD
%         Alejandro Ojeda, Jan 2013, SCCN/INC/UCSD
%         Christian Kothe, Jan 2013, SCCN/INC/UCSD

% FIXME: Need to add a for loop over the calls to Loreta (or add it in the
% flt_loreta function) which loops over sub-partitions of the data block,
% estimating CSD for each partition (i.e. assuming local stationarity of 
% the inverse operator within each partition). The invOperator results are 
% stored in a 3D array and the mean invOperator is then returned in state
% obj as well as in signal obj.
% FIXME: Implement option to only update inverse operatior infrequently,
% otherwise using the inv op stored in state object (i.e. mean over last
% block).

if ~exp_beginfun('filter'), return; end

declare_properties('name','SourceLocalization', 'experimental',true, 'follows',{'flt_ica','flt_fir','flt_iir'}, 'cannot_precede',{'flt_clean_settings'},'independent_channels',false, 'independent_trials',false);

headmodel_default = 'resources:/headmodels/standard-Colin27-385ch.mat';

% try to get ROI names supported by head model (for use in arg definition)
ROINames = {};
if ~onl_isonline
    hmObj = arg_extract(varargin,{'hmObj','HeadModelObject'},[],headmodel_default);
    if isstruct(hmObj) || (ischar(hmObj) && ~isempty(hmObj))
        hmObj = hlp_validateHeadModelObject(hmObj);
        if ~isempty(hmObj)
            % get the unique ROI names
            [tmp, idx] = unique(hmObj.atlas.label); %#ok<ASGLU>
            ROINames = hmObj.atlas.label(sort(idx))';
        end
    end
end

arg_define(varargin, ...
    arg_norep({'signal','Signal'},[],[],'Signal structure. Must contain .data field with channel data'), ...
    arg({'hmObj','HeadModelObject'},headmodel_default,[],'Head model object generated by MOBILAB. Usually a filename, or a cell array of the form {query,file,query,file,...} where a query is a wildcard pattern that, if matching the sourcedata filename, determines that the subsequent file in the list shall be used  See MOBILAB''s headModel class.','type','expression'), ...
    arg({'channels','Channels'},[],[],'Cell array of channel labels to retain. If empty, all channels in signal.chanlocs.labels will be used. Channels must be a subset of hmObj.getChannelLabels()','type','cellstr','shape','row'), ...
    arg_subswitch({'invMethod','InverseMethod'},{'LORETA'}, ...
    [{...
      'LORETA', @flt_loreta, ...
      'Evidence Approximation', @flt_evidence_approx_tikhonov_reg, ...
      'Variational LORETA', @flt_vbloreta, ...
      'Beamforming', @flt_beamforming}, ...
      quickif(exist('flt_champagne','file'),{'Champagne',@flt_champagne},{})], ...
     'Source Localization Algorithm'), ...
    arg({'sourceAtlasLabels','SourceAtlasLabels','SourceAtlasROI'},ROINames,ROINames,'Source regions of interest (atlas labels). This is a cell array of strings corresponding to a subset of the labels stored in hmObj.atlas.label. Current source density will be estimated only for these ROIs.','type','logical','rangecheck',true), ...
    arg({'colRoiCsd','CollapseRoiCsd'},'sum',{'none','mean','sum','max','maxmag','median'},'Method for computing ROI CSD. Return the (mean, integral, max) current source density within each ROI. signal.srcpot matrix will be reduced to [num_rois x num_samples].'), ...
    arg({'keepFullCsd','KeepFullCsd'},true,[],'Keep a copy of the csd for each vertex. This is stored in srcpot_all'), ...
    arg({'roiAtlasLabels','ROIAtlasLabels','AtlasROI','ROI'},ROINames,ROINames,'Regions of interest atlas labels. This is a cell array of strings corresponding to a subset of the labels stored in hmObj.atlas.label. If CollapseRoiCsd is set, current source density will be integrated over these ROIs. If empty the all ROIs in SourceAtlasLabels are selected.','type','logical','rangecheck',true), ...
    arg({'combineROIs','CombineROI'},[],[],'Combine ROIs in ROIAtlasLabels. This is a cell array of (name, value) pairs. The name entry is a string denoting the name of a new ROI. The value part is a cell array of strings indicating which of the ROIs in ROIAtlasLabels should be combined to form a new ROI. Any ROIs that were merged into a new ROI are removed from the original set of ROIs. To preserve an ROI that was merged into a new ROI, specify a new ROI containing only this ROI. All new ROIs are then appended to the list of existing (unmerged) ROIs in ROIAtlasLabels. If CollapseRoiCsd is set, current source density will be integrated over all ROIs.','type','expression','shape','row'), ...
    arg({'combineMode','CombineMode'},'merge',{'replace','append','merge'},'Mode for combining MetaROIs with original ROIs'), ...
    arg_nogui({'appendROI','AppendROI'},false,[],'Append superROIs or replace'), ...
    arg({'reference','Reference'},[],[],'Reference scheme. This is a cell array of channel labels defining the reference. If more than one channel, average reference is assumed','type','cellstr','shape','row'), ...
    arg_nogui({'roiVertices','ROIVertices'},[],[],'Regions of interest (vertices). This is supplementary to the ''RegionsOfInterest'' option. ''roiVertices'' is a cell array where each cell contains the indices of all vertices within a region of interest. The vertex indices correspond to respective columns of the lead field matrix. These ROIs will be added to the set of ROIs defined in the ''RegionsOfInterest'' option.'), ...
    arg_nogui({'roiVerticesLabels','ROIVerticesLabels'},[],[],'Labels for the supplementary ROIs. This is a cell array of same dimension as ''roiVertices'' with labels for each ROI','type','cellstr'), ...
    arg({'normalCSDWeight','NormalCSDWeight'},0.7,[],'Weight of normal direction in CSD. When calculating the CSD at a given vertex from a lead-field matrix that is not restricted to the surface normal this is the weighting of the current normal to the surface (and tangential current will be weighted by 1 minus this value).','guru',true),...
    arg({'leadfieldOrientation','LeadfieldOrientation'},'axisparallel',{'axisparallel','tangentspace'},'Orientation of local leadfield components. This is for a given source vertex (if orientations are flexible). If axisparallel, the projections will point in X/Y/Z directions, if tangentspace, the projections will point in surface normal, tangent and binormal directions.','guru',true),...
    arg({'makeDipfitStruct','MakeDipfitStruct'},true,[],'Make dipfit structure. If selected, a dipfit structure will be created in signal.dipfit containing the locations (.posxyz) and moments (.mom) of each dipole or center of mass of ROI.'), ... 
    arg({'do_transform','TransformData','transform'},false,[],'Transform the data rather than annotate. By default, source reconstructions are added as annotations (.srcpot) to the data set.'),...
    arg({'usecar','CAR'},true,[],'Rereference to common average. This should always be enabled, unless you are really sure about it'), ...
    arg({'verb','Verbosity'},0,[],'Verbosity level.','typecheck',false), ...
    arg_nogui({'state','State'}));


if isempty(roiAtlasLabels) %#ok
    roiAtlasLabels = ROINames;
end
if isempty(sourceAtlasLabels) %#ok
    sourceAtlasLabels = ROINames;
end


% Intialization block
% -------------------------------------------------------------------------
if isempty(state)
    % || (isfield(state,'roiAtlasLabels') && ~isequal(state.roiAtlasLabels,roiAtlasLabels))
    
    if iscellstr(hmObj)    
        % per-subject headmodels given as a cell array of query-model pairsl resolve which one to use
        loader = utl_find_filter(signal,'io_loadset');
        filename = loader.parts{1};
        for k=1:2:length(hmObj)
            pattern = regexptranslate('wildcard',hmObj{k});
            if regexp(filename,pattern)
                hmObj = hmObj{k+1};
                break;
            end
        end
        if iscellstr(hmObj)
            error('No matching headmodel found for the given data filename: %s (patterns: %s)',filename,hlp_tostring(hmObj(1:2:end))); end
    end
    
    % Either we are initializing our adaptive estimator or user-specified 
    % ROIs have changed.
    state = hlp_microcache('sourcespace',@build_state,hmObj,channels,sourceAtlasLabels,roiAtlasLabels,roiVertices,roiVerticesLabels,makeDipfitStruct,verb,signal.chanlocs,combineROIs,combineMode,usecar,normalCSDWeight,leadfieldOrientation);
end

% Estimate Current Source Density
% -------------------------------------------------------------------------

% Current source density is returned in signal.srcpot 
% [num_vertices x num_samples] 
% Inverse operator is returned in signal.srcweights   
% [num_vertices x num_channels]

if usecar
    % apply common average reference
    signal.data = bsxfun(@minus,signal.data,mean(signal.data));
end
switch lower(invMethod.arg_selection)
    case 'loreta'
        [signal, state.solverState] = hlp_scope({'disable_expressions',true}, ...
                                    @flt_loreta,                    ...
                                    'signal',signal,                ...
                                    invMethod,                      ...
                                    'state',state.solverState,      ...
                                    'K',state.leadFieldMatrix,      ...
                                    'L',state.laplacianOperator,    ...
                                    'verb',max(0,verb-1),'arg_direct',true);
        signal.srcweights = state.solverState.srcweights;
    case 'evidence approximation'
        [signal, state.solverState] = hlp_scope({'disable_expressions',true}, ...
                                    @flt_evidence_approx_tikhonov_reg,                  ...
                                    'signal',signal,                ...
                                    invMethod,                      ...
                                    'state',state.solverState,      ...
                                    'K',state.leadFieldMatrix,      ...
                                    'L',state.laplacianOperator,    ...
                                    'verb',verb,'arg_direct',true);
        signal.srcweights = state.solverState.srcweights;
    case 'variational loreta'
        [signal, state.solverState] = hlp_scope({'disable_expressions',true}, ...
                                    @flt_vbloreta,                  ...
                                    'signal',signal,                ...
                                    invMethod,                      ...
                                    'state',state.solverState,      ...
                                    'K',state.leadFieldMatrix,      ...
                                    'L',state.laplacianOperator,    ...
                                    'verb',max(0,verb-1),'arg_direct',true);
        signal.srcweights = state.solverState.srcweights;
    case 'dynamic loreta'
        [signal, state.solverState] = hlp_scope({'disable_expressions',true}, ...
                                    @flt_dynamicloreta,             ...
                                    'signal',signal,                ...
                                    invMethod,                      ...
                                    'state',state.solverState,      ...
                                    'K',state.leadFieldMatrix,      ...
                                    'L',state.laplacianOperator,    ...
                                    'verb',max(0,verb-1),'arg_direct',true);
        signal.srcweights = state.solverState.srcweights;
    case 'beamforming'
        [signal, state.solverState] = hlp_scope({'disable_expressions',true}, ...
                                    @flt_beamforming,                  ...
                                    'signal',signal,                ...                                    
                                    invMethod, ...
                                    'lead_field', state.leadFieldMatrix, ...
                                    'state',state.solverState,      ...
                                    'arg_direct',true);
    case 'champagne'
        [signal, state.solverState] = hlp_scope({'disable_expressions',true}, ...
                                    @flt_champagne,                  ...
                                    'signal',signal,                ...
                                    invMethod,                      ...
                                    'state',state.solverState,      ...
                                    'head_file',struct('leadfield',reshape(state.leadFieldMatrix,state.nChannels,state.nVertices,state.orientationDegrees), ...
                                                       'chanlabels',{{signal.chanlocs.labels}}), ...
                                    'arg_direct',true);
                                
    otherwise
        error('flt_sourceLocalize:badInverseMethod','Invalid inverse method %s',invMethod.arg_selection);
end


% collapse srcpot from 3d to 1d based on state.srcweightCollapseDoFWeights
if state.orientationDegrees == 3
    siz = size(signal.srcpot);
    signal.srcpot = reshape(sum(reshape(bsxfun(@times,signal.srcpot,state.srcweightCollapseDoFWeights),[3 siz(1)/3 siz(2:end)]),1),[siz(1)/3 siz(2:end)]);
end

% Collapse CSD within each ROI
% -------------------------------------------------------------------------
% FIXME: We may wish to replace avg with the surface integral over the ROI
% We can use numerical double integration i.e. quad2d() or more crudely:
% sum(signal.CSD(x,:)) * dx * dy;
% where dx, dy are the voxel x,y dimensions.
% If the csd is not constrained to the surface of a mesh, then we need to
% use triple integration i.e. triplequad(). Or
% sum(signal.CSD(x,:)) * dx * dy * dz

if ~strcmp(colRoiCsd,'none') && ~isempty(signal.srcpot)
    if verb
        fprintf('Computing %s CSD for each ROI \n', colRoiCsd); 
    end
    
    if keepFullCsd
        % store backup with all vertices
        % also, register this as a time-series field
        signal = utl_register_field(signal,'timeseries','srcpot_all',signal.srcpot);
        signal.srcweights_all = signal.srcweights;
    end
        
    % collapse current density
    signal.srcpot = hlp_colsrc(signal.srcpot,state.roiVerticesReduced,colRoiCsd);
    % collapse weights (note: if weights are time-varying, then whatever snapshot this is isn't
    % going to be too useful)
    if isfield(signal,'srcweights')
        signal.srcweights = hlp_colsrc(signal.srcweights,state.roiVerticesReduced,colRoiCsd); end
else
    signal.srcpot_all       = [];
    signal.srcweights_all   = [];
end


% Store additional outputs in signal
% -------------------------------------------------------------------------
if makeDipfitStruct
    signal.dipfit   = state.dipfit;
end
signal.roiLabels          = state.roiLabels;
signal.roiVertices        = state.roiVertices;
signal.roiVerticesReduced = state.roiVerticesReduced;
signal.rmIndices          = state.rmIndices;

if do_transform
    signal.data = signal.srcpot;
    signal.nbchan = size(signal.data,1);
    signal.chanlocs = struct('labels',cellfun(@num2str,num2cell(1:signal.nbchan,1),'UniformOutput',false));
end
signal.hmObj = hmObj;
signal.leadFieldMatrix = state.leadFieldMatrix;

    
exp_endfun;


function state = build_state(hmObj,channels,sourceAtlasLabels,roiAtlasLabels,roiVertices,roiVerticesLabels,makeDipfitStruct,verb,chanlocs,combineROIs,combineMode,usecar,normalCSDWeight,leadfieldOrientation)
% validate the head model and construct the source space

orilen = length(roiAtlasLabels);
[res ia] = intersect(roiAtlasLabels,sourceAtlasLabels);
roiAtlasLabels = roiAtlasLabels(sort(ia));  % FIXME: double-check this order
if length(roiAtlasLabels)~=orilen
    warn_once('BCILAB:AtlasLabelMismatch','Some elements of ROIAtlasLabels were not contained in SourceAtlasLabels. These were removed.');
end

% Validate Inputs
hmObj     = hlp_validateHeadModelObject(hmObj);
if isempty(hmObj)
    error('HeadModelObject was improperly defined. Exiting');
end
if ~isempty(roiVerticesLabels) && ~isempty(roiVertices) ...
        && length(roiVertices) ~= length(roiVerticesLabels)
    fprintf('WARNING: The number of ROI labels in ''roiVerticesLabels'' does not match the number of ROIs in ''roiVertices''. Reverting to generic labeling scheme {''ROI1'' ''ROI2'' ... ''ROIN''}.');
    roiVerticesLabels = [];
end


if verb
    fprintf('Initializing source reconstructor \n');
end

% Prune lead field matrix and laplacian operator
% to contain only regions of interest
% If there are no specific ROIs selected, this will "open" the
% surface mesh by removing a few vertices
% (the surface mesh cannot be a closed surface)
if verb
    fprintf('Constructing source space \n');
end

% remove undesired brain structures 
% note: we assume that the leadFieldMatrix produced by this step is 2d, but
% contain stacked pages for X, Y and Z dipole components (so that a reshape
% to [#ch,#vox,3] would yield the correct result.
brainStructsToRemove = setdiff(unique(hmObj.atlas.label),sourceAtlasLabels);
[   reducedSpace,                      ...
    state.leadFieldMatrix,             ...
    state.laplacianOperator            ...
    state.rmIndices                    ...
    ] = getSourceSpace4PEB(hmObj,brainStructsToRemove,roiVertices);

state.nVertices = size(reducedSpace.vertices,1);
state.nChannels = length(chanlocs);

% determine the number of orientation degrees of freedom
if size(state.leadFieldMatrix,2) == state.nVertices
    state.orientationDegrees = 1;
elseif size(state.leadFieldMatrix,2) == 3*state.nVertices
    state.orientationDegrees = 3;
else
    error('Your lead-field matrix has a # of columns that does not correspond to the # of source vertices.');
end

% project leadfield matrix onto tangent space
if state.orientationDegrees==3
    % get tangent space at each vertex
    if isprop(hmObj,'surfNormal')
        state.surfNormals = blk_diag(hmObj.surfNormal,1)';
        % prune removed indices
        state.surfNormals(state.rmIndices,:)=[];
    else
        % FIXME: these data contains NaN's (presumably near deleted triangles)
        state.surfNormals = geometricTools.getSurfaceNormals(reducedSpace.vertices,reducedSpace.faces,true);
    end
    state.surfTangents = cross(state.surfNormals,ones(length(state.surfNormals),1)*[1 0 0],2);
    state.surfBinormals = cross(state.surfNormals,state.surfTangents,2);    

    % also calc final projection operator
    if strcmp(leadfieldOrientation,'axisparallel')
        state.srcweightCollapseDoFWeights = vec((state.surfNormals*normalCSDWeight + (state.surfTangents+state.surfBinormals)*(1-normalCSDWeight))'); % does assume axis-parallel lead-field    
    elseif strcmp(leadfieldOrientation,'tangentspace')
        state.srcweightCollapseDoFWeights = vec((ones(length(state.surfNormals),1)*[normalCSDWeight (1-normalCSDWeight)/2 (1-normalCSDWeight)/2])');  % assumes rotated lead-field
        % turn into triplet order
        tmp = permute(reshape(state.leadFieldMatrix,state.nChannels,state.nVertices,state.orientationDegrees),[1 3 2]);
        % transpose for recombination
        tmp = tmp(:,:)'; siz = size(tmp);
        % recombine into tangent space in pages
        new_matrix = permute(cat(3,reshape(sum(reshape(bsxfun(@times,tmp,vec(state.surfNormals')),[3 siz(1)/3 siz(2:end)]),1),[siz(1)/3 siz(2:end)]), ...
            reshape(sum(reshape(bsxfun(@times,tmp,vec(state.surfTangents')),[3 siz(1)/3 siz(2:end)]),1),[siz(1)/3 siz(2:end)]), ...
            reshape(sum(reshape(bsxfun(@times,tmp,vec(state.surfBinormals')),[3 siz(1)/3 siz(2:end)]),1),[siz(1)/3 siz(2:end)])),[2 1 3]);
        % flatten again
        state.leadFieldMatrix = new_matrix(:,:);
    else
        error('Unsupported LeadfieldOrientation specified: %s',leadfieldOrientation);
    end
end

% We also store the indices of the vertices of each ROI (in the full
% source space) in a cell array. This allows us to obtain dipole
% centroids for each ROI
for k=1:length(roiAtlasLabels)
    state.roiVertices{k} = find(indices4Structure(hmObj,roiAtlasLabels{k}))';
    %         [~,state.roiVertices{k}] = removeStructureFromSourceSpace(hmObj,state.brainStructsToRemove{k});
end
% append any additional ROI defined by 'roiVertices' input argument
state.roiVertices = [state.roiVertices roiVertices];

% get ROI indices into reduced source space. These are used for
% integration over current source density (CSD) within each ROI
LFMcols = 1:length(hmObj.atlas.colorTable);
LFMcols(state.rmIndices) = [];
for k=1:length(state.roiVertices)
    state.roiVerticesReduced{k} = find(ismember(LFMcols,state.roiVertices{k}));
end


% set up the labeling for any vertex-index-defined ROIs
if isempty(roiVerticesLabels) && ~isempty(roiVertices)
    % use generic labeling scheme {'ROI1' 'ROI2' ... 'ROIN'}
    roiVerticesLabels = cellstr(num2str((1:length(roiVertices))'))';
    roiVerticesLabels = cellfun(@(x)['ROI' x],roiVerticesLabels,'UniformOutput',false);
end
state.roiLabels = [roiAtlasLabels roiVerticesLabels];


% if needed, combine ROIs according to rules in combineROIs
if ~isempty(combineROIs)
    if mod(length(combineROIs),2)
        error('BCILAB:BadNameValuePair','CombineROI must be a cell array of name,value pairs');
    end
    newRoiLabels = combineROIs(1:2:end-1);
    roiSubLabels = combineROIs(2:2:end);
    numroi = length(newRoiLabels);
    newRoiVerticesReduced = cell(1,numroi);
    newRoiVertices        = cell(1,numroi);
    roiIdxInAtlas  = cell(1,numroi);
    for k=1:numroi
        matchedLabels = ismember(roiSubLabels{k},roiAtlasLabels);
        if ~all(matchedLabels)
            error('BCILAB:BadROINames', ...
                  'The following ROIs in ''CombineROIs'' are not in the ROIAtlasLabels list: %s', ...
                          hlp_tostring(roiSubLabels{k}(~matchedLabels)));
        end
        roiIdxInAtlas{k}  = find(ismember(roiAtlasLabels,roiSubLabels{k}));
        A = cell2mat(state.roiVerticesReduced(roiIdxInAtlas{k}));
        [tmp, idx] = unique(A);
        newRoiVerticesReduced{k} = A(sort(idx));
        % do the same for roiVertices
        A = cell2mat(state.roiVertices(roiIdxInAtlas{k}));
        [tmp, idx] = unique(A);
        newRoiVertices{k} = A(sort(idx));
    end
    
    switch combineMode
        case 'merge'
            % remove all merged ROIs ...
            mergedROI = unique(cell2mat(roiIdxInAtlas));
            state.roiVerticesReduced(mergedROI) = [];
            state.roiVertices(mergedROI)        = [];
            state.roiLabels(mergedROI)          = [];
            
            % ... and append the new ROIs
            state.roiVerticesReduced = [state.roiVerticesReduced newRoiVerticesReduced];
            state.roiVertices        = [state.roiVertices newRoiVertices];
            state.roiLabels          = [state.roiLabels newRoiLabels];
        case 'replace'
            % keep only the new rois
            state.roiVerticesReduced = newRoiVerticesReduced;
            state.roiVertices        = newRoiVertices;
            state.roiLabels          = newRoiLabels;
        case 'append'
            % append the new ROIs
            state.roiVerticesReduced = [state.roiVerticesReduced newRoiVerticesReduced];
            state.roiVertices        = [state.roiVertices newRoiVertices];
            state.roiLabels          = [state.roiLabels newRoiLabels];
    end
end


% prune the lead field matrix to contain only desired channels
hmChanlabels  = lower(hmObj.channelLabel);
if isempty(channels)
    channels = lower({chanlocs.labels});
end

% use only selected channels that are in the head model
% chaninds is also an integer permutation vector so that rows of 
% LFM match the ordering of 'channels' variable
[~, ia, ib] = intersect(lower(hmChanlabels),lower(channels),'stable');
if max(ib)>length(ia)
    error('The following channels could not be matched to the headmodel: %s', ...
          hlp_tostring(channels(ib>length(ia))));
end
chaninds = ia; %ia(ib);
% if nnz(chaninds)~=length(channels)
%     error('Some channels could not be matched to the headmodel');
% end
state.leadFieldMatrix = state.leadFieldMatrix(chaninds,:,:);

if usecar
    % apply common-average rereference
    state.leadFieldMatrix = bsxfun(@minus,state.leadFieldMatrix,mean(state.leadFieldMatrix));
end

% load the original source space (non-reduced)
tmp = load(hmObj.surfaces);
fn  = fieldnames(tmp);
state.sourceSpace = tmp.(fn{1})(end);  % dim end = cortical surface

% make dipfit structure containing centroids (dipfit.model.posxyz) and
% surface mesh (dipfit.model.surfmesh) for each ROI as well as complete
% surface mesh (dipfit.surfmesh). We also store the indices of each ROI
% into the complete surface mesh (dipfit.model.meshVertices)
if makeDipfitStruct
    state.dipfit = hlp_makeDipfitStruct(state.sourceSpace,state.roiVertices,reducedSpace);
end

% initialize the state of the solver
state.solverState = [];
