function [lia, locb, llocb] = ismemberImpl(A, B, isRows, locType)
%ISMEMBERIMPL does the core calculation for the distributed ismember
%function.
%   Inputs: A - Potential member values
%           B - Reference set of values
%           useRows - Make comparisons by row
%           locType - What type of location data to return
%               "none" - No location data requested
%               "first" - Return index of first occurrence of A in B
%               "all" - Return logical array: all occurrences of A in B
%   Outputs: lia - local part of lia
%            locb - part of locb corresponding to localA with offsets
%                   applied if locb is numeric
%            llocb - logical array of size numel(localA) x numel(B)
%                    indicating where A and B match

%   Copyright 2022-2023 The MathWorks, Inc.

locb = [];
llocb = [];
if isempty(A)
    lia = false(0);
    return
end

% if one array is a cellstr or string, the other must be distributed by rows.
[A, B] = iCheckCellStrOrCategorical(A, B);

[sA, locType] = iMakeArrayStruct(A, isRows, locType);
[sB, locType] = iMakeArrayStruct(B, isRows, locType, sA);
if ischar(sA.local) && isrow(sA.local) && sA.size(1) > 1 
    sA.local = char(deblank(string(sA.local)));
elseif ischar(sB.local) && isrow(sB.local) && sB.size(1) > 1
    sB.local = char(deblank(string(sB.local)));
end

switch locType
    case "none"
        lia = iDoIsmember(sA, sB, isRows, false);
    case "first"
        [lia, locb] = iDoIsmember(sA, sB, isRows, true);
    case "all"
        [lia, llocb] = iDoCompleteIsmember(sA, sB, isRows);
    otherwise
        assert(false, "Unrecognized locType option")
end
end

function [sX, locType] = iMakeArrayStruct(X, isRows, locType, sA)
% Calculate properties of a (possibly distributed) array and return them.

sX = struct('codist', [], 'offset', [], 'local', [], 'size', [], 'distdim', []);

if iscodistributed(X)
    sX.codist = getCodistributor(X);
    sX.distdim = sX.codist.Dimension;
    if isRows && sX.distdim == 2
        locType = "all";
        if nargin > 3 && sA.distdim == 2 && ...
                ~all(sX.codist.Partition == sA.codist.Partition)  % Partitions must match
            sX.codist = codistributor1d(sX.distdim, sA.codist.Partition, size(X));
            X = redistribute(X, sX.codist);
        end
    end
    offset = zeros(ndims(X), spmdSize+1);
    offset(sX.distdim, :) = [0 cumsum(sX.codist.Partition)];
    sX.offset = offset;
    sX.local = getLocalPart(X);
else
    sX.offset = zeros(ndims(X), spmdSize + 1);
    sX.local = X;
    sX.distdim = 0;
end
sX.size = size(X);
end

function [lia, locb] = iDoIsmember(sA, sB, isRows, getLocb)
% Calls local ismember on A and B
if isRows
    flag = {'rows'};
else
    flag = {};
end
isBcodist = sB.distdim>0;

[lia, locb] = iLocalIsMember(sA.local, sB.local, flag, getLocb);
% Adjust locb if necessary
if getLocb && isBcodist
    locb = iConvertLoc(sB, spmdIndex, lia, locb, isRows);
end
if isBcodist  % Rotate parts of B through workers and evaluate rest of lia and locb.
    Blocal = sB.local;
    for ii = 1:spmdSize-1
        Blocal = iRotateB(Blocal);
        [newLia, newLocb] = iLocalIsMember(sA.local, Blocal, flag, getLocb);
        lia = lia | newLia;
        if getLocb
            newSpmdIndex = mod(spmdIndex-ii-1, spmdSize) + 1;
            newLocb = iConvertLoc(sB, newSpmdIndex, newLia, newLocb, isRows);
        end
        % Replace nonzero locb with newLocb if newLocb is smaller.
        % Replace zero locb with newLocb if newLocb is nonzero.
        locb(locb == 0 & newLocb > 0) = newLocb(locb == 0 & newLocb > 0);
        locb((locb>newLocb)&(newLocb>0)) = newLocb((locb>newLocb)&(newLocb>0));
    end
end
end

function [lia, locb] = iLocalIsMember(A, B, flag, getLocb)
if ~getLocb
    locb = [];
    lia = ismember(A, B, flag{:});
else
    [lia, locb] = ismember(A, B, flag{:});
end
end

function locb = iConvertLoc(sB, indexB, lia, locb, isRows)
% Convert linear indices in part of a distributed array to the entire array.
if isRows
    locb(lia) = locb(lia) + sB.offset(1, indexB);
else
    codistDim = sB.codist.Dimension;
    szB = sB.size;
    szPartOfB = sB.size;
    szPartOfB(codistDim) = sB.codist.Partition(indexB);
    rowcolB = cell(1,length(szB));
    [rowcolB{:}] = ind2sub(szPartOfB, locb(lia));
    rowcolB{codistDim} = rowcolB{codistDim} + sB.offset(codistDim, indexB);
    locb(lia) = sub2ind(szB, rowcolB{:});
end
end

function [lia, llocb] = iDoCompleteIsmember(sA, sB, isRows)
% Ands elements of A and B to get all occurrences of A in B.

if isRows
    llocb = false(size(sA.local, 1), sB.size(1));
    llocb = iDoLocalRowIsmember(llocb, sA, sB, sB.local, spmdIndex);
    if sB.distdim % Rotate parts of B through workers and update llocb.
        Blocal = sB.local;
        for ii = 1:spmdSize-1
            newLocb = false(size(sA.local, 1), sB.size(1));
            Blocal = iRotateB(Blocal);
            remoteIdx = mod(spmdIndex-ii-1, spmdSize) + 1;
            newLocb = iDoLocalRowIsmember(newLocb, sA, sB, Blocal, remoteIdx);
            if sB.distdim == 1
                llocb = llocb | newLocb;
            else
                llocb = llocb & newLocb;  % corresponding subrows must match
            end
        end
    end  % llocb has all rows evaluated in local A.
    if sA.distdim == 2
        % Combine results from all workers.  Not a match unless all workers
        % match.
        llocb = spmdReduce(@and, llocb);
    end
    llocb = reshape(llocb, size(sA.local, 1), sB.size(1));
else
    llocb = false(numel(sA.local), max(1, prod(sB.size)));
    llocb = iDoLocalIsmember(llocb, sA, sB, sB.local, spmdIndex);
    if sB.distdim % Rotate parts of B and fill in llocb.
        newLocb = false(numel(sA.local), max(1, prod(sB.size)));
        Blocal = sB.local;
        for ii = 1:spmdSize-1
            Blocal = iRotateB(Blocal);
            remoteIdx = mod(spmdIndex-ii-1, spmdSize) + 1;
            newLocb = iDoLocalIsmember(newLocb, sA, sB, Blocal, remoteIdx);
            llocb = llocb | newLocb;
        end
    end  % llocb has matches for all elements in local A with all of B.
end
lia = any(llocb, 2);
if ~isRows
    lia = reshape(lia, size(sA.local));
end
end

function llocb = iDoLocalRowIsmember(llocb, sA, sB, Blocal, remoteIdx)
% Calculate matrix of matches between overlapping parts of each local row,
% placing the results in llocb, whose size is rowsA x rowsB.

% calculate column limits
globalStart = max(sA.offset(2, spmdIndex), sB.offset(2, remoteIdx))+1;
globalEnd = min(sA.offset(2, spmdIndex) + size(sA.local, 2), sB.offset(2, remoteIdx) + size(Blocal, 2));
Astart = globalStart - sA.offset(2, spmdIndex); 
Aend = globalEnd - sA.offset(2, spmdIndex);
Acols = Astart:Aend;
Bstart = globalStart - sB.offset(2, remoteIdx);
Bend = globalEnd - sB.offset(2, remoteIdx);
Bcols = Bstart:Bend;
Boffset = sB.offset(1, remoteIdx);
Brows = (1:size(Blocal, 1)) + Boffset;
for rowA = 1:size(sA.local, 1)
     llocb(rowA, Brows) = (all(sA.local(rowA, Acols) == Blocal(:, Bcols), 2))';
end
end

function llocb = iDoLocalIsmember(llocb, sA, sB, Blocal, remoteIdx)
% Calculate matrix of matches between A.local and the local part of B from
% this worker or another one.

if iscellstr(sA.local)
    sA.local = string(sA.local);
end
if iscellstr(Blocal) %#ok<ISCLSTR>
    Blocal = string(Blocal);
end

% Flatten multidimensional arrays
Alocal = reshape(sA.local, size(sA.local, 1), []);
Blocal = reshape(Blocal, size(Blocal, 1), []);
Boffset = sB.offset(:, remoteIdx)*prod(sB.size(2:(sB.distdim-1)));
[nrowsB, ncolsB] = size(Blocal);

ixB = reshape(1:prod(sB.size), sB.size);
ixLocB = ixB((1:nrowsB) + Boffset(1), (1:ncolsB) + Boffset(2));
llocb(:, ixLocB(:)) = Alocal(:) == (Blocal(:)).';
end

function localB = iRotateB(localB)
to = mod(spmdIndex, spmdSize) + 1;    % worker to right with wrap
from = mod(spmdIndex-2, spmdSize) + 1; % worker to left with wrap
localB = spmdSendReceive(to, from, localB);
end

function [A, B] = 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.
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)
    if cellstrA
        B = redistribute(B, codistributor1d(1));
    else
        A = redistribute(A, codistributor1d(1));
    end
end
end