function [localA, b, M1, M2, x0] = iterSetupArgs(A, b, M1, M2, x0, restart)
%ITERSETUPARGS Redistribute data inputs so that b and x0 are distributed by
% rows, M1 and M2 are distributed by rows in the same way as x0 if they are
% matrices, and A is distributed by columns in the same way as b if it is a
% matrix.  Also provide distributed transposes of the matrix inputs and
% return local parts of all except b unless b and x0 are different sizes,
% in which case both b and x0 are returned distributed.
% If we're using the GMRES solver, the restart argument will be given and
% we need to assure that the partition used for distribution has at
% least restart elements on the first worker.
    
%   Copyright 2015-2022 The MathWorks, Inc.

% Check preconditioner functions for validity before redistributing A
% and b because deadlock can occur if a test fails on one worker while
% another is still involved in redistribution - g1229052
n = size(x0, 1);
m = size(b, 1);
isLsqr = (m ~= n);
localA = A;   % Parts of A other than A.fun aren't codistributed.
% If M1 is a function, make sure it works
if ~strcmp(M1.type, 'matrix')
    distributedutil.iterapp('mldivide', M1.fun, M1.type, ...
        zeros(size(x0)), 'notransp');
    M1.tfun = M1.fun;
end
% If M2 is a function, make sure it works
if ~strcmp(M2.type, 'matrix')
    distributedutil.iterapp('mldivide', M2.fun, M2.type, ...
        zeros(size(x0)), 'notransp');
    M2.tfun = M2.fun;
end

% redistribute A and b if necessary
inCodistrB = getCodistributor(b);
if isa(inCodistrB, 'codistributor1d') && inCodistrB.Dimension == 1
    newPartB = inCodistrB.Partition;
else
    newPartB = codistributor1d.defaultPartition(length(b));
end
inCodistrX = getCodistributor(x0);
if isLsqr
    if isa(inCodistrX, 'codistributor1d') && inCodistrX.Dimension == 1
        newPartX = inCodistrX.Partition;
    else
        newPartX = codistributor1d.defaultPartition(length(x0));
    end
else
    newPartX = newPartB;
end
if strcmp(A.type, 'matrix')
    inCodistrA = getCodistributor(A.fun);
    if isa(inCodistrA, 'codistributor1d') && inCodistrA.Dimension == 1
        newPartB = inCodistrA.Partition;
    end
    if isa(inCodistrA, 'codistributor1d') && inCodistrA.Dimension == 2
        newPartX = inCodistrA.Partition;
    end
    if nargin > 5  && newPartB(1)<restart+1  % GMRES
        newPartB = iGmresPartition(size(A.fun, 1), restart);
    end
    if ~(isa(inCodistrA, 'codistributor1d') && inCodistrA.Dimension == 1 ...
            && ~any(inCodistrA.Partition - newPartB))
        workCodistrA = codistributor1d(1, newPartB, size(A.fun));
        A.fun = redistribute(A.fun, workCodistrA);
    end
    if ~isLsqr
        newPartX = newPartB;
    end
    stack = dbstack;
    caller = stack(2).name;
    if any(strcmpi(caller, {'bicg', 'lsqr', 'qmr'}))  % Need A.tfun
        localA.tfun = getLocalPart(A.fun');
    end
    localA.fun = getLocalPart(A.fun);
elseif nargin > 5 && newPartB(1)<restart+1
    newPartB = iGmresPartition(length(b), restart);
    newPartX = newPartB;
else
    localA.tfun = localA.fun;
end
if ~(isa(inCodistrB, 'codistributor1d') && inCodistrB.Dimension == 1 ...
        && ~any(inCodistrB.Partition - newPartB))
    workCodistrB = codistributor1d(1, newPartB, size(b));
    b = redistribute(b, workCodistrB);
end
if ~(isa(inCodistrX, 'codistributor1d') && inCodistrX.Dimension == 1 ...
        && ~any(inCodistrX.Partition - newPartX))
    workCodistrX = codistributor1d(1, newPartX, size(x0));
    x0 = redistribute(x0, workCodistrX);
end

if ~iIsEmptyArg(M1) || ~iIsEmptyArg(M2)
    workCodistrM = codistributor1d(1, newPartX, [n n]);
    stack = dbstack;
    caller = stack(2).name;
    needTranspose = any(strcmpi(caller, {'bicg', 'lsqr', 'qmr'}));          
    if strcmp(M1.type, 'matrix') && ~isempty(M1.fun)
        try
            M1 = iSetPreconditionMatrix(M1, workCodistrM, needTranspose);
        catch ME
            throwAsCaller(ME);
        end
    else
        % Function preconditioner should be replicated
        M1 = gather(M1);
    end
    if strcmp(M2.type, 'matrix') && ~isempty(M2.fun)
        try
            M2 = iSetPreconditionMatrix(M2, workCodistrM, needTranspose);
        catch ME
            throwAsCaller(ME);
        end
    else
        % Function preconditioner should be replicated
        M2 = gather(M2);
    end
end
end

function M = iSetPreconditionMatrix(M, codistr, needTranspose)
% Distribute M by rows and put the matrix in M.fun and its transpose in
% M.tfun
    M.fun = codistributed(M.fun, codistr);
    if needTranspose
        M.tfun = redistribute(M.fun', codistr);
    end
end

function newPart = iGmresPartition(lenX, restart)
% Create a new partition that has restart+1 elements on the first worker iff
% the pool has more than 1 worker.
if spmdSize>1
    firstPart = min(lenX, restart+1);
    remaining = lenX - firstPart;
    newPart = zeros(1, spmdSize);
    newPart(1) = firstPart;
    newPart(2:end) = floor(double(remaining)/(spmdSize-1));
    newPart(2:(rem(remaining,(spmdSize-1))+1)) = ceil(double(remaining)/(spmdSize-1));
else
    % In case of a single worker, newPart holds all of X
    newPart = lenX;
end
end

function tf = iIsEmptyArg(M)
% Determines whether M is an empty matrix
tf = strcmpi('matrix', M.type) && isempty(M.fun);
end
