%% Copyright (C) 2020 Tony Richardson <richardson.tony@gmail.com>
%% 
%% This program 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.
%% 
%% This program 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 this program.  If not, see
%% <https://www.gnu.org/licenses/>.

%% -*- texinfo -*- 
%% @deftypefn {Function File} {} colfilt2 (@var{A}, @var{SELECT}, @var{NEIGHBORS}, @var{FUNC})
%% Process only certain neighbors of selected elements of a matrix with a 
%% user-supplied function.
%%
%% Executes the function FUNC on neighbors of selected elements of A.
%% SELECT is a logical index matrix that is the same size as A.  The TRUE
%% elements of SELECT correspond to elements of A that will be replaced by
%% a result from FUNC.
%% 
%% NEIGHBORS is an NN x 2 matrix where NN is the number of neighbors that 
%% are processed.  Each row of NEIGHBORS is a row and column pair of offsets 
%% from the center pixel. For example NEIGHBORS = [ -1 -1; -1 1; 1 -1 ; 1 1 ] 
%% would represent the four diagonal neighbors.  The left and right neighbors 
%% would be indicated using [0 -1; 0 1].  
%%
%% Zero-padding is used at the edges. (If another type of padding is desired
%% pad the input matrix before calling colfilt2 and strip the padding
%% afterward.  SELECT must be zero-padded to the same size as the input.)
%%
%% If there are NS elements selected then the argument passed to FUNC will be
%% an NN x NS matrix.  Each column represents a neighborhood of one of the
%% selected elements.  FUNC must return a 1 x NS result.
%%
%% colfilt2 returns a result that is the same size as A.
%%
%% colfilt2 is useful if you want to replace certain elements of a
%% matrix with results obtained by processing the neighborhoods of those
%% elements.
%%
%% colfilt2 may use less memory than colfilt if you need to process only a
%% few members of the neighborhood.
%%
%% Example: Replace all elements with values less than 10 with the median 
%% of their 4-neighbors.
%%
%% @example
%% idx = (A < 10);
%% NEIGHBORS = [-1 0; 0 -1; 0 1; 1 0 ];
%% B = colfilt2(A, idx, NEIGHBORS, @median);
%% @end example
%% 
%% @seealso{nlfilter, colfilt}
%% @end deftypefn

function r = colfilt2(A, SELECT, NEIGHBORS, FUNC)
  % Determine the amount of padding needed
  tpad = max([0 -NEIGHBORS(:, 1)']);  % top padding
  bpad = max([0  NEIGHBORS(:, 1)']);   % bottom padding
  lpad = max([0 -NEIGHBORS(:, 2)']);  % left padding
  rpad = max([0  NEIGHBORS(:, 2)']);   % right padding
  % APAD is the padded matrix
  MP = size(A,1) + tpad + bpad;  % number of rows in padded matrix
  NP = size(A,2) + lpad + rpad;  % number of cols
  APAD = zeros(MP, NP, class(A));
  APAD((1+tpad):(end-bpad),(1+lpad):(end-rpad)) = A;
  % SELECT matrix must be padded to the same size
  SPAD = zeros(MP, NP, 'logical');
  SPAD((1+tpad):(end-bpad),(1+lpad):(end-rpad)) = SELECT;

  % Find the row and column indexes of the SELECT'ed elements in APAD
  % First find the row and column indexes of all elements in APAD (RP and CP)
  % [CP RP] = meshgrid(1:NP, 1:MP);
  RP = (1:MP)'*ones(1, NP);
  CP = ones(MP, 1)*(1:NP);
  % Now get the row and column indexes of only the selected elements
  ri_p = RP(SPAD); ci_p = CP(SPAD);
  % NS is the number of selected elements
  NS = length(ri_p);
  
  % NN is the number of neighbors
  NN = size(NEIGHBORS, 1);
  % Formed stacked vector image.  Each row corresponds to a particular
  % neighbor of a selected element.  This loop cold be eliminated by
  % stacking the ri_p+neighbors(k,1) and ci_p+neighbors(k,2) vectors
  % and reshaping the result, but I am not sure any speed gain would
  % be worth the loss in code clarity.  Note that the loop is over the
  % number of neighbors (not the number of selected elements in A).
  BPAD = zeros(NN, NS, class(A));
  for k = 1:NN
    BPAD(k,:) = APAD(sub2ind([MP NP], ri_p+NEIGHBORS(k,1), ci_p+NEIGHBORS(k,2)))';
  end
  % FUNC must return one element for each column in B.
  r = A;
  r(SELECT) = FUNC(BPAD);
end

%!demo
%! %% Replace diagonal elements with average of left and right neighbors
%! %% Zero-padding is used at the edges
%! A = magic (12)
%! SELECT = eye(size(A), 'logical');
%! NEIGHBORS = [0 -1; 0 1];
%! colfilt2 (A, SELECT, NEIGHBORS, @mean)

