function varargout = find(varargin)
%FIND Find indices of nonzero elements of codistributed array
%   I = find(X)
%   I = find(X,K)
%   I = find(X,K,'first')
%   I = find(X,K,'last')
%   [I,J] = find(X,...)
%   [I,J,V] = find(X,...)
%   
%   Example:
%   spmd
%       N = 1000;
%       X = codistributed.rand(N) > 0.5 % build random array of ones and zeros
%       I = find(X) % find the indices where elements of X are non-zero
%   end
%   
%   See also FIND, CODISTRIBUTED, CODISTRIBUTED/RAND.


%   Copyright 2006-2022 The MathWorks, Inc.

narginchk(1, 3);

distributedutil.CodistParser.verifyNonCodistributedInputs(varargin);

A = varargin{1};
argList = distributedutil.CodistParser.gatherElements(varargin(2:end));
if ~isa(A, 'codistributed')
    % Only the optional arguments were codistributed.
    [varargout{1:nargout}] = find(A, argList{:});
    return;
end

if nargin > 1
    k = argList{1};
    if ~isscalar(k) || (~isPositiveIntegerValuedNumeric(k) && k ~= inf)
        error(message('parallel:distributed:FindNotScalarInt'))
    end
end
if nargin > 2
    searchOpt = argList{2};
    if ~strcmpi(searchOpt,'first') && ~strcmpi(searchOpt,'last')
        error(message('parallel:distributed:FindInvalidOption'))
    end
end

needJ = nargout > 1;
needV = nargout > 2;
if nargout > 3
    error(message('parallel:distributed:FindMaxlhs'))
end

if isempty(A) && size(A,1) == 0 && size(A,2) == 0
    [varargout{1:nargout}] = handleEmptyAndScalarEdgeCase(A);
    return
end

aDist = getCodistributor(A);
if ~isa(aDist, 'codistributor1d') || (aDist.Dimension < ndims(A) && ~isvector(A))
    A = redistribute(A,codistributor('1d',ndims(A)));
    aDist = getCodistributor(A);
end
[e,ignore] = globalIndices(A, aDist.Dimension, spmdIndex); %#ok<ASGLU> % Ignore the second output argument.
m = size(A,1);
if needV
    [i, j, v] = find(getLocalPart(A),varargin{2:nargin});
else
    [i, j] = find(getLocalPart(A),varargin{2:nargin});
end
if aDist.Dimension == 1
    i = i+e-1;
elseif aDist.Dimension == 2
    j = j+e-1;
else
    p = aDist.Partition;
    j = j+sum(p(1:spmdIndex-1))*size(A,2);
end

% Get the sizes of all the local i's, j's and v's.
part = spmdCat(numel(i));

if isrow(A)
    % Input is a row vector, so i, j and v all have to be row vectors.
    sz = [1, sum(part)];
    dist = codistributor1d(2, part, sz);
    i = codistributed.pDoBuildFromLocalPart(i(:).', dist);
    j = codistributed.pDoBuildFromLocalPart(j(:).', dist);
    if needV
        v = codistributed.pDoBuildFromLocalPart(v(:).', dist);
    end
    
else
    % Input is not a row vector, so i, j and v all have to be column vectors.
    sz = [sum(part), 1];
    dist = codistributor1d(1, part, sz);
    i = codistributed.pDoBuildFromLocalPart(i(:), dist);
    j = codistributed.pDoBuildFromLocalPart(j(:), dist);
    if needV
        v = codistributed.pDoBuildFromLocalPart(v(:), dist);
    end
    
end
if nargin == 2 || ( nargin == 3 && strcmpi(searchOpt,'first') )
    if needV
        [i,j,v] = iGetFirstKVectorElements(k,i,j,v);
    else
        [i,j] = iGetFirstKVectorElements(k,i,j);
    end
elseif nargin == 3 % && strcmpi(searchOpt,'last')
    if needV
        [i,j,v] = iGetLastKVectorElements(k,i,j,v);
    else
        [i,j] = iGetLastKVectorElements(k,i,j);
    end
end

if nargout <= 1
    i = i + (j-1)*m;
end

if isscalar(A) && isempty(i)
    [varargout{1:nargout}] = handleEmptyAndScalarEdgeCase(A);
else
    varargout{1} = i;
    if needJ
        varargout{2} = j;
        if needV
            varargout{3} = v;
        end
    end
end
end

function varargout = handleEmptyAndScalarEdgeCase(A)
%handle special find edge case
dist = codistributor1d(codistributor1d.unsetDimension, codistributor1d.defaultPartition(0), [0, 0]);
varargout{1} = zeros(0, dist);
if nargout > 1
    varargout{2} = varargout{1};
    if nargout > 2
        % We construct an empty array of the same class this way instead of
        % using the 'like' syntax because this way supports non-numeric types.
        emptyA = getLocalPart(A);
        emptyA(:) = [];
        if issparse(emptyA)
            emptyA = full(emptyA);
        end
        varargout{3} = codistributed.build(emptyA, dist);
    end
end
end

function varargout = iGetFirstKVectorElements(k,varargin)
% Helper to extract the first K elements of a vector, assuming 1D distribution.

% All inputs must have exactly the same distribution, so we just use the
% first to work out the partitioning.
inDist = getCodistributor(varargin{1});
assert(isvector(varargin{1}));
inNumDims = ndims(varargin{1});
dim = inDist.Dimension;

% Work out how much data is in each local part. Any partition that
% starts after K is left empty.
partEnd = cumsum(inDist.Partition);
newPartEnd = min(partEnd, k);
newPart = diff([0, newPartEnd]);
% Since we are dealing with a distributed vector, all dimensions are 1
% except the distribution dimension.
newSize = ones(1, inNumDims);
newSize(dim) = sum(newPart);
newDist = codistributor1d(dim, newPart, newSize);
elementsForThisPart = newPart(spmdIndex);

% Now work out the new local part for each input
varargout = cell(1,numel(varargin));
for ii=1:numel(varargin)
    localPart = getLocalPart(varargin{ii});
    newLocalPart = iVectorIdx(localPart, dim, 1, elementsForThisPart);
    varargout{ii} = codistributed.pDoBuildFromLocalPart(newLocalPart, newDist);
end
end

function varargout = iGetLastKVectorElements(k,varargin)
% Helper to extract the last K elements of a vector, assuming 1D distribution.

% All inputs must have exactly the same distribution, so we just use the
% first to work out the partitioning.
inDist = getCodistributor(varargin{1});
assert(isvector(varargin{1}));
inNumDims = ndims(varargin{1});
dim = inDist.Dimension;

% Work out the first and last element in each local part.
firstElement = cumsum([1, inDist.Partition(1:end-1)]);
lastElement = cumsum(inDist.Partition);
firstNeededElement = max(1, lastElement(end)-k+1);

% Now work out which elements from each local part are needed
firstElementInPart = max(1, firstNeededElement - firstElement + 1);
lastElementInPart = inDist.Partition;
newPart = max(0, lastElementInPart - firstElementInPart + 1);
% Since we are dealing with a distributed vector, all dimensions are 1
% except the distribution dimension.
newSize = ones(1, inNumDims);
newSize(dim) = sum(newPart);
newDist = codistributor1d(dim, newPart, newSize);

% Now create the new local part for each input
varargout = cell(1, numel(varargin));
for ii=1:numel(varargin)
    localPart = getLocalPart(varargin{ii});
    newLocalPart = iVectorIdx(localPart, dim, firstElementInPart(spmdIndex), lastElementInPart(spmdIndex));
    varargout{ii} = codistributed.pDoBuildFromLocalPart(newLocalPart, newDist);
end
end

function out = iVectorIdx(in,dim,first,last)
% Extract elements from a vector, preserving the orientation
if dim==1 % Column vector
    out = in(first:last, :);
else % Row vector
    out = in(:, first:last);
end
end