function [LPA, LPB, codistr] = elementWiseOpRedistributeInputs(A, B)
%elementWiseOpRedistributeInputs redistribute A and B such that one can
% call an element-wise binary function on LPA and LPB, obtaining a result
% distributed using codistr codistributor.

%   Copyright 2020 The MathWorks, Inc.

    dimsA = size(A);
    dimsB = size(B);
    args = {A, B};

    % There's a wildcard in operation when there is precisely one scalar-double
    isSDA = iIsScalarDouble(A);
    isSDB = iIsScalarDouble(B);
    isWildcard = sum([isSDA; isSDB]) == 1;
    
    % Only redistribute here if there's no wildcard in operation.
    if isequal(dimsA, dimsB) && ~isWildcard
        [LP, codistr] = codistributed.pRedistSameSizeToSingleDist(args);
        [LPA, LPB] = LP{:};
        return
    end
    ndimsA = ndims(A);
    ndimsB = ndims(B);
    sz = ones(2, max(ndimsA, ndimsB));
    sz(1, 1:ndimsA) = dimsA(1:ndimsA);
    sz(2, 1:ndimsB) = dimsB(1:ndimsB);
    outputSz = max(sz);
    
    % modify outputSz if one of the inputs was empty
    outputSz(min(sz) == 0) = 0;
    
    % determine output codistributor
    codistr = iGetOutputCodistributor(args, outputSz);
    
    % get local parts
    LPA = iGetLocalPart(codistr, A, isWildcard && isSDA);
    LPB = iGetLocalPart(codistr, B, isWildcard && isSDB);
end

function tf = iIsScalarDouble(x)
clz = underlyingType(x);
tf = isscalar(x) && isequal(clz, 'double');
end

function codistr = iGetOutputCodistributor(args, outputSz)
    isCodistr = cellfun(@iscodistributed, args, 'UniformOutput', true);
    codistrs = cellfun(@getCodistributor, args(isCodistr), 'UniformOutput', ...
        false);
    validCodistrs = cellfun(@(x)(x.hSupportsDimensionality(length(outputSz))), ...
        codistrs, 'UniformOutput', true);
    switch nnz(validCodistrs)
        case 0
            [maxSz, maxDim] = max(outputSz);
            codistr = codistributor1d(maxDim, ...
                codistributor1d.defaultPartition(maxSz), outputSz);
        case 1
            codistr = codistrs{validCodistrs};
            codistr = codistr.hGetNewForSize(outputSz);  %resize used codistributor
        case 2
            if numel(args{1})>=numel(args{2})
                codistr = codistrs{1};
            else
                codistr = codistrs{2};
            end
            codistr = codistr.hGetNewForSize(outputSz);  %resize used codistributor        
    end
end

function LP = iGetLocalPart(codistr, A, isThisArgWildcard)
    szA = size(A);
    
    % We must always gather args if they are the sole wildcard argument, as per
    % g1317175.
    if isThisArgWildcard || codistr.hDoGatherBsxfunInput(szA)
        replA = gather(A);
        LP = codistr.hGetBsxfunLocalPartFromReplicated(replA);
        return
    end
   
    % Here if gathering is unnecessary - redistribute
    destCodistr = codistr.hGetCompleteForSize(szA);
    if iscodistributed(A)
        codistrA = getCodistributor(A);
        LPA = getLocalPart(A);
        LP = distributedutil.Redistributor.redistribute(codistrA, LPA, ...
            destCodistr);
    else
        LP = destCodistr.hBuildFromReplicatedImpl(0, A);
    end
end
