function [y,counts] = combineVectorElements(f,x,identity,dim,omitnan,vlen,vstart)
%MATLAB Code Generation Library Function

%   Performs dimension-collapsing operations like SUM and PROD. For all
%   vectors v along dimension DIM, the corresponding scalar value is
%   computed
%
%       s = cast(v(1),'like',identity);
%       for k = 2:size(x,dim)
%           s = f(s,cast(v(k),'like',identity));
%       end
%
%   For sparse inputs, the output is initialized to 0 if the current vector
%   contains any zeros. This may result in behavior differences from dense
%   matrices but matches MATLAB behavior.
%
%   The IDENTITY value therefore determines both the output type and the
%   "accumulator" type. Note that IDENTITY must be complex if the output of
%   F is complex.
%
%   For SUM, F = @plus, and IDENTITY should be 0 of the output type.
%   For PROD, use F = @times, and IDENTITY should be 1 of the output type.
%
%   VLEN is an optional input that defaults to size(X,DIM). VLEN must be
%   in the range 0 <= VLEN <= size(X,DIM). Only the first VLEN elements of
%   each vector are combined to form the output Y.
%
%   VSTART is an optional input that defaults to 1. Each vector is indexed
%   starting at VSTART, i.e. the vector is indexed from VSTART to VSTART +
%   VLEN - 1. VSTART is not supported for sparse x.
%
%   The second output, COUNTS, is an array the same size as Y. COUNTS(k)
%   is the number of elements that contributed to the value of Y(k).

%   Copyright 2002-2022 The MathWorks, Inc.
%#codegen

coder.internal.prefer_const(identity,dim,omitnan);
coder.internal.allowEnumInputs;
coder.internal.allowHalfInputs;
coder.internal.assert(coder.internal.isConst(dim), ...
    'Coder:toolbox:dimNotConst');

% Preserves all dimensions except dim
ND = coder.internal.indexInt(coder.internal.ndims(x));
coder.unroll();
for i = 1:(dim - 1)
    coder.internal.preservesInputSize('y',i,'x',i);
end
coder.unroll()
for j = (dim + 1):ND
    coder.internal.preservesInputSize('y',j,'x',j);
end

if nargin < 7
    vstart = coder.internal.indexInt(1);
end
if nargin < 6
    vlen = coder.internal.indexInt(size(x,dim));
end
if nargin < 5
    omitnan = false;
end
if nargin < 4
    dim = coder.internal.constNonSingletonDim(x);
end
if nargin < 3
    % scalarEg(x,x) --> double(0) for logical and char types.
    identity = coder.internal.scalarEg(x,x);
end
coder.internal.prefer_const(identity,dim,omitnan,vlen,vstart);
% If the operation is a non-sparse floating-point summation, use a blocked
% algorithm to mitigate round-off error over long sums. GPU Coder
% transforms the simple loop formulation and would not generate good code
% with the blocked algorithm.
FLOATSUM = coder.const(isequal(f,@plus) && isfloat(identity));
if FLOATSUM && ~issparse(x) && ~coder.gpu.internal.isGpuEnabled
    % The block size in MATLAB is 4096, but MATLAB uses a SIMD formulation
    % that divides each block into 4 pieces, yielding an effective block
    % size of 4096/4 = 1024.
    BLOCKSIZE = coder.internal.indexInt(1024);
    if nargout == 2
        [y,counts] = blockedSummation(BLOCKSIZE,x,dim,omitnan, ...
            identity,vlen,vstart);
    else
        y = blockedSummation(BLOCKSIZE,x,dim,omitnan,identity,vlen,vstart);
    end
    return
end
% Avoid counting during iterations unless counts are requested and there is
% a chance they will turn out to be anything but vlen.
DOCOUNTS = nargout == 2 && isfloat(x) && omitnan;
FLATTEN = ~coder.internal.avoidArrayFlattening;
% Perform the operation.
if isempty(x) || vlen == 0
    % If size(x,dim) is zero, y can be non-empty when x is empty. It's
    % cleanest and safest to deal with this possibility as a special case.
    if dim >= 1 && dim <= ND
        if issparse(x)
            if dim == 1
                y = coder.internal.sparse.genericLike(identity,ONE,x.n);
            else
                y = coder.internal.sparse.genericLike(identity,x.m,ONE);
            end
        else
            sz = size(x);
            sz(dim) = 1;
            y = eml_expand(identity,sz);
        end
    else
        if issparse(x)
            y = coder.internal.sparse.genericLike(identity,x.m,x.n);
        else
            y = eml_expand(identity,size(x));
        end
    end
    if DOCOUNTS
        counts = eml_expand(vlen,size(y));
    else
        counts = coder.nullcopy(zeros(size(y),'like',vlen));
    end
elseif dim > ND
    if DOCOUNTS
        [y,counts] = castLikeAndReplaceNaNs(x,identity);
    elseif omitnan
        y = castLikeAndReplaceNaNs(x,identity);
        counts = coder.nullcopy(zeros(size(y),'like',vlen));
    else
        if issparse(x)
            y = cast(x,'like',sparse(identity));
        else
            y = cast(x,'like',identity);
        end
        counts = coder.nullcopy(zeros(size(y),'like',vlen));
    end
elseif issparse(x)
    % We don't support vstart with sparse inputs.
    coder.internal.assert(vstart == 1,'MATLAB:maxrhs');
    if DOCOUNTS
        [y,counts] = sparseIter(f,identity,x,dim,omitnan);
    else
        y = sparseIter(f,identity,x,dim,omitnan);
        counts = coder.nullcopy(zeros(size(y),'like',vlen));
    end
elseif FLATTEN
    if DOCOUNTS
        [y,counts] = colMajorFlatIter(f,identity,x,dim,omitnan,vlen,vstart);
    else
        y = colMajorFlatIter(f,identity,x,dim,omitnan,vlen,vstart);
        counts = coder.nullcopy(zeros(size(y),'like',vlen));
    end
else
    if DOCOUNTS
        [y,counts] = nestedIter(f,identity,x,dim,omitnan,vlen,vstart);
    else
        y = nestedIter(f,identity,x,dim,omitnan,vlen,vstart);
        counts = coder.nullcopy(zeros(size(y),'like',vlen));
    end
end
if nargout == 2 && ~DOCOUNTS
    % We did not do counts during iterations, but the output is requested.
    counts = eml_expand(vlen,size(y));
end

%--------------------------------------------------------------------------

function [y,counts] = colMajorFlatIter(f,identity,x,dim,omitnan,vlen,vstart)
coder.inline('always');
coder.internal.prefer_const(identity,dim,omitnan,vlen,vstart);
TWO = coder.internal.indexInt(2);
OMITNANS = isfloat(x) && omitnan;
DOCOUNTS = nargout == 2;
vstride = coder.internal.prodsize(x,'below',dim);
pagesize = vstride*coder.internal.indexInt(size(x,dim));
npages = coder.internal.prodsize(x,'above',dim);
voffset = (vstart - 1)*vstride;
% Preallocate outputs.
if coder.const(dim >= 1 && dim <= coder.internal.ndims(x))
    sz = size(x);
    sz(dim) = 1;
    y = coder.nullcopy(eml_expand(identity,sz));
else
    y = coder.nullcopy(eml_expand(identity,size(x)));
end
if DOCOUNTS
    counts = coder.nullcopy(zeros(size(y),coder.internal.indexIntClass));
end
for i = ONE:npages
    xpageoffset = (i - 1)*pagesize + voffset;
    ypageoffset = (i - 1)*vstride;
    % Initialize the current page of y with the first elements of
    % each vector to be processed.
    for j = ONE:vstride
        ix = xpageoffset + j; % index of the first element of x vector
        iy = ypageoffset + j; % index of the output value
        if OMITNANS && isnan(x(ix))
            y(iy) = identity;
            if DOCOUNTS
                counts(iy) = 0;
            end
        else
            y(iy) = cast(x(ix),'like',identity);
            if DOCOUNTS
                counts(iy) = 1;
            end
        end
    end
    % Accumulate remaining elements of each vector on this page.
    for k = TWO:vlen
        xoffset = xpageoffset + (k - 1)*vstride;
        for j = ONE:vstride
            ix = xoffset + j;
            iy = ypageoffset + j;
            if ~(OMITNANS && isnan(x(ix)))
                y(iy) = f(y(iy),cast(x(ix),'like',identity));
                if DOCOUNTS
                    counts(iy) = counts(iy) + 1;
                end
            end
        end
    end
end

%--------------------------------------------------------------------------

function [y,counts] = nestedIter(f,identity,x,dim,omitnan,vlen,vstart)
% Assumes dim <= length(size(x)).
coder.inline('always');
coder.internal.prefer_const(identity,dim,omitnan,vlen,vstart);
DOCOUNTS = nargout == 2;
ND = coder.internal.indexInt(coder.internal.ndims(x));
if coder.const(dim >= 1 && dim <= ND)
    sz = size(x);
    sz(dim) = 1;
    y = coder.nullcopy(eml_expand(identity,sz));
else
    y = coder.nullcopy(eml_expand(identity,size(x)));
end
if DOCOUNTS
    counts = coder.nullcopy(zeros(size(y),coder.internal.indexIntClass));
else
    counts = coder.internal.indexInt([]);
end
if coder.isRowMajor
    OuterDim = ONE;
else
    OuterDim = ND;
end
[y,counts] = genloops(f,identity,x,dim,omitnan, ...
    vlen,vstart,y,counts,OuterDim);

%--------------------------------------------------------------------------

function [y,counts] = genloops(f,identity,x,dim,omitnan, ...
    vlen,vstart,y,counts,loopdim,varargin)
% Iterate using full ND-indexing.
coder.inline('always');
coder.internal.prefer_const(identity,dim,omitnan,vlen,vstart,loopdim);
ROWMAJOR = coder.isRowMajor;
OMITNANS = isfloat(x) && omitnan;
EMPTYCOUNTS = coder.internal.isConstTrue(isempty(counts));
DOCOUNTS = nargout == 2 && ~EMPTYCOUNTS;
TWO = coder.internal.indexInt(2);
ND = coder.internal.indexInt(coder.internal.ndims(x));
voffset = vstart - 1;
if coder.const(loopdim == dim)
    if ROWMAJOR
        % At this point, varargin{:} is dim-1 elements long and contains
        % the leading indices of the page of x we need to process. For
        % example, if dim == 3 and x is 5-D, the page is
        % x(varargin{1},varargin{2},:,:,:). At this point we need to
        % initialize y(varargin{1},varargin{2},1,:,:) =
        % x(varargin{1},varargin{2},1,:,:), though taking account of
        % omitnan behavior, if applicable.
        [y,counts] = initPage(f,identity,x,dim,omitnan, ...
            y,counts,loopdim + 1,varargin{:},vstart);
        % Now accumulate the 2:vlen remaining elements of each vector.
        for k = TWO:vlen
            [y,counts] = genloops(f,identity,x,dim,omitnan, ...
                vlen,vstart,y,counts,loopdim + 1,varargin{:},voffset + k);
        end
    else
        % At this point, varargin{:} is ND-dim elements long and contains
        % the trailing indices of the page of x we need to process. For
        % example, if dim == 3 and x is 5-D, the page is
        % x(:,:,:,varargin{2},varargin{3}). At this point we need to
        % initialize y(:,:,1,varargin{1},varargin{2}) =
        % x(:,:,1,varargin{1},varargin{2}), though taking account of
        % omitnan behavior, if applicable.
        [y,counts] = initPage(f,identity,x,dim,omitnan, ...
            y,counts,loopdim - 1,vstart,varargin{:});
        % Now accumulate the 2:vlen remaining elements of each vector.
        for k = TWO:vlen
            [y,counts] = genloops(f,identity,x,dim,omitnan, ...
                vlen,vstart,y,counts,loopdim - 1,voffset + k,varargin{:});
        end
    end
elseif coder.const(loopdim >= ONE && loopdim <= ND)
    for k = ONE:size(x,loopdim)
        if ROWMAJOR
            [y,counts] = genloops(f,identity,x,dim,omitnan, ...
                vlen,vstart,y,counts,loopdim + 1,varargin{:},k);
        else
            [y,counts] = genloops(f,identity,x,dim,omitnan, ...
                vlen,vstart,y,counts,loopdim - 1,k,varargin{:});
        end
    end
elseif vlen >= TWO
    % At this point we are in the innermost loop and need to add the
    % currently-indexed element of x to its accumulator in y.
    xk = x(varargin{:});
    if ~OMITNANS || ~isnan(xk)
        yidx = varargin;
        yidx{dim} = ONE;
        y(yidx{:}) = f(y(yidx{:}),cast(xk,'like',identity));
        if DOCOUNTS
            counts(yidx{:}) = counts(yidx{:}) + 1;
        end
    end
end

%--------------------------------------------------------------------------

function [y,counts] = initPage(f,identity,x,dim,omitnan, ...
    y,counts,loopdim,varargin)
% Initialize the current page of row-major data.
coder.inline('always');
coder.internal.prefer_const(identity,dim,omitnan,loopdim);
ROWMAJOR = coder.isRowMajor;
OMITNANS = isfloat(x) && omitnan;
EMPTYCOUNTS = coder.internal.isConstTrue(isempty(counts));
DOCOUNTS = nargout == 2 && ~EMPTYCOUNTS;
ONE = coder.internal.indexInt(1);
ND = coder.internal.indexInt(coder.internal.ndims(x));
if coder.const(loopdim >= ONE && loopdim <= ND)
    for k = ONE:size(x,loopdim)
        if ROWMAJOR
            [y,counts] = initPage(f,identity,x,dim,omitnan, ...
                y,counts,loopdim + 1,varargin{:},k);
        else
            [y,counts] = initPage(f,identity,x,dim,omitnan, ...
                y,counts,loopdim - 1,k,varargin{:});
        end
    end
else
    xk = x(varargin{:});
    outsubs = varargin;
    outsubs{dim} = ONE;
    if OMITNANS && isnan(xk)
        y(outsubs{:}) = identity;
        if DOCOUNTS
            counts(outsubs{:}) = 0;
        end
    else
        y(outsubs{:}) = x(varargin{:});
        if DOCOUNTS
            counts(outsubs{:}) = 1;
        end
    end
end

%--------------------------------------------------------------------------

function [y,counts] = sparseIter(f,identity,x,dim,omitnan)
coder.inline('always');
coder.internal.prefer_const(identity,dim,omitnan);
if dim == 1
    [y,counts] = sparseColumnIter(f,identity,x,dim,omitnan);
else
    % Row reduce
    % If the number of rows is bounded by the size of an already allocated
    % array, we use a fast map-like algorithm. Otherwise, fall back to
    % slower methods which use less memory.
    if canUseRowMap(x)
        y = sparseRowMapIter(f,identity,x,dim,omitnan);
    else
        y = sparseRowReduceIter(f,identity,x,dim,omitnan);
    end
end
coder.internal.sparse.sanityCheck(y);

%--------------------------------------------------------------------------

function [y,counts] = sparseColumnIter(f,identity,x,dim,omitnan)
% Simple column iteration for sparse
coder.inline('always');
coder.internal.prefer_const(identity,dim,omitnan);
OMITNANS = isfloat(x) && omitnan;
counts = [];
ncol = x.n;
nalloc = eml_min(ncol,nnzInt(x));
y = coder.internal.sparse.nullcopyLike(ONE,ncol,nalloc,identity);
y.colidx(1) = 1;
outidx = ONE;
for col = 1:x.n
    xstart = x.colidx(col);
    xendp1 = x.colidx(col+1);
    xend = xendp1-1;
    anyzero = xendp1-xstart < x.m;
    % Initialize
    if anyzero
        r = zeros('like',identity);
    else
        r = identity;
    end
    % Accumulate
    for xp = xstart:xend
        if ~(OMITNANS && isnan(x.d(xp)))
            r = f(r,cast(x.d(xp),'like',identity));
        end
    end
    if r ~= zeros('like',y.d)
        y.d(outidx) = r;
        outidx = outidx+1;
    end
    y.colidx(col+1) = outidx;
end
for k = 1:nnzInt(y)
    y.rowidx(k) = 1;
end

%--------------------------------------------------------------------------

function [y,counts] = sparseRowMapIter(f,identity,x,dim,omitnan)
% Row iteration accumulating results in an array having size(x,1) elements
coder.inline('always');
coder.internal.prefer_const(identity,dim,omitnan);
OMITNANS = isfloat(x) && omitnan;
counts = [];
if isequal(f,@times)
    yt = coder.nullcopy(zeros(x.m,1,'like',identity));
    rowcounts = nonzeroRowCounts(x);
    for k = 1:x.m
        if rowcounts(k) == x.n
            yt(k) = identity;
        else
            yt(k) = 0;
        end
    end
else
    yt = zeros(x.m,1,'like',identity);
end
nzx = nnzInt(x);
for xp = 1:nzx
    if ~(OMITNANS && isnan(x.d(xp)))
        yt(x.rowidx(xp)) = f(yt(x.rowidx(xp)),cast(x.d(xp),'like',identity));
    end
end
% Count non-zeros
nzy = zeros('like',ONE);
for k = 1:x.m
    if yt(k) ~= zeros('like',yt)
        nzy = nzy+1;
    end
end
y = coder.internal.sparse.nullcopyLike(x.m,ONE,nzy,identity);
y.colidx(1) = 1;
y.colidx(end) = nzy+1;
yidx = ONE;
for k = 1:x.m
    if yt(k) ~= zeros('like',yt)
        y.rowidx(yidx) = k;
        y.d(yidx) = yt(k);
        yidx = yidx+1;
    end
end

%--------------------------------------------------------------------------

function [y,counts] = sparseRowReduceIter(f,identity,x,dim,omitnan)
% Row iteration when allocating an array having size(x,1) elements may not
% be possible.
coder.inline('always');
coder.internal.prefer_const(identity,dim,omitnan);
if isfloat(x) && omitnan
    pred = @reducePredOmitNaNs;
    nzinit = @(x)reduceNzinitOmitNaNs(x,identity);
else
    pred = @reducePredIncludeNaNs;
    nzinit = @reduceNzinitIncludeNaNs;
end
counts = [];
y = rowReduction(f,pred,nzinit,x,identity);

%--------------------------------------------------------------------------

function y = ONE
coder.inline('always');
y = coder.internal.indexInt(1);

%--------------------------------------------------------------------------

function p = reducePredOmitNaNs(x)
coder.inline('always');
p = ~isnan(x);

%--------------------------------------------------------------------------

function p = reducePredIncludeNaNs(~)
coder.inline('always');
p = true;

%--------------------------------------------------------------------------

function y = reduceNzinitOmitNaNs(x,identity)
coder.inline('always');
coder.internal.prefer_const(identity);
if isnan(x)
    y = identity;
else
    y = x;
end

%--------------------------------------------------------------------------

function y = reduceNzinitIncludeNaNs(x)
coder.inline('always');
y = x;

%--------------------------------------------------------------------------
