function [localX, szLeftPad, localT] = stencilCommSamples(X, padSz, T)
% STENCILCOMMSAMPLES Adds padding to local parts of a distributed array X,
% returning the padded array and the number of elements added at its
% beginning (left end). T provides sample points for the calculation. A
% subset of T which contains the sample points that correspond to the
% values in localX is returned as localT. X can be an array of any
% dimensionality distributed 1-d. T must be a column vector distributed
% with the same partition as X if it is distributed. The communication
% pattern is optimized to avoid forwarding data. The algorithm used here
% must be kept in-sync with the similar algorithm in stencilCommConv and
% stencilCommPad.

% Example: consider an array spread over four workers:
%
% (value)        X: 1  2  3  |  4   5   6   |  |  7  8  9
% (sample point) T: 1  2  3  |  3.5 4.0 4.5 |  |  5  6  7
%
% If padSz is [2 2] (k=4), then
% localX: 1  2  3  4  5  6  7  | 2  3  4  5  6  7  8  |     |  3  4  5  6  7  8  9
% szLeftPad:       0           |          2           |  0  |           4
% localT: 1  2  3 3.5 4 4.5 5  | 2  3 3.5 4 4.5 5  6  |     |  3 3.5 4 4.5 5  6  7
%
%   Copyright 2022 The MathWorks, Inc.

szLeftPad = 0;
ourParts = cell(2, spmdSize);
needToRecv = false(1, spmdSize);
localX = getLocalPart(X);
codX = getCodistributor(X);
partitionX = codX.Partition;
cumPartitionX = [0 cumsum(partitionX)];
szXInDim = cumPartitionX(spmdSize+1);
partToSend = cell(2, 1);
dim = codX.Dimension;
distT = false;

% We need the minimum and maximum sample point values in each partition,
% whether T is distributed or not.
if iscodistributed(T)
    distT = true;
    localT = getLocalPart(T);
    if ~isempty(localT)
        limitsOfT = [localT(1), localT(end)];
    else
        limitsOfT = [NaN, NaN];
    end
    limT = spmdCat(limitsOfT, 1);
else
    localT = T((cumPartitionX(spmdIndex)+1):cumPartitionX(spmdIndex+1));
    limT = [T(min(cumPartitionX(1:end-1)+1, szXInDim)), T(cumPartitionX(2:end))];
end

% Figure out who is sending to us.
if ~isempty(localT)  % Receive only if we contribute to result.
    minT = localT(1) - padSz(1);
    maxT = localT(end) + padSz(2);
    for spmdIx = 1:spmdIndex-1  % Loop over workers to our left.
        needToRecv(spmdIx) = minT <= limT(spmdIx, 2) && partitionX(spmdIx);
    end
    for spmdIx = spmdIndex+1:spmdSize  % Loop over workers to our right.
        needToRecv(spmdIx) = maxT >= limT(spmdIx, 1) && partitionX(spmdIx);
    end
else  % if localT is empty, localX is also, so nothing to send
    return
end

% Do the sends and receives.
mwTag = 2^15 - mod(sum(mfilename),2^15);
for spmdIx = 1:spmdSize
    if partitionX(spmdIx) == 0  % No need to send to or receive from empty workers.
        continue
    end
    if spmdIx < spmdIndex
        maxT2Send = limT(spmdIx, 2) + padSz(2);
        startIndex = 1;
        endIndex = find(localT<=maxT2Send, 1, 'last');
    elseif spmdIx == spmdIndex
        continue  % Don't need to send anything to ourself.
    else % spmdIx > spmdIndex
        minT2Send = limT(spmdIx, 1) - padSz(1);
        startIndex = find(localT>=minT2Send, 1);
        endIndex = partitionX(spmdIndex);
    end
    % Communication happens here.
    needToSend = endIndex>=startIndex;
    if needToSend
        partToSend{1} = extractSlices(localX, startIndex:endIndex, dim);
        if distT
            partToSend{2} = localT(startIndex:endIndex);
        end
        if needToRecv(spmdIx)
            ourParts(:, spmdIx) = spmdSendReceive(spmdIx, spmdIx, partToSend, mwTag+spmdIndex+spmdIx+spmdSize);
        else
            spmdSend(partToSend, spmdIx, mwTag+spmdIndex);
        end
    elseif needToRecv(spmdIx)
        ourParts(:, spmdIx) = spmdReceive(spmdIx, mwTag+spmdIx);
    else
        % Here if nothing to send or receive
    end
end

% Attach the received elements to localX and track how many were added to
% the beginning of localX so we can remove them later.
szUnPaddedLocalX = size(localX, dim);
for spmdIx = spmdIndex-1:-1:1  % Apply padding to left side of localX.
    localX = cat(dim, ourParts{1, spmdIx}, localX);
end
szLeftPad = size(localX, dim) - szUnPaddedLocalX;
for spmdIx = spmdIndex+1:spmdSize  % Apply padding to right side of localX.
    localX = cat(dim, localX, ourParts{1, spmdIx});
end

if distT
    leftPadT = cat(1, ourParts{2, 1:spmdIndex-1});
    rightPadT = cat(1, ourParts{2, spmdIndex+1:spmdSize});
    localT = [leftPadT; localT; rightPadT];
else
    localT = T(cumPartitionX(spmdIndex)-szLeftPad+(1:size(localX, dim)));
end
end
