## Copyright (C) 2020 ar63
## 
## 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 {} {@var{retval} =} nlfilter2 (@var{A}, @var{BLOCK_SIZE}, @var{FUNC})
## Process matrix in sliding blocks with user-supplied function.
##
## This processes a region of size BLOCK_SIZE around each element of A.  It processes
## the region using FUNC and replaces the element with the processed result.
## If BLOCK_SIZE = [m n] then FUNC must accept an argument with mxn rows and a
## number of columns equal to the number of elements in A.  The matrix is
## zero-padded near the edges.
##
## Each column in the argument to FUNC is a copy of the neighborhood around an
## element in A.  The neighborhood elements are arranged in left-to-right then
## top-to-bottom order.
##
## This is an attempt to create a 'vectorized' version of nlfilter.
## Additional work needs to be done on this function to make it as robust
## and general as nlfilter.
##
## @seealso{nlfilter}
## @end deftypefn

## Author: Tony Richardson richardson.tony@gmail.com
## Created: 2020-07-17

function r = nlfilter2(A, BLOCK_SIZE, FUNC)
  [M N] = size(A);
  % NS is the number of pixels in image
  NS = M*N;
  % NN is the number of neighbors
  NN = BLOCK_SIZE(1)*BLOCK_SIZE(2);
  % Number of rows and columns to pad on each edge
  rows_to_pad = round(BLOCK_SIZE(1)-1)/2;
  cols_to_pad = round(BLOCK_SIZE(2)-1)/2;
  
  % APAD is the zero-padded image
  APAD = [
    zeros(rows_to_pad, N+2*cols_to_pad, class(A));
    zeros(M, cols_to_pad, class(A)) A zeros(M, cols_to_pad, class(A));
    zeros(rows_to_pad, N+2*cols_to_pad, class(A))
  ];
  
  [MP NP] = size(APAD);
      
  % Find the row and column indexes of the pixels in the padded image that
  % were in the original image.(Use meshgrid instead?)
  RP = vec((rows_to_pad+(1:M))'*ones(1,N));
  CP = vec(ones(M,1)*(cols_to_pad+(1:N)));
  
  % Find row and column offsets (relative) to center pixel of all pixels in block
  rind = vec(ones(BLOCK_SIZE(2),1)*((1:BLOCK_SIZE(1))-(rows_to_pad+1)));
  cind = vec(((1:BLOCK_SIZE(2))-(cols_to_pad+1))'*ones(1,BLOCK_SIZE(1)));
  % Formed stacked vector image.  
  B = zeros(NN, NS, class(A));
  for k = 1:NN
    B(k,:) = APAD(sub2ind([MP NP], RP+rind(k), CP+cind(k)));
  end
  % FUNC must return one element for each column in B.
  r = reshape(FUNC(B), [M N]);
end