function [idx, distance] = knnsearch(X,Y,varargin)
%KNNSEARCH Find K nearest neighbors.
%   IDX = KNNSEARCH(X,Y)
%   [IDX,D] = KNNSEARCH(X,Y)
%   [IDX,D] = KNNSEARCH(X,Y,'NAME1',VALUE1,...,'NAMEN',VALUEN)
%
%   'IncludeTies' and 'NSMethod' name arguments are not supported.
%
%   Example:
%     % Find 2 nearest neighbors in X and the corresponding values to each
%     % point in Y using the distance metric 'cityblock'
%     X = gpuArray.randn(100,5);
%     Y = gpuArray.randn(25,5);
%     [idx, dist] = knnsearch(X,Y,'dist','cityblock','k',2);
%
% See also knnsearch, gpuArray

% Copyright 2015-2023 The MathWorks, Inc.

% error checking is done here to output the correct error IDs. We will call
% directly at pdist2mex to calculate the indices and (optionally) the
% distances
[nX,nDims] = size(X);
[nY,nDims2]= size(Y);

if nDims2 ~= nDims
    error(message("stats:KDTreeSearcher:knnsearch:SizeMisMatch", nDims));
end

pnames = [ "k" "Distance"  "nsmethod"   "bucketsize", "includeties" "sortindices" "cachesize", "P" "Cov" "Scale"];
dflts =  { 1   "euclidean" "exhaustive" []            false          true         1e3          []  []    []};

[numNN, distMetric, nsmethod, ~, includeTies, doSort, ~, minExp, mahaCov, seucInvWgt] = ...
    internal.stats.parseArgs(pnames, dflts, varargin{:});

% nsmethod, if specified, must be 'exhaustive'
if ~matlab.internal.datatypes.isScalarText(nsmethod)
    error(message("stats:createns:BadMethod"))
else
    nsmethodNames = ["exhaustive","kdtree","hnsw"];
    nsMethodIdx = strncmpi(nsmethod,nsmethodNames,1);
    if ~any(nsMethodIdx)
        error(message("stats:createns:UnknownMethod", nsmethod));
    end
    nsmethod = nsmethodNames(nsMethodIdx);

    if nsmethod ~= "exhaustive"
        error(message("stats:gpu:KnnsearchBadNSMethod"));
    end
end

% includeTies, if specified, must be false.
if ~islogical(includeTies)
     error(message("stats:ExhaustiveSearcher:knnsearch:InvalidIncludeTies"));
elseif includeTies
     error(message("stats:gpu:KnnsearchBadIncludeTies"));
end

% integer/logical/char/anything data will be converted to float. Complex
% floating point data can't be handled.
try
    outClass = superiorfloat(X,Y);
catch
    outClass = underlyingType(X);
    warning(message('stats:KDTreeSearcher:knnsearch:DataConversion', underlyingType(Y), outClass));
end

% check the input types
Y = cast(Y,outClass);
if ~strcmp(outClass, underlyingType(X))
    % only happens when X is double and Y is single
    X = cast(X, outClass);
end

if ~isreal(Y)
    error(message("stats:KDTreeSearcher:knnsearch:ComplexData"));
end

if issparse(Y)
    warning(message("stats:KDTreeSearcher:knnsearch:DataConversion", "sparse", "full"));
    Y = full(Y);
end

% check the number of nearest neighbors requested
if ~isscalar(numNN) || ~isnumeric(numNN) ||  ...
        numNN < 1 || numNN~=round(numNN)
    error(message("stats:KDTreeSearcher:knnsearch:BadK"));
end

numNN = min(numNN,nX);

% degenerate case, just return an empty matrix or cell of the proper size.
if (nY == 0 || numNN ==0)
    idx = gpuArray.zeros(nY, numNN,outClass);
    distance = gpuArray.zeros(nY, numNN, outClass);
    return;
end

% distance metric
distList = ["euclidean"; "cityblock"; "chebychev"; "minkowski";...
    "fasteuclidean"; "fastseuclidean";...
    "mahalanobis"; "seuclidean"; "cosine"; "correlation"; ...
    "spearman"; "hamming"; "jaccard"];
distMetric = statslib.internal.getParamVal(distMetric,distList,"Distance");

% calculate the additional argument for minkowski, seuclidean, and
% mahalanobis distances
additionalArg = [];

% minkowski -> exponent
if strncmp(distMetric,"minkowski",3)
    if  ~isempty(minExp)
        additionalArg = gather(minExp); % the exponent must be a scalar
    else
        additionalArg = 2;
    end
elseif ~isempty(minExp)
    error(message("stats:ExhaustiveSearcher:InvalidMinExp"));
end

% seuclidean -> scale
if strncmp(distMetric,"seuclidean",3)
    if ~isempty(seucInvWgt)
        additionalArg = seucInvWgt;
    else
        additionalArg = std(X,[],1, "omitnan");
    end
elseif ~isempty(seucInvWgt)
    error(message("stats:ExhaustiveSearcher:InvalidScale"));
end

% mahalanobis -> cov matrix
if strncmp(distMetric, "mahalanobis",3)
    if ~isempty(mahaCov)
        if ~isequal(size(mahaCov),[nDims,nDims])
            error(message("stats:ExhaustiveSearcher:BadCov"));
        end
        %use cholcov because we also need to check whether the matrix is symmetric
        [~,flag] = cholcov(mahaCov,0);
        if flag ~= 0
            error(message("stats:ExhaustiveSearcher:SingularCov"));
        end
        additionalArg = mahaCov;
    else
        if nX == 1
            error(message("stats:ExhaustiveSearcher:TooFewXRowsForMah"));
        end
        additionalArg = cov(X, "omitrows");
    end
elseif ~isempty(mahaCov)
    error(message("stats:ExhaustiveSearcher:InvalidCov"));
end

% numNN must be a scalar on the host
numNN = gather(numNN);
arg = {};
if ~isempty(additionalArg)
    arg = {additionalArg};
end

[distance,idx] = pdist2(X,Y, distMetric, arg{:}, "sortIndices", doSort, "smallest",numNN);

distance = distance';
idx = idx';

end
