function R = pSumAndProd(fcn, A, argList)
%pSumAndProd is a common implementation of sum and prod for (co)distributed 
%   matrices
%   


%   Copyright 2012-2022 The MathWorks, Inc.

distributedutil.CodistParser.verifyNonCodistributedInputs({A, argList});
argList = distributedutil.CodistParser.gatherElements(argList);
if ~isa(A, 'codistributed')
    % Dimension or accumType are codistributed, but not the array itself.
    try
        R = fcn(A, argList{:});
        return;
    catch E
        throwAsCaller(E);
    end
end

try
    % Validate inputs with in-memory function
    distributedutil.validateSyntax(fcn, [{A}, argList]);
    [argList, trailingArgs] = iGetTrailingArgs(fcn, argList);
    [A, trailingArgs] = iPreprocessIfOmitnan(fcn, A, trailingArgs);

    % Forward trailing args to sum or prod, then reduce.
    accumFcn = @(varargin) fcn(varargin{:}, trailingArgs{:});
    R = codistributed.pReductionOpAlongDim(accumFcn, A, argList{:});
catch E
    throwAsCaller(E); % Strip ourselves from the stack.
end
end


function [argList, trailingArgs] = iGetTrailingArgs(fcn, argList)
% Helper to split data and dimension args from trailing flags. Unlike the
% other reduction functions (any, all), sum and prod take extra string
% arguments that control the type and missing value handling. We peel those
% off and hide them from our generic reduction processing.

idx = find(~cellfun(@isScalarString, argList), 1, 'last');
if isempty(idx)
    % No non-char args
    trailingArgs = argList;
    argList = {};
else
    % Strip trailing args
    trailingArgs = argList(idx+1:end);
    argList = argList(1:idx);
end

% Special case for "all" - treat as the dim arg, not a flag
if ~isempty(trailingArgs) ...
        && matlab.internal.datatypes.isScalarText(trailingArgs{1}) ...
        && strncmpi(trailingArgs{1}, "all", strlength(trailingArgs{1}))
    argList{1,end+1} = "all";
    trailingArgs(1) = [];
end

% If more than one trailing arg (i.e. more than just dim) we have a
% problem.
if numel(argList) > 1
    error(message(['MATLAB:',func2str(fcn),':unknownFlag']));
end

% Use a local call to sanity check the syntax
fcn(1, argList{:}, trailingArgs{:});
end

function [A, trailingArgs] = iPreprocessIfOmitnan(fcn, A, trailingArgs)
% Helper to preprocess the inputs if "omitnan" is specified. If "omitnan"
% or "omitmissing" we replace NaN entries by 0 or 1 before reducing and
% swallow that argument. At this point, we know the provided flags are
% consistent and are not repeated.

isOmitnanOpt = cellfun(@(arg) matlab.internal.datatypes.isText(arg) && ...
    any(strncmpi(arg, ["omitnan", "omitmissing"], strlength(arg))), trailingArgs);
if nnz(isOmitnanOpt) == 1
    localA = getLocalPart(A);
    codA = getCodistributor(A);

    assert(isequal(fcn, @sum) || isequal(fcn, @prod), ...
        "Function handle must be SUM or PROD.")

    if istabular(localA)
        % Replace values per variable as they might have different
        % datatypes.
        fillValue = varfun(@(x) iExtractFillValue(x, fcn), localA, "OutputFormat", "cell");
        for k = 1:size(localA, 2)
            localA{:,k} = fillmissing(localA{:,k}, "constant", fillValue{k});
        end
    else
        fillValue = iExtractFillValue(localA, fcn);
        localA = fillmissing(localA, "constant", fillValue);
    end

    A = codistributed.pDoBuildFromLocalPart(localA, codA);
    trailingArgs = trailingArgs(~isOmitnanOpt);
end
end

function fillValue = iExtractFillValue(x, fcn)
if isduration(x)
    % DURATION is supported for SUM only, for PROD we flow-through to error
    % later.
    fillValue = hours(0);
elseif isequal(fcn, @sum)
    % If SUM, replace NaN by 0
    fillValue = 0;
else
    % If PROD, replace NaN by 1
    fillValue = 1;
end
end