% EEG_PVAF - Compute EEG.data 'percent variance accounted for' (pvaf) by specified components. 
%              Can omit specified components and channels from the computation. Can draw a plot 
%              of the scalp distribution of pvaf, or progressively compute the pvaf for comps
%              1:k, where k = 1 -> the total number of components.  Note: pvaf's of spatially
%              non-orthogonal independent components may not add to 100%, and individual component 
%              pvaf could be < 0%.
% Usage:
%              >> [pv] = eeg_pvaf(EEG,comps);s
%              >> [pvaf,pvafs,vars] = eeg_pvaf(EEG, comps,'key', val);
% Inputs:
%    EEG       - EEGLAB dataset. Must have icaweights, icasphere, icawinv, icaact.
%    comps     - vector of component indices to sum {default|[] -> progressive mode}
%                In progressive mode, comps is first [1], then [1 2], etc. up to
%                [1:size(EEG.icaweights,2)] (all components); here, the plot shows pvaf.
%
% Optional inputs:
%    'artcomps'  - [integer] vector of artifact component indices to remove from data before
%                  computing pvaf {default|[]: none}
%    'omitchans' - [integer] channels to omit from the computation (e.g. off-head, etc) 
%                  {default|[]: none}
%    'chans'     - [integer] only compute pvaf at selected channels. Overwrite omitchans above.
%    'fraction'  - [0<real<=1] fraction of the data to randomly select {default|[]: 1=all}
%    'plot'      - ['on'|'off'] Plot scalp map of channel pvafs. {default: Plot only if no 
%                  output arguments}
%
% Outputs:
%    pvaf      - (real) percent total variance accounted for by the summed back-projection of
%                the requested components. If comps is [], a vector of pvafs for the sum of 
%                components 1:k (k=1:ncomps).
%    pvafs     - (real vector) percent variance accounted for by the summed back-projection of
%                the requested components to each data channel. If comps is [], a matrix of
%                pvafs (as for pv above).
%    vars      - variances of the requested channels
%
% Fields:  
%    Assumes existence of the following EEG fields: EEG.data, EEG.pnts, EEG.nbchan, EEG.trials,
%          EEG.icaact, EEG.icaweights, EEG.icasphere, EEG.icawinv, and for plot, EEG.chanlocs
%
% Author: Scott Makeig & Arnaud Delorme, SCCN, INC, UCSD, Fri Feb 13, 2004

% Copyright (C) 2004- Scott Makeig & Arnaud Delorme, SCCN, INC, UCSD
%
% This file is part of EEGLAB, see http://www.eeglab.org
% for the documentation and details.
%
% Redistribution and use in source and binary forms, with or without
% modification, are permitted provided that the following conditions are met:
%
% 1. Redistributions of source code must retain the above copyright notice,
% this list of conditions and the following disclaimer.
%
% 2. Redistributions in binary form must reproduce the above copyright notice,
% this list of conditions and the following disclaimer in the documentation
% and/or other materials provided with the distribution.
%
% THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
% AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
% IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
% ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
% LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
% CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
% SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
% INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
% CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
% ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF
% THE POSSIBILITY OF SUCH DAMAGE.

function [pvaf,pvafs,pvall] = eeg_pvaf(EEG,comps, varargin)

if nargin < 1
   help eeg_pvaf
   return
end

g = finputcheck(varargin, { 'artcomps'   'integer'    []         [];
                            'omitchans'  'integer'    []         [];
                            'chans'      'integer'    []         [];
                            'fraction'   'real'       []         1;
                            'plot'       'string'     { 'on';'off';'def' } 'def' }, 'eeg_pvaf');
if ischar(g), error(g); end

numcomps = size(EEG.icaact,1);
if round(g.fraction*EEG.pnts*EEG.trials)<1
   error('g.fraction of data specified too small.')
   return
end
if strcmpi(g.plot, 'def')
    if nargout > 0, g.plot = 'on';
    else            g.plot = 'off';
    end
end 

numchans = EEG.nbchan;
chans = 1:numchans;
if ~isempty(g.chans)
    g.omitchans = setdiff_bc([1:EEG.nbchan], g.chans);
end
if ~isempty(g.omitchans)
 if max(g.omitchans)>numchans
  help eeg_pvaf
  error('at least one channel to omit > number of channels in data');
 end
 if min(g.omitchans)<1
  help eeg_pvaf
  error('channel numbers to omit must be > 0');
 end
 chans(g.omitchans) = [];
end

progressive = 0; % by default, progressive mode is off
if nargin < 2 || isempty(comps) || comps(1) == 0
  comps = [];
  progressive = 1;  % turn progressive mode on
end

if isempty(EEG.icaweights)
   help eeg_pvaf
   return
end
if isempty(EEG.icasphere)
   help eeg_pvaf
   return
end

if isempty(EEG.icawinv)
     EEG.icawinv = pinv(EEG.icaweights*EEG.icasphere);
end
if isempty(EEG.icaact)
    help eeg_pvaf
    fprintf('EEG.icaact not present.\n');
    % EEG.icaact = EEG.icaweights*EEG.icasphere*EEG.data; % remake it like this
end
if max(comps) > size(EEG.icawinv,1)
    help eeg_pvaf
   fprintf('Only %d components in this dataset. Cannot project component %d.\n',numcomps,max(comps));
   error('bad comps input');
end
if ~isempty(g.artcomps) && max(g.artcomps) > numcomps
    help eeg_pvaf
   fprintf('Only %d components in this dataset. Cannot project artcomp %d.\n',numcomps,max(g.artcomps));
   error('bad artcomps input')
end

npts = EEG.trials*EEG.pnts;
allcomps = 1:numcomps;
if progressive
   fprintf('Considering components up to: ');
   cum_pvaf  = zeros(1,numcomps);
   cum_pvafs = zeros(numcomps,numchans);
end

for comp = 1:numcomps %%%%%%%%%%%%%%% progressive mode %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
if progressive
   comps = allcomps(1:comp); % summing components 1 to current comp
   fprintf('%d ',comp)
end

if ~isempty(g.artcomps)
   [a b c] = intersect_bc(g.artcomps,comps);
   if ~isempty(a)
      if ~progressive
        if length(a)>1
          fprintf('eeg_pvaf(): not back-projecting %d comps already in the artcomps.\n',length(c));
        else
          fprintf('eeg_pvaf(): not back-projecting comp %d already in the artcomps.\n',comps(c));
        end
      end
      comps(c) = [];
   end
end
if ~isempty(g.artcomps) && min([comps g.artcomps]) < 1
   error('comps and artcomps must contain component indices');
end

%
%%%%%%%%%%%%%%%%%%%%%%%% compute variance accounted for by specified components %%%%%%%%%%%%%
%
if ~progressive || comp == 1 % pare out g.omitchans and artcomps from EEG.data
  if ~isempty(g.artcomps)
    EEG.data = EEG.data(chans,:) - EEG.icawinv(chans,g.artcomps)*EEG.icaact(g.artcomps,:);
  else
    EEG.data = EEG.data(chans,:); 
  end

  nsel = round(g.fraction*npts);
  varpts = randperm(npts);
  varwts = ones(size(varpts));
  if nsel<npts
    varwts(varpts(nsel+1:npts)) = 0;
  end
  pvall = var(EEG.data(:,:)',varwts);
end

pvdiff =  var((EEG.data(:,:) - EEG.icawinv(chans,comps)*EEG.icaact(comps,:))', varwts);

%
%%%%%%%%%%%%%%%%%%%%%%%% compute percent variance accounted for %%%%%%%%%%%%%%%
%
pvafs =  pvdiff ./ pvall;
pvafs = 100-100*pvafs;
pvaf = sum(pvdiff) ./ sum(pvall);
pvaf = 100-100*pvaf;

if ~progressive
  break
else
  cum_pvaf(comp) = pvaf;
  cum_pvafs(comp,:) = pvafs;
end

end %%%%%%%%%%%%%% end progressive forloop %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if progressive % output accumulated results
  fprintf('\n');
  pvaf = cum_pvaf;
  pvafs = cum_pvafs;
  if strcmpi(g.plot, 'on');
     plot(1:numcomps,pvaf);
     xl = xlabel('Components Included (1:n)');
     yl = ylabel('Percent Variance Accounted For (pvaf)');
     set(xl,'fontsize',15);
     set(yl,'fontsize',15);
     set(gca,'fontsize',14);
  end
elseif strcmpi(g.plot, 'on')
%
%%%%%%%%%%%%%%%%%%%%%%%% plot the scalp distribution of pvaf %%%%%%%%%%%%%
%
 if isfield(EEG,'chanlocs')
   chanlocs = EEG.chanlocs;
   if ~isempty(g.omitchans)
     chanlocs(g.omitchans) = [];
   end
   if length(chanlocs) > 1
       topoplot(pvafs',chanlocs);  % plot pvaf here
   end
   
   if length(comps)>5        % add text legend
     if length(g.artcomps)>3
        tlstr=sprintf('Pvaf by %d comps in data minus %d comps',length(comps),length(g.artcomps));
     elseif isempty(g.artcomps)
        tlstr=sprintf('Pvaf by %d comps in data',length(comps));
     elseif length(g.artcomps)==1 % < 4 g.artcomps, list them
        tlstr=sprintf('Pvaf by %d comps in data (less comp ',length(comps));
        tlstr = [tlstr sprintf('%d ',g.artcomps) ')'];
     else
        tlstr=sprintf('Pvaf by %d comps in data (less comps ',length(comps));
        tlstr = [tlstr sprintf('%d ',g.artcomps) ')'];
     end
   else %  < 6 comps, list them
     if length(comps)>1
        tlstr=sprintf('Pvaf by comps ');
     else
        tlstr=sprintf('Pvaf by comp ');
     end
     if length(g.artcomps)>3
        tlstr = ...
[tlstr sprintf('%d ',comps) sprintf('in data minus %d comps',length(comps),length(g.artcomps))];
     else
        if isempty(g.artcomps)
           tlstr = [tlstr sprintf('%d ',comps) 'in data'];
        elseif length(g.artcomps)==1
           tlstr = [tlstr sprintf('%d ',comps) 'in data (less comp '];
           tlstr = [tlstr int2str(g.artcomps) ')'];
        else
           tlstr = [tlstr sprintf('%d ',comps) 'in data (less comps '];
           tlstr = [tlstr sprintf('%d ',g.artcomps) ')'];
        end
     end
   end
   tl=title(tlstr);
   if max(pvafs)>100, 
      maxc=max(pvafs)
   else 
      maxc=100; 
   end

   pvstr=sprintf('Total pvaf: %3.1f%%',pvaf);
   tx=text(-0.9,-0.6,pvstr);

   caxis([-100 100]);
   cb=cbar('vert',33:64,[0 100]); % color bar showing >0 (green->red) only
 else
    fprintf('EEG.chanlocs not found - not plotting scalp pvaf\n');
 end
end % end plot