function [localX, flag, relres, iter, resvec] = hBicgstabImpl(codistr, A, b, tol, maxit, M1, M2, x)
% HBICGSTABIMPL Implementation of bicgstab 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 2017-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);

resvec = zeros(2*maxit+1, 1);      % Preallocate vector for norms of residuals
iter = 0;

% Compute with initial guess
% Step 2
r = localB - iterApplyA(A, localX, localX, 'notransp', partition);
% Step 3
norms = distributedutil.vector2NormFromLocalParts([{localB} {r}]);
normB = norms(1);
tolb = tol*normB;                  % Relative tolerance
% Step 4
normRes = norms(2);
resvec(1) = normRes;               % resvec(1) = norm(b-A*x0)
resMin = normRes;
% Check if we have a solution
% Step 5
if resMin <= tolb                % Initial guess is a good enough solution
    flag = 0;
    relres = resMin/tolb;
    resvec = resvec(1);
    return
end
% Step 6
rt = r;                            % Shadow residual
p = r;
% Step 7
flag = 1;
xmin = localX;
imin = 0;
stag = 0;
moreSteps = 0;
maxMoreSteps = min([floor(n/50), 5 ,n-maxit]);
rho = normRes^2;
% Call iPrecondition the first time such that the condition number is
% checked.
op = 'mldivideWithRcondCheck';
% Step 8
for ii = 1:maxit
    % Step 8.1
    ph = iPrecondition(op, p, M1, M2, codistr);
    % Don't check the condition number in all further calls to
    % iPrecondition.
    op = 'mldivide';
    % Step 8.2 
    v = iterApplyA(A, ph, ph, 'notransp', partition);
    % Step 8.3
    rtv = distributedutil.vectorProdFromLocalParts(rt, v);
    % Step 8.4
    flag = iterCheckParam(rtv, flag);
    if flag > 1
        resvec = resvec(1:2*ii-1);
        break
    end
    % Step 8.5
    alpha = rho / rtv;
    flag = iterCheckParam(alpha, flag, 2);
    if flag > 1
        resvec = resvec(1:2*ii-1);
        break
    end
    norms = distributedutil.vector2NormFromLocalParts([{ph} {localX}]);
    stag = iterCheckStag(stag, norms, alpha);
    % Step 8.6
    xh = localX + alpha*ph;      % form the "half" iterate
    % Step 8.7
    s = r - alpha*v;        % residual associated with xh
    % Step 8.8
    normRes = distributedutil.vector2NormFromLocalParts(s);
    resvec(2*ii) = normRes;
    % Step 8.9
    [resvec, normResAct, flag, stag, moreSteps, imin, xmin, resMin] =...
        iCheckConv(A, localB, xh, resvec, normRes, 2*ii-1, tolb, stag, moreSteps, ...
        maxMoreSteps, imin, xmin, resMin, flag, partition);
    if flag ~= 1
        iter = ii - 0.5;
        break
    end
    % Step 8.10
    sh = iPrecondition('mldivide', s, M1, M2, codistr);
    % Step 8.11
    t = iterApplyA(A, sh, sh, 'notransp', partition);
    % Step 8.12
    tt = distributedutil.vectorProdFromLocalParts(t, t);
    % Step 8.13
    flag = iterCheckParam(tt, flag);
    if flag > 1
        resvec = resvec(1:2*ii);
        break
    end
    % Step 8.14
    omega = distributedutil.vectorProdFromLocalParts(t, s)/tt;
    % Step 8.15
    flag = iterCheckParam(omega, flag, 2);
    if flag > 1
        resvec = resvec(1:2*ii);
        break
    end
    norms = distributedutil.vector2NormFromLocalParts([{sh} {xh}]);
    stag = iterCheckStag(stag, norms, omega);
    % Step 8.16
    localX = xh + omega*sh;
    % Step 8.17
    r = s - omega*t;
    % Step 8.18
    normRes = distributedutil.vector2NormFromLocalParts(r);
    resvec(2*ii+1) = normRes;
    % Step 8.19
    [resvec, normResAct, flag, stag, moreSteps, imin, xmin, resMin] =...
        iCheckConv(A, localB, localX, resvec, normRes, 2*ii, tolb, stag, moreSteps, ...
        maxMoreSteps, imin, xmin, resMin, flag, partition);
    if flag ~= 1
        iter = ii;
        break
    end
    % Step 8.20
    rho1 = rho;
    % Step 8.21
    rho = distributedutil.vectorProdFromLocalParts(rt, r);
    % Step 8.22
    flag = iterCheckParam(rho, flag);
    if flag > 1
        resvec = resvec(1:2*ii+1);
        break
    end
    % Step 8.23
    beta = (rho/rho1)*(alpha/omega);
    % Step 8.24
    flag = iterCheckParam(beta, flag);
    if flag > 1
        resvec = resvec(1:2*ii+1);
        break
    end
    % Step 8.25
    p = r + beta*(p - omega*v);
end
if isempty(ii)
    ii = 0;
    normResAct = normRes;
end

% Returned solution is first with minimal residual
if flag == 0
    relres = normResAct/normB;
    localX = xmin;
else
    % Step 8.9
    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
end

function [resvec, normResAct, flag, stag, moreSteps, imin, xmin, normResMin] =...
    iCheckConv(A, b, x, resvec, normRes, mm, tolb, stag, moreSteps, ...
    maxMoreSteps, imin, xmin, normResMin, flag, partition)
% ICHECKCONV checks a half-iteration step of the BICGSTAB algorithm for 
% convergence.  If the iteration has not converged, it also checks for
% stagnation and updates the best result so far.

maxStagSteps = 3;
normResAct = normRes;
if normRes <= tolb || stag >= maxStagSteps || moreSteps
    r = b - iterApplyA(A, x, x, 'notransp', partition);
    normResAct = distributedutil.vector2NormFromLocalParts(r);
    resvec(mm+1) = normResAct;
    if (normResAct <= tolb)
        flag = 0;
        resvec = resvec(1:mm+1);
    else
        if stag >= maxStagSteps && moreSteps == 0
            stag = 0;
        end
        moreSteps = moreSteps + 1;
        if moreSteps >= maxMoreSteps
            flag = 3;
            resvec = resvec(1:mm+1);
        end
    end
end
if normResAct < normResMin
    normResMin = normResAct;
    xmin = x;
    imin = mm/2;
end
if stag >= maxStagSteps
    flag = 3;
    return
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
