function [lia, locb] = ismember(A, B, varargin)
%ISMEMBER True for member of codistributed set
%   LIA = ISMEMBER(A,B)
%   LIA = ISMEMBER(A,B,'rows')
%   [LIA,LOCB] = ISMEMBER(A,B)
%   [LIA,LOCB] = ISMEMBER(A,B,'rows')
%   
%   Limitation:
%   The 'legacy' option is not supported.
%   Table and timetable inputs are not supported.
%   
%   See also ISMEMBER, CODISTRIBUTED.
%   
%   


%   Copyright 2012-2023 The MathWorks, Inc.

% Allow 4 inputs so the right error can be thrown for the "legacy" flag"
narginchk(2, 4)
distributedutil.CodistParser.verifyNonCodistributedInputs(varargin);

argList = distributedutil.CodistParser.gatherElements(varargin);

% Use the local function to check syntax and throw appropriate errors
distributedutil.validateSyntax(@ismember, [{A}, {B}, argList]);
% We don't support tables or timetables yet, so immediately error for these
if istabular(A)
    error(message('parallel:array:UnsupportedFunction', 'ISMEMBER', underlyingType(A)));
end
if istabular(B)
    error(message('parallel:array:UnsupportedFunction', 'ISMEMBER', underlyingType(B)));
end

useRows = cellfun(@(x) iIsPartialMatchFor("rows", x), argList);
if ~isempty(useRows)
    unsupportedFlags = argList(~useRows);
    if contains(unsupportedFlags, "legacy")
        error(message("MATLAB:ISMEMBER:LegacyTrailing"))
    end
end

isRows = any(useRows);
if isRows && (size(A, 2) ~= size(B, 2))
    error(message("MATLAB:ISMEMBER:AandBColnumAgree"))
end
 
% if one array is a cellstr or string, the other must be distributed by rows.
[A, B, Aredist] = iCheckCellStrOrCategorical(A, B);

locTypes = ["none"; "first"; "all"];
locType = locTypes(nargout);
locb = [];
llocb = [];
try
    if nargout <= 1
        lia = ismemberImpl(A, B, isRows, locType);
    else
        [lia, locb, llocb] = ismemberImpl(A, B, isRows, locType);
    end
catch E
   throw(E);
end

if iscodistributed(A)
    codistA = getCodistributor(A);
    szLocalA = size(getLocalPart(A));
else
    szLocalA = size(A);
end

if nargout > 1 && ~isempty(llocb)  % locb is returned as a logical array
    if isRows % llocb is rowsA x rowsB
        lia = any(llocb, 2);
        locb = zeros(size(lia));
        [allLoca, allLocb] = find(llocb);
        allLoca = flipud(allLoca);   % Apply the first matches last -
        allLocb = flipud(allLocb);   % otherwise legacy results
        locb(allLoca) = allLocb;
    else      % llocb is A(:) x B(:)
        locb = zeros(size(llocb, 1), 1);
        [allLoca, allLocb] = find(llocb);
        allLoca = flipud(allLoca);   % Apply the first matches last -
        allLocb = flipud(allLocb);   % otherwise legacy results
        locb(allLoca) = allLocb;
        locb = reshape(locb, szLocalA);
        lia = reshape(lia, szLocalA);
    end
end

if iscodistributed(A) && ~(isRows && codistA.Dimension == 2)
    if isRows || Aredist
        sizeOut = [size(A, 1), 1];
    else
        sizeOut = size(A);
    end
    lia = iCodistOutput(lia, codistA, sizeOut);
    if nargout > 1
        locb = iCodistOutput(locb, codistA, sizeOut);
    end
    return
end
lia = codistributed.pConstructFromReplicated(lia);
if nargout > 1
    locb = codistributed.pConstructFromReplicated(locb);
end
end

function out = iCodistOutput(arr, codistA, sizeOut)
codistr = codistributor1d(codistA.Dimension, codistA.Partition, sizeOut);
out = codistributed.pDoBuildFromLocalPart(arr, codistr);
end

function tf = iIsPartialMatchFor(str, arg)
tf = matlab.internal.datatypes.isScalarText(arg) && length(arg)>1 ...
    && startsWith(str, arg, "IgnoreCase", true);
end

function [A, B, Aredist] = iCheckCellStrOrCategorical(A, B)
% If one input is a cellstr, string or categorical and the other is
% distributed but is not a cellstr, string or categorical, it must be
% distributed by rows.
Aredist = false;
cellstrA = iscellstr(A) || isstring(A) || iscategorical(A);
cellstrB = iscellstr(B) || isstring(B) || iscategorical(B);
isdistA = iscodistributed(A);
isdistB = iscodistributed(B);
if isdistA
    codistA = getCodistributor(A);
    if ~isa(codistA, "codistributor1d")
        A = redistribute(A, codistributor1d(2));
        codistA = getCodistributor(A);
    end
end
if isdistB
    codistB = getCodistributor(B);
    if ~isa(codistB, "codistributor1d")
        B = redistribute(B, codistributor1d(2));
        codistB = getCodistributor(B);
    end
end
if ~xor(cellstrA, cellstrB) || cellstrA && ~isdistB || cellstrB && ~isdistA
    return
end
% Here only if cellstrA and isdistB or cellstrB and isdistA
if (cellstrA && codistB.Dimension == 2) || (cellstrB && codistA.Dimension == 2)
    Aredist = true;
    if cellstrA
        B = redistribute(B, codistributor1d(1));
    else
        A = redistribute(A, codistributor1d(1));
    end
end
end
