function C = convDistDist(A, B, convFcnHandle, outClass, isResultRow)
% Perform convolution of two distributed arrays A and B with appropriate
% shape. isResultRow is an optional parameter used by conv (1D convolution)
% to return a result with the right shape.  This routine returns the local
% part of A convolved with all of B.

% Copyright 2016-2022 The MathWorks, Inc.

codistrB = getCodistributor(B);
partB = codistrB.Partition;
cumPartB = cumsum(partB);

szB = codistrB.Cached.GlobalSize;
dim = codistrB.Dimension;

isConv = isequal(convFcnHandle, @conv);
if isConv
    assert(nargin==5, "Expected 5 inputs to CONV");
    [dim,szB] = iReshapeToCat(szB,isResultRow);
    if isvector(B) && cumPartB(end) == 1
        partB = partB*szB(dim);
    end
end

firstIteration = true;

% For every worker that contains a part of B, the algorithm goes through it by
% turn and then replicates the data across all workers. Following this it
% calls convDistRepl, stores the answer and adds it to an accumulator C.
for ii = 1:spmdSize
    if partB(ii) == 0
        % Local part empty so skip
        continue;
    end

    if ii == spmdIndex
        dataShared = spmdBroadcast(ii,getLocalPart(B));
    else
        dataShared = spmdBroadcast(ii);
    end

    if isConv
        temp = convDistRepl(A,dataShared,convFcnHandle,outClass,isResultRow);
    else
        temp = convDistRepl(A,dataShared,convFcnHandle,outClass);
    end

    if firstIteration % Set up accumulator C
        szC = [size(temp) ones(1,dim-ndims(temp))];
        if szC(dim) == 0
            szC(dim) = szB(dim);
        else
            szC(dim) = szC(dim) - partB(ii) + szB(dim);
        end
        C = zeros(szC, 'like', temp);
        firstIteration = false;
        C = assignSlices(C, 1:size(temp, dim), temp, dim);
    else
        offset = cumPartB(ii-1);
        changedSlices = offset + (1:size(temp, dim));
        priorC = extractSlices(C, changedSlices, dim);
        C = assignSlices(C, changedSlices, priorC + temp, dim);
    end
end
end

function [dim,szB] = iReshapeToCat(szB,isResultRow)
% Reshape szB to correctly pad the result

if isResultRow
    szB = sort(szB,'ascend');
    dim = 2;
else
    szB = sort(szB,'descend');
    dim = 1;
end
end
