function [B,I] = topkrows(X,k,varargin)
%MATLAB Code Generation Library Function

%   Limitations:
%       The first input cannot be a cell array, table, categorical,
%       duration, or datetime


%   Copyright 2019-2022 The MathWorks, Inc.
%#codegen
coder.inline('always');
coder.internal.prefer_const(k,varargin);
narginchk(2,inf);
if nargout == 2
    [B,II] = TopKRows(X,k,varargin{:});
    I = double(II);
else
    B = TopKRows(X,k,varargin{:});
end
end


function [B,I] = TopKRows(X,k,varargin)
coder.inline('never');
coder.internal.prefer_const(k,varargin);
ONE = coder.internal.indexInt(1);

% Parse inputs
narginchk(2,inf);

% Validate type of X
coder.internal.assert(~iscell(X),'MATLAB:topkrows:CellsNotSupported');
coder.internal.assert(coder.internal.isBuiltInNumeric(X) || islogical(X) ...
    || ischar(X) || isa(X,'categorical') || isduration(X),'MATLAB:topkrows:InvalidX');

% X must be a vector or a matrix
coder.internal.assert(ismatrix(X),'MATLAB:topkrows:MustBeMatrix')

% X cannot be a complex integer
coder.internal.assert(isreal(X) || isfloat(X) || isstring(X) ||...
    isa(X,'categorical') ,'MATLAB:topkrows:complexIntegers');

% k must be numeric and a non-negative, real, finite, scalar, integer
coder.internal.assert(coder.internal.isBuiltInNumeric(k) && isreal(k) && ...
    isscalar(k) && k>=0 && round(k)==k && isfinite(k),'MATLAB:topkrows:InvalidK')
numRows = coder.internal.indexInt(size(X,1));
numColumns = coder.internal.indexInt(size(X,2));

colDefined = false;
numColToSort = numColumns;
method = AUTO; % Default comparison method

% Parse optional inputs
if(nargin>2)
    nextArg = ONE;
    lastArg = nargin-2;
    % Check for column input
    if ~coder.internal.isTextRow(varargin{nextArg}) && ~iscellstr(varargin{nextArg})
        % Validate column
        % Skip if input is empty
        if ~isempty(varargin{nextArg})
            % Column must be a vector of positive integer values,
            % and the maximum value cannot be greater than the
            % number of columns in X
            coder.internal.assert(isnumeric(varargin{nextArg}) && isvector(varargin{nextArg}) &&...
                isreal(varargin{nextArg}) && all(round(varargin{nextArg}) == varargin{nextArg}) && ...
                all(varargin{nextArg}>0) && max(varargin{nextArg})<=numColumns, 'MATLAB:topkrows:ColNotIndexVec');
            % Negative columns indicate sorting in descending order
            col = -1*varargin{nextArg};
            colDefined = true;
            numColToSort = coder.internal.indexInt(numel(col));
        else
            col = [];
        end
        nextArg = nextArg + ONE;
    else
        % If not column input is given, use all columns
        col = -1*(1:numColToSort); 
    end
    % Check for direction input
    if nextArg <= lastArg && (coder.internal.isTextRow(varargin{nextArg}) || iscellstr(varargin{nextArg}))
        if ~iscell(varargin{nextArg})
            len = coder.internal.indexInt(strlength(varargin{nextArg}));
            numDirGiven = ONE;
        else
            numDirGiven = coder.internal.indexInt(numel(varargin{nextArg}));
        end
        if iscell(varargin{nextArg}) || strncmpi(varargin{nextArg},'Ascend',len) || ...
                strncmpi(varargin{nextArg},'Descend',len)
            if colDefined
                coder.internal.assert(isvector(varargin{nextArg}) && ...
                    (numDirGiven == numColToSort || numDirGiven == ONE),'MATLAB:topkrows:NumDirectionsFourth');
            else
                coder.internal.assert((numDirGiven == numColToSort || numDirGiven == ONE),'MATLAB:topkrows:NumDirectionsThird');
            end
            % Check that directions are valid
            if numDirGiven>ONE
                coder.unroll();
                for jj = 1:numDirGiven
                    col(jj) = col(jj)*checkDirection(varargin{nextArg}{jj});
                    % As arrays of strings are not supported in codegen,
                    % direction must be given as a cell array to make it to
                    % this point
                end
            else
                col = col*checkDirection(varargin{nextArg});
            end
            nextArg = nextArg + ONE;
        end
    end
    % Check for comparison method input
    if nextArg <= lastArg
        method = parseNVPairs(varargin{nextArg:end});
    end
else
    col = -1*(1:numColToSort);
end

intK = eml_min( ...
    coder.internal.indexInt(size(X,1)), ...
    coder.internal.indexInt(k));

% Return appropriately sized empty matrices if k is 0 or if x is blank
if intK == 0
    if isempty(X)
        B = X;
    else
        B = zeros(0,numColumns,'like',X);
    end
    I = zeros(0,1,coder.internal.indexIntClass);
    return
end

% Begin sorting
% Initialize heap by sorting first k rows
heap = buildHeap(X, intK, col, method);

% Iterate over remaining rows, adding each new row to the heap if it is
% larger than the lowest row currently in the heap
for ii = intK+ONE:numRows
    heap = checkAndInsertHeap(X,ii, intK, heap,col,method);
end

I = coder.nullcopy(heap');
if ~issparse(X)
    B = coder.nullcopy(zeros(numel(heap),size(X,2),'like',X));
else
    B = cast(sparse(double(intK),double(numColumns)),class(X));
end
% Sort heap
numToRemove = ONE;
for ii = intK:-1:ONE
    I(ii) = heap(ONE);
    B(ii,:) = X(heap(ONE),:);
    tmp = heap(ii);
    heap(ii) = heap(ONE);
    heap(ONE) = tmp;
    % By decreasing the k given to percolateDown, a subset of the heap can
    % be sorted each iteration without copying the heap
    heap = percolateDown(heap,ONE,X,intK-numToRemove,col,method);
    numToRemove = numToRemove+ONE;
end
end
%--------------------------------------------------------------------------

function heap = checkAndInsertHeap(X,newRow,k, heap,col,method)
coder.internal.prefer_const(k,col,method);
ONE = coder.internal.indexInt(1);
% See if new row is higher than the lowest row in the heap
newRowIsHigher = rowIsHigher(X,newRow,heap(ONE),col,method);
if newRowIsHigher
    % Pop the lowest row out of the heap and replace it with the new one
    heap(ONE) = newRow;
    heap = percolateDown(heap,ONE,X,k,col,method);
end
end

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

function heap = buildHeap(X,k,col,method)
coder.internal.prefer_const(k,col,method)
ONE = coder.internal.indexInt(1);
% Initialize heap
heap = coder.internal.indexInt(k:-1:1);
% This is a min-heap, so initializing it backwards reduces the cost when dealing with
% already sorted data
lastNonLeafNode = coder.internal.indexInt(ceil(single(k-1)/2));
% Build a min-heap of using the first k row indices
for ii = lastNonLeafNode:-1:ONE
    heap = percolateDown(heap,ii,X,k,col,method);
end

end

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

function heap = percolateDown(heap,parent,X,k,col,method)
coder.internal.prefer_const(k,col,method);
ONE = coder.internal.indexInt(1);

valuesSwapped = false;
firstIteration = true;
firstChild = coder.internal.indexInt(2*(parent-1)+2);
% Only check the next level of the tree if it contains the children of
% the current node or if a swap took place in the previous level
while firstChild<=k && (firstIteration || valuesSwapped)
    if firstIteration
        firstIteration = false;
    end
    % Compare to children
    child = firstChild;
    % Check if there are two children, compare to smaller child
    if firstChild+ONE <= k && rowIsHigher(X,heap(firstChild),heap(firstChild+ONE),col,method)
        child = firstChild+1;
    end
    if rowIsHigher(X,heap(parent),heap(child),col,method)
        % If parent is higher, perform swap
        tmp = heap(parent);
        heap(parent) = heap(child);
        heap(child) = tmp;
        valuesSwapped = true;
        parent = child;
        firstChild = coder.internal.indexInt(2*(parent-1)+2);
    else
        valuesSwapped = false;
    end
end
end


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

function out = checkDirection(in)
% Output -1 for ascending, 1 for descending (the order is descending by
% default, so this equates to fliping the sort order for ascending or
% keeping it the same for descending)
len = coder.internal.indexInt(strlength(in));
isAscend = strncmpi(in,'Ascend',len);
isDescend = strncmpi(in,'Descend',len);
coder.internal.assert(isAscend || isDescend,'MATLAB:topkrows:SortRowsMode');
out = 1;
if isAscend
    out = -1;
end
end

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

function tf = rowIsHigher(X,row1,row2,col,method)
coder.internal.prefer_const(col,method);
% If there is a tie, choose the higher row index. This is needed to make
% the sort stable
if row1<row2
    tf = coder.internal.sortLE(X,col,row1,row2,method,'L');
else
    tf = ~coder.internal.sortLE(X,col,row2,row1,method,'L');
end
end

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

function method = parseNVPairs(varargin)
coder.internal.prefer_const(varargin);
ONE = coder.internal.indexInt(1);

coder.unroll();
for ii = ONE:2:nargin
    % Check that a value is present
    coder.internal.assert(ii+ONE<=nargin,'MATLAB:topkrows:NameValuePairs');
    % Check name parameter
    coder.internal.assert(coder.internal.isTextRow(varargin{ii}),'MATLAB:topkrows:NameValueNames');
    len = coder.internal.indexInt(strlength(varargin{ii}));
    coder.internal.assert(strncmpi(varargin{ii},'ComparisonMethod',eml_max(len,1)),'MATLAB:topkrows:NameValueNames');
    % Check value parameter
    coder.internal.assert(coder.internal.isTextRow(varargin{ii+ONE}),'MATLAB:topkrows:InvalidComparison');
    len = coder.internal.indexInt(strlength(varargin{ii+ONE}));
    ISAUTO = strncmpi(varargin{ii+ONE},'auto',eml_max(len,2));
    ISREAL = strncmpi(varargin{ii+ONE},'real',eml_max(len,1));
    ISABS = strncmpi(varargin{ii+ONE},'abs',eml_max(len,2));
    coder.internal.assert(ISAUTO || ISREAL || ISABS, ...
        'MATLAB:topkrows:InvalidComparison');
    if ISREAL
        method = REAL;
    elseif ISABS
        method = ABS;
    else
        method = AUTO;
    end
end
end

% Enumerate comparision method options

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

function m = AUTO()
coder.inline('always');
m = 'A';
end

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

function m = REAL()
coder.inline('always');
m = 'L';
end


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

function m = ABS()
coder.inline('always');
m = 'M';
end