function outCodistr = pagefunGetOutputsCodistr(vars, numOutPages)
% Helper to determine the codistributor of the outputs. If all inputs are
% already distributed 1D in the same dimension and it's dimension 3 or
% higher then we can use that. Otherwise choose a new dimension.

%   Copyright 2020 The MathWorks, Inc.

inCodistrs = iGet1DDists(vars);
inDistDims = unique(cellfun(@(x) x.Dimension, inCodistrs));
outPart = [];
if isscalar(inDistDims) && inDistDims > 2
    % All 1D distributed inputs have the same distribution dimension so use
    % it.
    outDistDim = inDistDims;
    % If all distributed inputs also have the same partitioning use that too.
    part1 = inCodistrs{1}.Partition;
    if numel(inCodistrs) == 1 ...
            || all(cellfun(@(x) isequal(x.Partition, part1), inCodistrs))
        outPart = part1;
    end
else
    % Inputs have mismatched or too low distribution dimensions. Choose a
    % new one.
    outDistDim = iChooseDistDim(numOutPages, inDistDims);
end
if isempty(outPart)
    outPart = codistributor1d.defaultPartition(numOutPages(outDistDim-2));
end
outCodistr = codistributor1d(outDistDim, outPart);
end

% Find all 1D distributed inputs and return their codistributors
function codistrs = iGet1DDists(args)
codistrs = {};
for ii=1:numel(args)
    if iscodistributed(args{ii})
        codistr = getCodistributor(args{ii});
        if isa(codistr, "codistributor1d")
             codistrs{1,end+1} = codistr; %#ok<AGROW>
        end
    end
end
end

% Choose how the output should be distributed
function distDim = iChooseDistDim(outPgSzs, inDistDims) %#ok<INUSD>
% For now just select the last dimension.
distDim = numel(outPgSzs)+2;
end
