function [C, ia, ib] = union(A, B, varargin)
%UNION Set union for codistributed array data.
%   C = union(A,B)    
%   [C,IA,IB] = union(A,B)    
%   [C,IA,IB] = union(A,B,'rows')    
%   [C,IA,IB] = union(A,B,'stable') 
%   [C,IA,IB] = union(A,B,'rows','stable') 
%   [C,IA,IB] = union(A,B,'sorted') 
%   [C,IA,IB] = union(A,B,'rows','sorted') 
%   
%   Limitations:
%   The 'legacy' flag is not supported.
%   1. Table, timetable, categorical, datetime and duration inputs are not supported.
%   2. Char and string inputs are not supported when A or B is a cell array of char vectors. Convert cell array of char vectors input arguments to string instead.
%   
%   See also union, codistributed.


%   Copyright 2015-2023 The MathWorks, Inc.

% Allow 5 inputs so the right error can be thrown for the "legacy" flag"
narginchk(2, 5)
nargoutchk(0, 3)

% Default outputs
ia = [];
ib = [];

try
    [flags, A, B] = setfunCheckArgs(@union, A, B, varargin{:});
catch ME
    throwAsCaller(ME)
end

isRows = false;
if ~isempty(flags)
    isRows = ismember("rows", flags);
end

if isRows
    fullC = [A; B];
    Acol = A;
else
    Acol = distributedutil.colonize(A);
    Bcol = distributedutil.colonize(B);
    fullC = [Acol; Bcol];
end

try
    if nargout < 2
        C = unique(fullC, flags{:});
    else
        [C, cia] = unique(fullC, flags{:});
        numA = size(Acol, 1);
        [ia, ib] = iGetIndex(cia, numA);
    end

    if isrow(A) && isrow(B) && iscolumn(C) && ~isRows
        C = C.';
    end

    % Finally restore correct sparsity if needed.
    if issparse(A) || issparse(B)
        C = sparse(C);
    end
    if issparse(ia)
        ia = full(ia);
    end
    if issparse(ib)
        ib = full(ib);
    end
catch E
    throw(E)
end
end

function [ia, ib] = iGetIndex(cia, numA)
% Convert indices into composite fullC into indices into A and B.
if iscodistributed(cia)
    if numA
        localCia = getLocalPart(cia);
        localIa = localCia(localCia<=numA);
        localIb = localCia(localCia>numA) - numA;
        partIx = spmdCat([size(localIa, 1); size(localIb, 1)], 2);
        iaCodist = codistributor1d(1, partIx(1, :), [sum(partIx(1, :)), 1]);
        ibCodist = codistributor1d(1, partIx(2, :), [sum(partIx(2, :)), 1]);
        ia = codistributed.pDoBuildFromLocalPart(localIa, iaCodist);
        ib = codistributed.pDoBuildFromLocalPart(localIb, ibCodist);
    else
        ia = codistributed.zeros(0, 1);
        ib = cia;
    end
else
    ia = cia(cia<=numA);
    ib = cia(cia>numA) - numA;
end
end
