function [input] = ft_apply_montage(input, montage, varargin)

% FT_APPLY_MONTAGE changes the montage (i.e. linear combination) of a set of
% electrode or gradiometer channels. A montage can be used for EEG rereferencing, MEG
% synthetic gradients, MEG planar gradients or unmixing using ICA. This function not
% only applies the montage to the EEG or MEG data, but also applies the montage to
% the input EEG or MEG sensor array, which can subsequently be used for forward
% computation and source reconstruction of the data.
%
% Use as
%   [sens]    = ft_apply_montage(sens,     montage,  ...)
%   [data]    = ft_apply_montage(data,     montage,  ...)
%   [freq]    = ft_apply_montage(freq,     montage,  ...)
%   [montage] = ft_apply_montage(montage1, montage2, ...)
%
% A montage is specified as a structure with the fields
%   montage.tra      = MxN matrix
%   montage.labelold = Nx1 cell-array
%   montage.labelnew = Mx1 cell-array
%
% As an example, a bipolar montage could look like this
%   bipolar.labelold  = {'1',   '2',   '3',   '4'}
%   bipolar.labelnew  = {'1-2', '2-3', '3-4'}
%   bipolar.tra       = [
%     +1 -1  0  0
%      0 +1 -1  0
%      0  0 +1 -1
%   ];
%
% The montage can optionally also specify the channel type and unit of the input
% and output data with
%   montage.chantypeold = Nx1 cell-array
%   montage.chantypenew = Mx1 cell-array
%   montage.chanunitold = Nx1 cell-array
%   montage.chanunitnew = Mx1 cell-array
%
% Additional options should be specified in key-value pairs and can be
%   'keepunused'    = string, 'yes' or 'no' (default = 'no')
%   'feedback'      = string, see FT_PROGRESS (default = 'text')
%   'warning'       = boolean, whether to show warnings (default = true)
%
% If the first input is a montage, then the second input montage will be
% applied to the first. In effect, the output montage will first do
% montage1, then montage2.
%
% See also FT_READ_SENS, FT_DATATYPE_SENS

% Please note that the options 'inverse' and 'balancename' have been
% removed with https://github.com/fieldtrip/fieldtrip/pull/2529 in favour
% of a more explicit construction of the inverse and of the bookkeeping.

% Copyright (C) 2008-2025, Robert Oostenveld
%
% This file is part of FieldTrip, see http://www.fieldtriptoolbox.org
% for the documentation and details.
%
%    FieldTrip is free software: you can redistribute it and/or modify
%    it under the terms of the GNU General Public License as published by
%    the Free Software Foundation, either version 3 of the License, or
%    (at your option) any later version.
%
%    FieldTrip is distributed in the hope that it will be useful,
%    but WITHOUT ANY WARRANTY; without even the implied warranty of
%    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%    GNU General Public License for more details.
%
%    You should have received a copy of the GNU General Public License
%    along with FieldTrip. If not, see <http://www.gnu.org/licenses/>.
%
% $Id$

if iscell(input)
  % this represents combined EEG, ECoG and/or MEG
  for i=1:numel(input)
    input{i} = ft_apply_montage(input{i}, montage, varargin{:});
  end
  return
end

% use "old/new" instead of "org/new"
montage = fixoldorg(montage);
input   = fixoldorg(input); % the input might be a montage or a sensor array

% get optional input arguments
keepunused    = ft_getopt(varargin, 'keepunused',  'no');
feedback      = ft_getopt(varargin, 'feedback',    'text');
showwarning   = ft_getopt(varargin, 'warning',     true);

% warn users about backward incompatible changes
if any(strcmp(varargin(1:2:end), 'inverse'))
  error('the option "inverse" is not supported any more, see FT_INVERSE_MONTAGE');
end
if any(strcmp(varargin(1:2:end), 'montagename'))
  error('the option "montagename" is not supported any more');
end

if istrue(showwarning)
  warningfun = @warning;
else
  warningfun = @nowarning;
end

% ensure that the input montage is correct, see https://github.com/fieldtrip/fieldtrip/issues/1718
assert(length(unique(montage.labelold))==length(montage.labelold), 'the montage is invalid');
assert(length(unique(montage.labelnew))==length(montage.labelnew), 'the montage is invalid');

% these are optional, at the end we will clean up the output in case they did not exist
haschantype = isfield(input, 'chantype') || isfield(input, 'chantypeold') || isfield(input, 'chantypenew');
haschanunit = isfield(input, 'chanunit') || isfield(input, 'chanunitold') || isfield(input, 'chanunitnew');

% make sure they always exist to facilitate the remainder of the code
if ~isfield(montage, 'chantypeold')
  montage.chantypeold = repmat({'unknown'}, size(montage.labelold));
  if isfield(input, 'chantype')
    [sel1, sel2] = match_str(montage.labelold, input.label);
    montage.chantypeold(sel1) = input.chantype(sel2);
  end
end

if ~isfield(montage, 'chantypenew')
  montage.chantypenew = repmat({'unknown'}, size(montage.labelnew));
  if isfield(input, 'chantype')
    [sel1, sel2] = match_str(montage.labelnew, input.label);
    montage.chantypenew(sel1) = input.chantype(sel2);
  end
end

if ~isfield(montage, 'chanunitold')
  montage.chanunitold = repmat({'unknown'}, size(montage.labelold));
  if isfield(input, 'chanunit')
    [sel1, sel2] = match_str(montage.labelold, input.label);
    montage.chanunitold(sel1) = input.chanunit(sel2);
  end
end

if ~isfield(montage, 'chanunitnew')
  montage.chanunitnew = repmat({'unknown'}, size(montage.labelnew));
  if isfield(input, 'chanunit')
    [sel1, sel2] = match_str(montage.labelnew, input.label);
    montage.chanunitnew(sel1) = input.chanunit(sel2);
  end
end

if ~isfield(input, 'label') && isfield(input, 'labelnew')
  % the input data structure is also a montage
  inputlabel = input.labelnew;
  if isfield(input, 'chantypenew')
    inputchantype = input.chantypenew;
  else
    inputchantype = repmat({'unknown'}, size(input.labelnew));
  end
  if isfield(input, 'chanunitnew')
    inputchanunit = input.chanunitnew;
  else
    inputchanunit = repmat({'unknown'}, size(input.labelnew));
  end
else
  % the input should describe the channel labels, and optionally the type and unit
  inputlabel = input.label;
  if isfield(input, 'chantype')
    inputchantype = input.chantype;
  else
    inputchantype = repmat({'unknown'}, size(input.label));
  end
  if isfield(input, 'chanunit')
    inputchanunit = input.chanunit;
  else
    inputchanunit = repmat({'unknown'}, size(input.label));
  end
end

% check the consistency of the montage
if ~iscell(montage.labelold) || ~iscell(montage.labelnew)
  ft_error('montage labels need to be specified in cell-arrays');
end

% check the consistency of the montage
if ~all(isfield(montage, {'tra', 'labelold', 'labelnew'}))
  ft_error('the second input argument does not correspond to a montage');
end

% check the consistency of the montage
if size(montage.tra,1)~=length(montage.labelnew)
  ft_error('the number of channels in the montage is inconsistent');
elseif size(montage.tra,2)~=length(montage.labelold)
  ft_error('the number of channels in the montage is inconsistent');
end

% use a default unit transfer from sensors to channels if not otherwise specified
if ~isfield(input, 'tra') && isfield(input, 'label')
  if     isfield(input, 'elecpos') && length(input.label)==size(input.elecpos, 1)
    nchan = length(input.label);
    input.tra = eye(nchan);
  elseif isfield(input, 'coilpos') && length(input.label)==size(input.coilpos, 1)
    nchan = length(input.label);
    input.tra = eye(nchan);
  elseif isfield(input, 'chanpos') && length(input.label)==size(input.chanpos, 1)
    nchan = length(input.label);
    input.tra = eye(nchan);
  end
end

% keep only the columns that are not empty
selcol = ~all(montage.tra==0, 1);
if istrue(keepunused)
  for i=find(selcol==false)
    % don't remove the column if it corresponds to one of the output channels
    selcol(i) = any(strcmp(montage.labelnew, montage.labelold{i}));
  end
end
montage.tra         = montage.tra(:,selcol);
montage.labelold    = montage.labelold(selcol);
montage.chantypeold = montage.chantypeold(selcol);
montage.chanunitold = montage.chanunitold(selcol);
clear selcol

% keep only the channels that are present in the input data
remove = setdiff(montage.labelold, intersect(montage.labelold, inputlabel));
selcol = match_str(montage.labelold, remove);
% we cannot just remove the columns, all rows that depend on it should also be removed
selrow = false(length(montage.labelnew),1);
for i=1:length(selcol)
  selrow = selrow & (montage.tra(:,selcol(i))~=0);
end
% convert from indices to logical vector
selcol = indx2logical(selcol, length(montage.labelold));
% remove rows and columns
montage.labelold    = montage.labelold(~selcol);
montage.labelnew    = montage.labelnew(~selrow);
montage.chantypeold = montage.chantypeold(~selcol);
montage.chantypenew = montage.chantypenew(~selrow);
montage.chanunitold = montage.chanunitold(~selcol);
montage.chanunitnew = montage.chanunitnew(~selrow);
montage.tra         = montage.tra(~selrow, ~selcol);
clear remove selcol selrow

% add columns for channels present in the input data but that are not specified in the montage, stick to the original order in the data
[dum, ix]   = setdiff(inputlabel, montage.labelold);
addlabel    = inputlabel(sort(ix));
addchantype = inputchantype(sort(ix));
addchanunit = inputchanunit(sort(ix));
m = size(montage.tra,1);
n = size(montage.tra,2);
k = length(addlabel);

% % check for NaNs in unused channels; these will be mixed in with the rest
% % of the channels and result in NaNs in the output even when multiplied
% % with zeros or identity
% if k > 0 && isfield(input, 'trial') % check for raw data now only
%   cfg = [];
%   cfg.channel = addlabel;
%   cfg.showcallinfo = showcallinfo;
%   data_unused = ft_selectdata(cfg, input);
%   % use an anonymous function to test for the presence of NaNs in the input data
%   hasnan = @(x) any(isnan(x(:)));
%   if any(cellfun(hasnan, data_unused.trial))
%     ft_error('FieldTrip:NaNsinInputData', ['Your input data contains NaNs in channels that are unused '...
%       'in the supplied montage. This would result in undesired NaNs in the '...
%       'output data. Please remove these channels from the input data (using '...
%       'ft_selectdata) before attempting to apply the montage.']);
%   end
% end

if istrue(keepunused)
  % add the channels that are not rereferenced to the input and output of the montage
  montage.tra((m+(1:k)),(n+(1:k))) = eye(k);
  montage.labelold    = cat(1, montage.labelold(:), addlabel(:));
  montage.labelnew    = cat(1, montage.labelnew(:), addlabel(:));
  montage.chantypeold = cat(1, montage.chantypeold(:), addchantype(:));
  montage.chantypenew = cat(1, montage.chantypenew(:), addchantype(:));
  montage.chanunitold = cat(1, montage.chanunitold(:), addchanunit(:));
  montage.chanunitnew = cat(1, montage.chanunitnew(:), addchanunit(:));
else
  % add the channels that are not rereferenced to the input of the montage only
  montage.tra(:,(n+(1:k))) = zeros(m,k);
  montage.labelold    = cat(1, montage.labelold(:), addlabel(:));
  montage.chantypeold = cat(1, montage.chantypeold(:), addchantype(:));
  montage.chanunitold = cat(1, montage.chanunitold(:), addchanunit(:));
end
clear addlabel addchantype addchanunit m n k

% determine whether all channels are unique
m = size(montage.tra,1);
n = size(montage.tra,2);
if length(unique(montage.labelnew))~=m
  ft_error('not all output channels of the montage are unique');
end
if length(unique(montage.labelold))~=n
  ft_error('not all input channels of the montage are unique');
end

% determine whether all channels that have to be rereferenced are available
if length(intersect(inputlabel, montage.labelold))~=length(montage.labelold)
  ft_error('not all channels that are required in the montage are available in the data');
end

% reorder the columns of the montage matrix
[selinput, selmontage] = match_str(inputlabel, montage.labelold);
montage.tra            = montage.tra(:,selmontage);
montage.labelold       = montage.labelold(selmontage);
montage.chantypeold    = montage.chantypeold(selmontage);
montage.chanunitold    = montage.chanunitold(selmontage);

% ensure that the montage is double precision
montage.tra = double(montage.tra);

% making the tra matrix sparse will speed up subsequent multiplications, but should not result in sparse output data
% this only makes sense for matrices with a lot of zero elements; for dense matrices it is faster to keep it full
if size(montage.tra,1)>1 && nnz(montage.tra)/numel(montage.tra) < 0.3
  montage.tra = sparse(montage.tra);
else
  montage.tra = full(montage.tra);
end

% update the channel scaling if the input has different units than the montage expects
if isfield(input, 'chanunit') && ~isequal(input.chanunit(:), montage.chanunitold(:))
  scale = ft_scalingfactor(input.chanunit, montage.chanunitold);
  montage.tra = montage.tra * diag(scale);
  montage.chanunitold = input.chanunit;
elseif isfield(input, 'chanunitnew') && ~isequal(input.chanunitnew(:), montage.chanunitold(:))
  scale = ft_scalingfactor(input.chanunitnew, montage.chanunitold);
  montage.tra = montage.tra * diag(scale);
  montage.chanunitold = input.chanunitnew;
end

if isfield(input, 'chantype') && ~isequal(input.chantype(:), montage.chantypeold(:))
  ft_error('inconsistent chantype in data and montage');
elseif isfield(input, 'chantypenew') && ~isequal(input.chantypenew(:), montage.chantypeold(:))
  ft_error('inconsistent chantype in data and montage');
end

if isfield(input, 'labelold') && isfield(input, 'labelnew')
  inputtype = 'montage';
elseif isfield(input, 'tra')
  inputtype = 'sens';
elseif ft_datatype(input, 'raw')
  inputtype = 'raw';
elseif ft_datatype(input, 'timelock')
  inputtype = 'timelock';
elseif ft_datatype(input, 'freq') && isfield(input, 'fourierspctrm')
  inputtype = 'freq';
elseif ft_datatype(input, 'freq') && isfield(input, 'crsspctrm')
  inputtype = 'freq_crsspctrm';

  % attempt to convert to a chan-chan representation
  input     = ft_checkdata(input, 'cmbstyle', 'full');
else
  inputtype = 'unknown';
end

switch inputtype
  case 'montage'
    % apply the montage on top of the other montage
    if isa(input.tra, 'single')
      % sparse matrices and single precision do not match
      input.tra = full(montage.tra) * input.tra;
    else
      input.tra = montage.tra * input.tra;
    end
    input.labelnew    = montage.labelnew;
    input.chantypenew = montage.chantypenew;
    input.chanunitnew = montage.chanunitnew;

  case 'sens'
    % apply the montage to an electrode or gradiometer description
    sens = input;
    clear input

    % apply the montage to the input
    if isa(sens.tra, 'single')
      % sparse matrices and single precision do not match
      sens.tra = full(montage.tra) * sens.tra;
    else
      sens.tra = montage.tra * sens.tra;
    end

    % The montage operates on the coil weights in sens.tra, but the output channels can be different.
    % If possible, we want to keep the old channel positions and orientations.
    [sel1, sel2] = match_str(montage.labelnew, inputlabel);
    keepchans = isequal(sel1(:)', 1:numel(montage.labelnew));

    posweight = abs(montage.tra);
    posweight = diag(1./sum(posweight,2)) * posweight;

    if isfield(sens, 'chanpos')
      if keepchans
        sens.chanpos = sens.chanpos(sel2,:);
      else
        if ~isfield(sens, 'chanposold')
          % add a chanposold only if it is not there yet
          sens.chanposold  = sens.chanpos;
          %  also keep the old label, type and unit for reference
          sens.labelold    = inputlabel;
          sens.chantypeold = inputchantype;
          sens.chanunitold = inputchanunit;
        end
        % compute the new channel positions as a weighted sum of the old ones
        sens.chanpos = posweight * sens.chanpos;
      end
    end

    if isfield(sens, 'chanori')
      if keepchans
        sens.chanori = sens.chanori(sel2,:);
      else
        if ~isfield(sens, 'chanoriold')
          sens.chanoriold = sens.chanori;
        end
        % compute the new channel orientations as a weighted sum of the old ones
        sens.chanori = posweight * sens.chanori;
      end
    end

    sens.label    = montage.labelnew;
    sens.chantype = montage.chantypenew;
    sens.chanunit = montage.chanunitnew;
  
    % rename the output variable
    input = sens;
    clear sens

  case 'raw'
    % apply the montage to the raw input data
    data = input;
    clear input

    % there are two challenges to deal with
    % 1) the input data can be single, and sparse(1) * single(0) fails
    % 2) the input data can contain nans, and 0*nan returns a nan
    % see http://bugzilla.fieldtriptoolbox.org/show_bug.cgi?id=3035 and https://github.com/fieldtrip/fieldtrip/issues/2169

    Ntrials = numel(data.trial);
    ft_progress('init', feedback, 'processing trials');
    for i=1:Ntrials
      ft_progress(i/Ntrials, 'processing trial %d from %d\n', i, Ntrials);

      if     ~isa(data.trial{i}, 'single') && ~any(isnan(data.trial{i}(:)))
        data.trial{i} = montage.tra * data.trial{i};
      elseif ~isa(data.trial{i}, 'single') &&  any(isnan(data.trial{i}(:)))
        % do not multiply 0 in the montage with nan in the data
        tmp = zeros(size(montage.tra,1), size(data.trial{i}, 2));
        for j=1:size(montage.tra,1)
          sel = montage.tra(j,:)~=0;
          tmp(j,:) = montage.tra(j,sel) * data.trial{i}(sel,:);
        end
        data.trial{i} = tmp;
      elseif  isa(data.trial{i}, 'single') && ~any(isnan(data.trial{i}(:)))
        % sparse matrices and single precision do not match
        data.trial{i} = full(montage.tra) * data.trial{i};
      elseif  isa(data.trial{i}, 'single') &&  any(isnan(data.trial{i}(:)))
        % sparse matrices and single precision do not match
        % do not multiply 0 in the montage with nan in the data, see http://bugzilla.fieldtriptoolbox.org/show_bug.cgi?id=3035
        tmp = zeros(size(montage.tra,1), size(data.trial{i}, 2));
        for j=1:size(montage.tra,1)
          sel = montage.tra(j,:)~=0;
          tmp(j,:) = full(montage.tra(j,sel)) * data.trial{i}(sel,:);
        end
        data.trial{i} = tmp;
      end
    end % for Ntrials
    ft_progress('close');

    data.label    = montage.labelnew;
    data.chantype = montage.chantypenew;
    data.chanunit = montage.chanunitnew;

    % rename the output variable
    input = data;
    clear data

  case 'timelock'
    % apply the montage to averaged data
    timelock = input;
    clear input

    fn = {'avg', 'trial', 'individual', 'cov'};
    for i=1:numel(fn)
      if isfield(timelock, fn{i})
        switch getdimord(timelock, fn{i})
          case 'chan_time'
            timelock.(fn{i}) = montage.tra * timelock.(fn{i});
          case 'rpt_chan_time'
            siz    = getdimsiz(timelock, fn{i});
            nrpt   = siz(1);
            nchan  = siz(2);
            ntime  = siz(3);
            output = zeros(nrpt, size(montage.tra,1), ntime);
            for rptlop=1:nrpt
              output(rptlop,:,:) = montage.tra * reshape(timelock.(fn{i})(rptlop,:,:), [nchan ntime]);
            end
            timelock.(fn{i}) = output; % replace the original field
          case 'chan_chan'
            timelock.(fn{i}) = montage.tra * timelock.(fn{i}) * montage.tra';
          case 'rpt_chan_chan'
            siz    = getdimsiz(timelock, fn{i});
            nrpt   = siz(1);
            nchan  = siz(2);
            output = zeros(nrpt, size(montage.tra,1), size(montage.tra,1));
            for rptlop=1:nrpt
              output(rptlop,:,:) = montage.tra * reshape(timelock.(fn{i})(rptlop,:,:), [nchan nchan]) * montage.tra';
            end
            timelock.(fn{i}) = output; % replace the original field
          otherwise
            ft_error('unsupported dimord for %s', fn{i});
        end % switch
      end % if
    end % for

    timelock.label    = montage.labelnew;
    timelock.chantype = montage.chantypenew;
    timelock.chanunit = montage.chanunitnew;

    % rename the output variable
    input = timelock;
    clear timelock

  case 'freq'
    % apply the montage to the spectrally decomposed data
    freq = input;
    clear input

    switch getdimord(freq, 'fourierspctrm')
      case 'rpttap_chan_freq'
        siz    = [getdimsiz(freq, 'fourierspctrm') 1];
        nrpt   = siz(1);
        nchan  = siz(2);
        nfreq  = siz(3);
        output = zeros(nrpt, size(montage.tra,1), nfreq);
        for foilop=1:nfreq
          output(:,:,foilop) = freq.fourierspctrm(:,:,foilop) * montage.tra';
        end
        freq.fourierspctrm = output; % replace the original Fourier spectrum

      case 'rpttap_chan_freq_time'
        siz    = getdimsiz(freq, 'fourierspctrm');
        nrpt   = siz(1);
        nchan  = siz(2);
        nfreq  = siz(3);
        ntime  = siz(4);
        output = zeros(nrpt, size(montage.tra,1), nfreq, ntime);
        for foilop=1:nfreq
          for toilop = 1:ntime
            output(:,:,foilop,toilop) = freq.fourierspctrm(:,:,foilop,toilop) * montage.tra';
          end
        end
        freq.fourierspctrm = output; % replace the original Fourier spectrum

      otherwise
        ft_error('unsupported dimord for fourierspctrm');
    end % switch

    freq.label    = montage.labelnew;
    freq.chantype = montage.chantypenew;
    freq.chanunit = montage.chanunitnew;

    % rename the output variable
    input = freq;
    clear freq
  case 'freq_crsspctrm'
    % input freq data has a chan-chan crsspctrm, montage needs to be
    % applied to both ends

    freq = input;
    clear input
    if contains(getdimord(freq, 'crsspctrm'), 'rpt')
      % first dimension is rpt-like, so the square is dimensions 2 and 3
      siz   = [getdimsiz(freq, 'crsspctrm') 1];
      nrpt  = siz(1);
      nchan = siz(2);
      nrest = prod(siz(4:end));
      output = zeros([nrpt size(montage.tra,1).*[1 1] siz(4:end)]);
      for rptlop = 1:nrpt
        for restlop = 1:nrest
          output(rptlop,:,:,restlop) = montage.tra*reshape(freq.crsspctrm(rptlop,:,:,restlop),[nchan nchan])*montage.tra';
        end
      end
      freq.crsspctrm = output;
    else
      % the square is dimensions 1 and 2
      siz   = [getdimsiz(freq, 'crsspctrm') 1];
      nrest = prod(siz(4:end));
      output = zeros([size(montage.tra,1).*[1 1] siz(4:end)]);
      for restlop = 1:nrest
        output(:,:,restlop) = montage.tra*freq.crsspctrm(:,:,restlop)*montage.tra';
      end
      freq.crsspctrm = output;
    end

    freq.label    = montage.labelnew;
    freq.chantype = montage.chantypenew;
    freq.chanunit = montage.chanunitnew;

    % rename the output variable
    input = freq;
    clear freq
  otherwise
    ft_error('unrecognized input');
end % switch inputtype

% only retain the chantype and/or chanunit if they were present in the input
if ~haschantype
  input = removefields(input, {'chantype', 'chantypeold', 'chantypenew'});
end
if ~haschanunit
  input = removefields(input, {'chanunit', 'chanunitold', 'chanunitnew'});
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% HELPER FUNCTION
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function y = indx2logical(x, n)
y = false(1,n);
y(x) = true;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% HELPER FUNCTION
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function nowarning(varargin)
return
