function P = pCumop(cumFcn, adjustFcn, initVal, A, varargin)
%pCumOp Template for CUMMAx, CUMMIN, CUMPROD, and CUMSUM for codistributed array


%   Copyright 2006-2022 The MathWorks, Inc.

distributedutil.CodistParser.verifyNonCodistributedInputs(varargin);
narginchk(4, 7); % Three optional args

% Tables support accumulations, but we don't yet!
if istabular(A)
    error(message("parallel:distributed:TabularMathUnsupported", upper(func2str(cumFcn))))
end

% Gather trailing arguments.
args = distributedutil.CodistParser.gatherElements(varargin);

% Set default for dimension, isReverse, and isIncludeNan
dim = distributedutil.Sizes.firstNonSingletonDimension(size(A));
isReverse = false;
isIncludeNan = false;
        
% Check optional arguments and extract dimension if present
if nargin > 4
    % If first args is not a scalar string, it is assumed that it is the
    % dimension.
    flagArgs = args;
    if ~isScalarString(varargin{1})
        dim = args{1};
        flagArgs(1) = [];
    else
        % Dimension was not provided, hence, include default dimension as
        % first element in args.
        args = [dim, args];
    end
    
    % Check additional inputs using in-memory cumFcn on an empty array and
    % throw errors as caller.
    try
        cumFcn([], args{:});
    catch ME
        throwAsCaller(ME);
    end

    if numel(flagArgs)>0
        % Check if "reverse" is the start of any of the additional inputs 
        isReverse = any(cellfun(@(x)startsWith("reverse", x, ...
            'IgnoreCase', true), flagArgs, 'UniformOutput', true));
        % Check if "includenan" or "includemissing" is the start of any of
        % the additional inputs
        isIncludeNan = any(cellfun(@(x)any(startsWith(["includenan" "includemissing"], x, ...
            'IgnoreCase', true)), flagArgs, 'UniformOutput', true));
    end
else
    % Use default dimension afterwards if not provided since the default
    % dimension for a local part might not be the desired default dimension
    % of the entire array.
    args = {dim};
end

% Check that we aren't trying to do maths on complex integers
% We need to build the MException outside throwIfComplexInt to be conform
% with core MATLAB error IDs although we might not throw it
ME = MException(['MATLAB:', func2str(cumFcn), ':complexIntegers'], ...
    getString(message('MATLAB:cumfun:complexIntegers')));
throwIfComplexInt(ME, A);

% Call the normal MATLAB version if the data is on the host
if ~isa(A, 'codistributed')
    try
        P = cumFcn(A, args{:});
        return;
    catch E
        throwAsCaller(E);
    end
end

try
    % This implementation only supports codistributor1d.
    codistributed.pVerifyUsing1d(func2str(cumFcn), A);
catch E
    % Error stack should only show cummax, cummin, cumprod, or cumsum, not pCumop.
    throwAsCaller(E);
end

if isempty(A) || (size(A, dim) == 1 && isIncludeNan)
    % cumop on empty or a singleton dimension when using includenan
    P = A;
else
    % cumop on a non-singleton dimension
    % first accumulate the local part
    localP = cumFcn(getLocalPart(A), args{:});
    
    % If the accumulation dimension is the distribution dimension we
    % also need to adjust for previous totals.
    aDist = getCodistributor(A);
    if aDist.Dimension == dim
        localP = iAdjustLocalPart(localP, adjustFcn, initVal, ndims(A), dim, isReverse, isIncludeNan);
    end
    P = codistributed.pDoBuildFromLocalPart(localP, aDist);
end
end % pCumOp

function localP = iAdjustLocalPart(localP, adjustFcn, initVal, ndimsA, dim, isReverse, isIncludeNan)
% Adjust the values in this local part using the final values from
% previous local parts.

% Form an index expression that will extract the values from each page
pageindex = cell(1, ndimsA);
[pageindex{:}] = deal(':');
if ~isempty(localP)
    if isReverse
        % We want the first element from each local part
        pageindex{dim} = 1;
    else
        % We want the final element from each local part
        pageindex{dim} = size(localP, dim);
    end
end
% Extract final values on each worker
scale = spmdCat({localP(pageindex{:})}, dim);
adjustFcnStr = func2str(adjustFcn);

if ~isempty(localP)
    setcumscale = false;
    
    % Initialize cumscale and make sure it has the same sparsity as the
    % local part.
    if issparse(localP) 
        cumscale = sparse(initVal);
    else
        cumscale = full(initVal);
    end
    
    % Scan all the preceding workers to build up the scaling factor
    if isReverse
        workersToScan = spmdIndex+1:spmdSize;
    else
        workersToScan = 1:spmdIndex-1;
    end
    for ii = workersToScan
        if ~isempty(scale{ii})
            % Create inputs to adjustFcn depending on the function
            fcnInputs = {scale{ii}, cumscale};
            if ismember(adjustFcnStr, {'max', 'min'})
                % For min and max, we have to include NaNs if isIncludeNan
                % is true and we are either not in the first step (cumscale
                % is not initVal anymore, i.e., setcumscale==true) or we
                % are in the first step and scale{ii} has a NaN.
                if isIncludeNan && (setcumscale || ...
                        ~setcumscale && any(isnan(scale{ii})))
                    flag = 'includenan';
                else
                    flag = 'omitnan';
                end
                fcnInputs = [fcnInputs, flag]; %#ok<AGROW>
            end
            cumscale = adjustFcn(fcnInputs{:});
            setcumscale = true;
        end
    end
    % If any worker set the scaling factor, apply it to our local part.
    if setcumscale
        % Create inputs to adjustFcn depending on the function
        fcnInputs = {localP, iExpandToMatch(cumscale, localP, dim, ndimsA)};
        if ismember(adjustFcnStr, {'max', 'min'})
             % For min and max, include NaNs if isInlcudenan is true.
            if isIncludeNan
                flag = 'includenan';
            else
                flag = 'omitnan';
            end
            fcnInputs = [fcnInputs, flag];
        end
        localP = adjustFcn(fcnInputs{:});
    end
end
end % iAdjustLocalPart

function  expandX = iExpandToMatch(reducedX, X, dim, ndimsA)
s = ones(1, ndimsA);
s(dim) = size(X, dim);
expandX = repmat(reducedX, s);
end % iExpandToMatch
