function [localX, flag, relres, iter, resvec] = hCgsImpl(codistr, A, b, tol, maxit, M1, M2, x)
% HCGSIMPL Implementation of cgs function with minimum communication using
% algorithm described in Codistributed_sparse_iterative_solvers_RFA on the wiki.
% b is codistributed 1d by rows; A.fun is the local part of A distributed 1d
% by rows, using the same partition scheme as b; M1 and M2, if they are
% matrices, are distributed by rows.  The "step" comments refer to lines in
% the algorithm on the wiki.

% The flag variable indicates the progress of the solution:
% 0  Converged
% 1  Not converged yet
% 2  There was an infinite value in the output from the preconditioner
% 3  Iteration has stagnated
% 4  A scalar parameter has become too small or too large

%   Copyright 2015-2022 The MathWorks, Inc.

% Message-passing info
[partitionStart, partitionEnd] = globalIndices(codistr, 1);
partition = matlab.internal.ColonDescriptor(partitionStart, partitionEnd);
n = codistr.Cached.GlobalSize(1);
localB = getLocalPart(b);
localX = getLocalPart(x);

normB = distributedutil.vector2NormFromLocalParts(localB);
tolb = tol*normB;
flag = 1;

resvec = zeros(maxit+1,1);         % Preallocate vector for norms of residuals
xmin = localX;                          % Iterate which has minimal residual so far
imin = 0;                          % Iteration at which xmin was computed

% Compute with initial guess
% Step 2
r = localB - iterApplyA(A, localX, localX, 'notransp', partition);
r0 = r;                                     % Shadow residual
% Check if we have a solution
normRes = distributedutil.vector2NormFromLocalParts(r);
resvec(1) = normRes;
if normRes <= tolb
    resvec = resvec(1);
    flag = 0;
    relres = normRes/normB;
    iter = 0;
    return
end
% Step 3
rho = normRes^2;
u = r0;
% Step 4
p = r0;
normResMin = normRes;
normResAct = normRes;
stag = 0;
moreSteps = 0;
maxMoreSteps = min([floor(n/50), 5, n-maxit]);
% Call iPrecondition the first time such that the condition number is
% checked.
op = 'mldivideWithRcondCheck';
% Step 5
for ii = 1:maxit
    % Step 5.1
    ph = iPrecondition(op, p, M1, M2, codistr);
    % Don't check the condition number in all further calls to
    % iPrecondition.
    op = 'mldivide';
    % Step 5.2
    [vh, rtvh] = iterApplyA(A, ph, r0, 'notransp', partition);
    flag = iterCheckParam(rtvh, flag, 1);
    if flag>1
        break
    end
    alpha = rho/rtvh;
    flag = iterCheckParam(alpha, flag, 2);
    if flag>1
        break
    end
    % Step 5.3
    q = u - alpha*vh;
    % Step 5.4
    uh = iPrecondition('mldivide', u+q, M1, M2, codistr);
    % Step 5.5
    norms = distributedutil.vector2NormFromLocalParts([{uh} {localX}]);
    stag = iterCheckStag(stag, norms, alpha);
    localX = localX + alpha*uh;
    % Step 5.6
    qh = iterApplyA(A, uh, r0, 'notransp', partition);
    r = r - alpha*qh;
    % Step 5.7
    rtr = distributedutil.vectorProdFromLocalParts(r0, r);
    normRes = distributedutil.vector2NormFromLocalParts(r);
    resvec(ii+1) = normRes;
    [resvec, flag, stag, moreSteps, imin, xmin, normResMin, normResAct] =...
        iterCheckConv(A, localB, localX, resvec, normRes, ii, tolb, stag, moreSteps, ...
        maxMoreSteps, imin, xmin, normResMin, flag, partition);
    if flag ~= 1
        break
    end
    % Step 5.8
    beta = rtr/rho;
    flag = iterCheckParam(beta, flag);
    rho = rtr;
    flag = iterCheckParam(rho, flag);
    if flag>1
        break
    end
    % Step 5.9
    u = r + beta*q;
    % Step 5.10
    p = u + beta*(q + beta*p);
end
if isempty(ii)
    ii = 0;
end

% Returned solution is first with minimal residual
if flag == 0
    iter = imin;
    relres = normResAct/normB;
else
    r = localB - iterApplyA(A, xmin, xmin, 'notransp', partition);
    normRes = distributedutil.vector2NormFromLocalParts(r);
    if normRes <= normResAct
        localX = xmin;
        iter = imin;
        relres = normRes/normB;
    else
        iter = ii;
        relres = normResAct/normB;
    end
end    
% Truncate the zeros from resvec
if flag <= 1 || flag == 3
    resvec = resvec(1:ii+1);
else
    resvec = resvec(1:ii);
end
end

function localZ = iPrecondition(op, localR, M1, M2, codistr)
% IPRECONDITION applies the preconditioners (if supplied) to r.  It
% uses local parts for input and output.
hasAPreconditioner = ~iterIsEmptyArg(M1) || ~iterIsEmptyArg(M2);
if hasAPreconditioner
    r = codistributed.build(localR, codistr, 'noCommunication');
    if ~iterIsEmptyArg(M1)
        z = distributedutil.iterapp(op, M1.fun, M1.type, ...
                           r, 'notransp');
    else
        z = r;
    end
    if ~iterIsEmptyArg(M2)
        z = distributedutil.iterapp(op, M2.fun, M2.type, ...
                           z, 'notransp');
    end
    localZ = getLocalPart(z);
else
    localZ = localR;
end
end
