function varargout = topkrows(A, k, varargin)
%TOPKROWS Top rows in sorted order
%   Y = TOPKROWS(X,K)
%   Y = TOPKROWS(X,K,COL)
%   Y = TOPKROWS(...,DIRECTION)
%   Y = TOPKROWS(...,'ComparisonMethod',C)
%   [Y,I] = TOPKROWS(...)
%   
%   Example:
%   spmd
%       N = 1000;
%       D = codistributed(magic(N))
%       m = topkrows(D,10)
%       m1 = topkrows(D,10,10)
%       m2 = topkrows(D,10,'ComparisonMethod','real')
%   end
%   
%   m, m1 and m2 are a set of K codistributed row vectors.
%   
%   See also TOPKROWS, CODISTRIBUTED, MAGIC.
%   


%   Copyright 2018-2022 The MathWorks, Inc.

narginchk(2, inf);
nargoutchk(0, 2);

% Make sure all arguments other than A are gathered.
distributedutil.CodistParser.verifyNonCodistributedInputs([{A}, {k}, varargin]);
argList = distributedutil.CodistParser.gatherElements([{k}, varargin]);

% If A isn't codistributed, call MATLAB function.
if ~iscodistributed(A)
    % Only the nondata arguments were codistributed.
    [varargout{(1:nargout)}] = topkrows(A, argList{:});
    return
end

% Use the local function to check syntax.
distributedutil.validateSyntax(@topkrows, [{A}, argList]);
[k, col, sortArgs] = iParseArgs(A, k, argList{2:end});
% If not 1D, redistribute to 1D.
codistrA = getCodistributor(A);
if ~isa(codistrA, 'codistributor1d')
    A = redistribute(A, codistributor1d);
    codistrA = getCodistributor(A);
end
distDim = codistrA.Dimension;

% Quick return if A is empty or not sorting any columns.
if isempty(A) || isempty(col)
    % Return A unmodified, but might need to create I
    varargout(1) = {A};
    if ~isempty(A)
        LPY = iDoSubsref(A, k);
        codistrY = codistrA.hGetNewForSize([k, size(A, 2)]);
        varargout(1) = {codistributed.pDoBuildFromLocalPart(LPY, codistrY)};
    end
    if nargout == 2
        I = (1:min(size(A, 1), k))'; % Always a column vector
        codistrI = codistrA.hGetNewForSize(size(I));
        varargout(2) = {codistributed.pConstructFromReplicated(I, codistrI)};
    end
    return
end

if k == size(A, 1)  % Just do sortrows if all results are tops.
    [varargout{(1:nargout)}] = sortrows(A, col, sortArgs{:});
    return
end

if distDim == 1 && k > size(A, 1)/spmdSize
    A = redistribute(A, codistributor1d(2));
end
[varargout{(1:nargout)}] = topkrowsImpl(A, k, col, sortArgs{:});
end

function [k, col, sortArgs] = iParseArgs(A, k, varargin)
% iParseArgs returns k (required arg) and default values for other args if
% they weren't specified.

% k shouldn't be greater than the size of A in the sorting dimension.
k = min(k, size(A, 1));
    
isSortableType = isnumeric(A) ...
    || islogical(A) || ischar(A) || iscellstr(A); %#ok<ISCLSTR>
if ~isSortableType
    ex = MException(message('MATLAB:topkrows:MustNotBeCell'));
    throwAsCaller(ex);
elseif isinteger(A) && ~isreal(A)
    ex = MException(message('MATLAB:topkrows:complexIntegers'));
    throwAsCaller(ex);
end

sortArgs = varargin;
% Capture col if supplied and supply if not.
if ~isempty(sortArgs) && isnumeric(sortArgs{1})
    col = sortArgs{1};
    sortArgs(1) = [];
    if iscolumn(col)
        col = col';
    end
elseif iscellstr(A) %#ok<ISCLSTR>
    col = 1;
else
    col = 1:size(A, 2);
end

hasdir = false;
% If we have an even number of arguments left, the next argument
% should be the direction, with the rest as param-value pairs.
if mod(numel(sortArgs), 2)==0
    sortArgs = [{'descend'}, sortArgs];
    hasdir = true;
end

% Check for strings ascend, descend, or cell containing them.
for ii = 1:numel(sortArgs)
    if matlab.internal.math.checkInputName(sortArgs{ii}, {'ascend'}) || ...
            matlab.internal.math.checkInputName(sortArgs{ii}, {'descend'}) || ...
            iscell(sortArgs{ii}) || (isstring(sortArgs{ii}))
        hasdir = true;
    end
end

codistrA = getCodistributor(A);
distDim = codistrA.Dimension;
if iscellstr(A) || ischar(A) || ...
        (distDim == 1 && k <= size(A, 1)/spmdSize && k ~= size(A, 1)) %#ok<ISCLSTR>
    % We will be calling topkrows.
    if ~hasdir
        sortArgs = {'descend'};
    end
else
    % We will be calling sortrows.
    if hasdir
        sortArgs = [sortArgs, {'MissingPlacement'}, {'last'}];
    else
        sortArgs = [{'descend'}, sortArgs, {'MissingPlacement'}, {'last'}];
    end
end
end

function [LPY, LPI] = iDoSubsref(Y, k, I)
% Function to carry out subsref operation on distributed Y and I 
wantI = nargout == 2;
codistrY = getCodistributor(Y);
LPY = getLocalPart(Y);
if codistrY.Dimension == 2
    if ~isempty(LPY)
        LPY = LPY(1:k, :);
    end
    if wantI
        LPI = getLocalPart(I);
        if ~isempty(LPI)
            LPI = LPI(1:k);
        end
    end
else
    gIndices = globalIndices(Y, 1);
    if ~isempty(LPY)
        LPY = LPY((gIndices <= k), :);
    end
    if wantI
        LPI = getLocalPart(I);
        if ~isempty(LPI)
            LPI = LPI(gIndices <= k);
        end
    end
end
end
