function [Y, I] = topkrowsImpl(A, k, col, varargin)
%TOPKROWSIMPL Implementation of TOPKROWS for codistributed arrays
%   We handle two cases: 
%   (1) A is distributed by columns
%   (2) A is distributed by rows and k is small enough so that the results
%   of local calls to topkrows can be combined by a spmdReduce
% In case (1) we sort on the set of columns that precede the first column
% that is not on worker 1.  If adjacent rows in the sorted array are
% identical, they are "tied", so sorting on columns not on worker 1 is
% necessary.  Thus, if any "ties" are true, we need to so a full sortrows
% and return the top k of those.

%   Copyright 2019-2022 The MathWorks, Inc.

LPA = getLocalPart(A);
codistrA = getCodistributor(A);
distDim = codistrA.Dimension;
wantI = nargout == 2;
if distDim > 1   % case (1)
    if spmdIndex == 1
        localCol = codistrA.Partition(1);
        idxLastSortedCol = find(col>localCol, 1)-1;
        if ~isempty(idxLastSortedCol)  % Some will not be sorted on.
            if idxLastSortedCol == 0 % Can't sort here
                LPI = (1:k)';
                ties = true(k, 1);
            else
                TrailingArgs = iAdjustArgsForTopkrows(col(1:idxLastSortedCol),...
                    varargin{:});
                [LPY, LPI] = topkrows(LPA, k+1, TrailingArgs{:});
                % Find out if we have any ties
                sortedCols = col(1:idxLastSortedCol);
                ties = diff(LPY(:, sortedCols))==0;
                ties = all(ties, 2);   % rows that match
            end
        else  % We can do a complete sort here
            TrailingArgs = iAdjustArgsForTopkrows(col, varargin{:});
            [LPY, LPI] = topkrows(LPA, k, TrailingArgs{:});
            ties = false(k, 1);
        end
        IandTie = [LPI(1:k) ties];
        spmdBroadcast(1, IandTie);
        if ~any(ties)
            LPY = LPY(1:k, :);
            LPI = LPI(1:k);
        end   
    else
        IandTie = spmdBroadcast(1);
        ties = IandTie(:, 2);
        if ~any(ties)
            LPI = IandTie(:, 1);
            LPY = LPA(LPI, :);
        end
    end
    if any(ties) % We need to do a full sort
        [LPY, LPI] = codistrA.hSortrowsGeneralImpl(LPA, col, varargin);
        LPY = LPY(1:k, :);
        LPI = LPI(1:k);
    end
else % case (2)
    % First the workers independently calculate the topkrows on their local part.
    if wantI
        [LPY, LPI] = topkrows(LPA, k, varargin{:});
    else
        LPY = topkrows(LPA, k, varargin{:});
    end
    [offset, ~] = globalIndices(codistrA, 1);

    % Then we combine results from different workers with spmdReduce.
    if wantI
        localYI = [{LPY}, {LPI + offset - 1}];
        topkrowsfun = @(x, y) iTopkrowsReduce(x, y, k, varargin);
        localYI = spmdReduce(topkrowsfun, localYI);
        LPY = localYI{1};
        LPI = localYI{2};
    else
        topkrowsfun = @(x, y) topkrows([x; y], k, varargin{:});
        LPY = spmdReduce(topkrowsfun, LPY);
    end
end
[codistrY, codistrI] = codistrA.hGetCodistForTopK(k, 1, distDim);
Y = codistributed.pDoBuildFromLocalPart(LPY, codistrY);
if wantI
    I = codistributed.pDoBuildFromLocalPart(LPI, codistrI);
end

end

function LI = iTopkrowsReduce(x, y, k, sortArgs)
% Reduction function for partial results of distributed topkrows.
% x and y are cell arrays containing {Y} and {I} from earlier topkrows calls.
% Input I's are global indices.

[outL, I] = topkrows([x{1}; y{1}], k, sortArgs{:});

inputIndices = [x{2}; y{2}];
outI = inputIndices(I, 1);

LI = [{outL}, {outI}];
end

function TrailingArgs = iAdjustArgsForTopkrows(col, varargin)
% This function makes sure that the arguments are appropriate for a call to
% topkrows on the local part
TrailingArgs = [{col}, varargin];
idxLastSortedCol = numel(col);
dir = varargin{1};
if iscell(dir) && numel(dir) >= idxLastSortedCol
    TrailingArgs(2) = {dir(1:idxLastSortedCol)};
end
% Remove "MissingPlacement"
location = find(strcmp(varargin, 'MissingPlacement'));
if ~isempty(location)  % Remove it and its value.
    TrailingArgs(location+1:location+2) = [];
end
end
