function varargout = pagefun(fcn, varargin)
%PAGEFUN Apply a function to each page of codistributed array
%   PAGEFUN iterates over the pages of a codistributed array, applying the same 
%   function to each page.   
%   
%   A = PAGEFUN(FUN, B) applies the function specified by FUN to each page
%   of the codistributed array B, and returns the results in codistributed array
%   A such that A(:,:,I,J,...) = FUN(B(:,:,I,J,...)).  FUN is a function handle 
%   to a function that takes one two-dimensional input argument.
%   
%   A = PAGEFUN(FUN, B, C, ...) evaluates the function specified by FUN using pages of the 
%   arrays B, C, ... as input arguments with scalar expansion enabled.  Any of the
%   input page dimensions that are scalar will be virtually replicated to match the 
%   size of the other arrays in that dimension so that 
%   A(:,:,I,J,...) = FUN(B(:,:,I,J,...), C(:,:,I,J,...), ...).  At least one of 
%   the inputs B, C, ... must be a codistributed array.  Any other inputs held in CPU 
%   memory are converted to a codistributed array before calling the function.
%   If an array is to be used in several different PAGEFUN calls it is more 
%   efficient to convert that array to a codistributed array before your series of 
%   PAGEFUN calls.  The input pages B(:,:,I, J, ...), C(:,:,I, J, ...), ... must 
%   satisfy all of FUN's input and output requirements.
%   
%   [A, B, ...] = PAGEFUN(FUN, C, ...), where FUN is a function handle
%   to a function that returns multiple outputs, returns codistributed arrays
%   A, B,..., each corresponding to one of the output arguments of FUN.  PAGEFUN
%   invokes FUN with as many outputs as there are in the call to PAGEFUN.
%   All elements of A must be the same class; B can be a different class from A, 
%   but all elements of B must be of the same class.
%   
%   PAGEFUN supports most element-wise codistributed array functions for FUN. These functions
%   are also supported:
%       @conv
%       @conv2
%       @ctranspose
%       @fliplr
%       @flipud
%       @inv
%       @lu
%       @mldivide
%       @mrdivide
%       @mtimes
%       @norm
%       @qr
%       @rot90
%       @svd
%       @transpose
%       @tril
%       @triu
%   
%   Examples:
%   spmd
%       M = 3;                                  % output number of rows
%       K = 6;                                  % matrix multiply inner dimension
%       N = 2;                                  % output number of columns
%       P1 = 10;                                % size of first page dimension
%       P2 = 17;                                % size of second page dimension
%       P3 = 4;                                 % size of third page dimension
%       P4 = 12;                                % size of fourth page dimension
%       A = codistributed.rand(M, K, P1,  1, P3);
%       B = codistributed.rand(K, N,  1, P2, P3, P4);
%       C = pagefun(@mtimes, A, B);
%       size(C)                                 % M x N x P1 x P2 x P3 x P4
%   end
%   
%   spmd
%       M = 300;                     % output number of rows
%       K = 500;                     % matrix multiply inner dimension
%       N = 1000;                    % output number of columns
%       P = 200;                     % number of pages
%       A = codistributed.rand(M, K);   
%       B = codistributed.rand(K, N, P);
%       C = pagefun(@mtimes, A, B);
%       size(C)                      % returns M x N x P 
%   end
%   
%   See also  CODISTRIBUTED, function_handle, arrayfun, bsxfun, gather, pagemtimes.
%   


%   Copyright 2013-2022 The MathWorks, Inc.

narginchk(2, Inf);

% First input must be a function handle
if ~isa(fcn, "function_handle")
    error(message("parallel:array:FunfunFirstArgFcnHandle"))
end

% Sparse is not allowed
if any(cellfun(@issparse, varargin))
    error(message("parallel:array:PagefunNoSparse"));
end

% Check for local data with a remote parameter.
if ~iscodistributed(varargin{1}) && iFcnRequiresFirstInputRemote(fcn)
    error(message("parallel:array:PagefunDataMustBeRemote", "codistributed"));
end

% Calculate the number of pages in dimensions 3 and up, referred as 
% trailing output dimensions. Will error if the corresponding page 
% dimensions of the inputs are not consistent.
trailingOutDims = parallel.internal.array.pagefunTrailingDimensions(varargin, ...
    "parallel:array:PagefunIncompatiblePages");

% Get the codistributor of the output.
outCodistr = pagefunGetOutputsCodistr(varargin, trailingOutDims);
% Redistribute or replicate all inputs, then get the local parts.
varargin = pagefunRedistributeInputsIfNeeded(varargin, outCodistr);
varargin = pagefunGetLocalParts(varargin);

% We should now have the correct pages on every worker. Call in-memory
% pagefun on the local parts.
try
    [LPout{1:max(1, nargout)}] = parallel.internal.array.inmemoryPagefun(fcn, varargin{:});
    err = [];
catch err
end
% Check for errors
err = spmdCat(err);
if ~isempty(err)
    throw(err(1));
end

% Now build the outputs
[varargout{1:max(1, nargout)}] = pagefunBuildOutputs(LPout, outCodistr, trailingOutDims);

end % pagefun

% Return true for functions that require the first input to be remote.
function tf = iFcnRequiresFirstInputRemote(fcn)
tf = isequal(fcn, @rot90) ...
    || isequal(fcn, @lu) ...
    || isequal(fcn, @qr) ...
    || isequal(fcn, @svd) ...
    || isequal(fcn, @tril) ...
    || isequal(fcn, @triu) ...
    || isequal(fcn, @norm);

end
