function [A, origDistA, isRedistributed] = redistributeIfNeeded(A, newDistA)
% REDISTRIBUTEIFNEEDED
% Function that checks if A is distributed using newDistA.
% Otherwise, A is redistributed and the original distribution scheme,
% as well as a flag that indicates the redistribution, is returned. If A is
% not distributed on entry, it is treated as if it was distributed using
% the default codistributor1d.

%   Copyright 2016-2019 The MathWorks, Inc.

isRedistributed = false;

% Check if A is distributed
if ~iscodistributed(A)
    A = codistributed(A, newDistA);
    
    % Solution should use default codistributor1d. If newDistA is
    % codistributor2dbc or a codistributor1d not using the default
    % distribution dimension, we will use default codistributor for
    % solution. Otherwise, we keep solution distributed via newDistA. 
    defaultDistDim = distributedutil.Sizes.lastNonSingletonDimension(...
        newDistA.Cached.GlobalSize);
    if isa(newDistA, 'codistributor2dbc') || ...
            (isa(newDistA, 'codistributor1d') && ...
            newDistA.Dimension ~= defaultDistDim)
        origDistA = codistributor1d(defaultDistDim, ...
            codistributor1d.unsetPartition, newDistA.Cached.GlobalSize);
        isRedistributed = true;
    else
        origDistA = newDistA;
    end
else
    % Save original distribution scheme for later use
    origDistA = getCodistributor(A);
    
    % Compare codistributors (all properties but the Cached.GlobalSize)
    classesMatch = strcmp(class(origDistA), class(newDistA));
    
    codistAre1dAndMatch = classesMatch && ...
        isa(origDistA, 'codistributor1d') && ...
        origDistA.Dimension==newDistA.Dimension && ...
        all(origDistA.Partition==newDistA.Partition);
    
    codistAre2dbcAndMatch = classesMatch && ...
        isa(origDistA,'codistributor2dbc') && ...
        all(origDistA.LabGrid==newDistA.LabGrid) && ...
        origDistA.BlockSize==newDistA.BlockSize && ...
        strcmp(origDistA.Orientation, newDistA.Orientation);
    
    % Redistribute if new distribution scheme is not the needed one
    if ~codistAre1dAndMatch && ~codistAre2dbcAndMatch
        A = redistribute(A, newDistA);
        isRedistributed = true;
    end
end
end
