function [localX, flag, relres, iter, resvec] = hBicgstablImpl(codistr, A, b, tol, maxit, M1, M2, x)
% HBICGSTABLIMPL Implementation of bicgstabl 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 0 or infinite

%   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);

% ell is fixed to be 2. The user is free to modify this parameter.
ell = 2; 
resvec = zeros(2*ell*maxit+1, 1);       % Preallocate vector for norms of residuals
iter = 0;

% Compute with initial guess
solution = iterApplyA(A, localX, localX, 'notransp', partition);
% Step 2
r0 = localB - solution;
x0 = localX;
% Step 3
normB = distributedutil.vector2NormFromLocalParts(localB);
tolb = tol*normB;                    % Relative tolerance
% Step 4
normRes = distributedutil.vector2NormFromLocalParts(r0);
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
r = r0;                            % Shadow residual
rk = r;
% Step 7
flag = 1;
rho0 = 1;
omega = 1;
% Step 8
alpha = 0;
ukm1 = localX*0;
xk = ukm1;
% Step 9
ut = zeros(size(localX, 1), ell+1, 'like', localX);
rt = ut;
% Step 10
tau = zeros(ell+1, ell+1);
sigma = zeros(ell+1, 1);
gamma = zeros(ell+1, 3);
xmin = localX;
imin = 0;
stag = 0;
moreSteps = 0;
maxMoreSteps = min([floor(n/50), 4*ell, n-maxit]);
% Call iPrecondition the first time such that the condition number is
% checked.
op = 'mldivideWithRcondCheck';
% Step 11
for kk = 1:maxit
    % Step 11.1
    ut(:, 1) = ukm1;
    rt(:, 1) = rk;
    xt = xk;
    % Step 11.2 
    rho0 = -omega*rho0;
    for jj = 1:ell
        % Step 11.3.1
        rho1 = distributedutil.vectorProdFromLocalParts(r0, rt(:, jj));
        % Step 11.3.2
        flag = iterCheckParam(rho0, flag);
        if flag > 1
            resvec = resvec(1:2*ell*(kk-1)+jj);
            break
        end
        % Step 11.3.3
        beta = alpha*rho1/rho0;
        % Step 11.3.4
        rho0 = rho1;
        % Step 11.3.5
        ut(:, 1:jj) = rt(:, 1:jj) - beta*ut(:, 1:jj);
        % Step 11.3.6
        put = iPrecondition(op, ut(:, jj), M1, M2, codistr);
        % Don't check the condition number in all further calls to
        % iPrecondition.
        op = 'mldivide';
        % Step 11.3.7
        ut(:, jj+1) = iterApplyA(A, put, put, 'notransp', partition);
        % Step 11.3.8
        gamma_s = distributedutil.vectorProdFromLocalParts(r0, ut(:, jj+1));
        % Step 11.3.9
        flag = iterCheckParam(gamma_s, flag);
        if flag > 1
            resvec = resvec(1:2*ell*(kk-1)+jj);
            break
        end
        % Step 11.3.10
        alpha = rho0/gamma_s;
        % Step 11.3.11
        rt(:, 1:jj) = rt(:, 1:jj) - alpha*ut(:, 2:jj+1);
        % Step 11.3.12
        prt = iPrecondition('mldivide', rt(:, jj), M1, M2, codistr);
        % Step 11.3.13
        rt(:, jj+1) = iterApplyA(A, prt, prt, 'notransp', partition);
        % Step 11.3.14
        normRes = distributedutil.vector2NormFromLocalParts(rt(:, 1));
        resvec((kk-1)*ell*2+1+jj) = normRes;
        norms = distributedutil.vector2NormFromLocalParts([{ut(:, 1)} {xt}]);
        stag = iterCheckStag(stag, norms, alpha);
        % Step 11.3.15
        xt = xt + alpha*ut(:, 1);          % Half-iterates
        % Step 11.3.16
        [resvec, normResAct, flag, stag, moreSteps, imin, xmin, resMin,...
            rt(:, 1), pxt] = iCheckConv(A, localB, xt, x0, resvec, normRes,...
            (kk-1)*ell*2+jj, ell, tolb, stag, moreSteps, M1, M2, maxMoreSteps,...
            imin, xmin, resMin, rt(:, 1), flag, partition, codistr);
        if flag ~= 1
            iter = ((kk-1)*ell*2+jj)/(ell*2);
            resvec = resvec(1:2*ell*(kk-1)+jj+1);
            break
        end
    end
    if flag ~= 1
        break
    end
    % Step 11.4
    for jj = 2:ell+1
        % Step 11.4.1
        for ii = 2:jj-1
            % Step 11.4.1.1
            tau(ii, jj) = distributedutil.vectorProdFromLocalParts(rt(:, jj),...
                rt(:, ii)/sigma(ii));
            % Step 11.4.1.2
            rt(:, jj) = rt(:, jj) - tau(ii, jj)*rt(:, ii);
        end
        % Step 11.4.2
        sigma(jj) = distributedutil.vectorProdFromLocalParts(rt(:, jj), rt(:, jj));
        % Step 11.4.3
        flag = iterCheckParam(sigma(jj), flag);
        if flag > 1
            resvec = resvec(1:2*ell*(kk-1)+ell+1);
            break
        end
        % Step 11.4.4
        gamma(jj, 2) = distributedutil.vectorProdFromLocalParts(rt(:, 1),...
            rt(:, jj))/sigma(jj);
    end
    if flag == 4
        break
    end
    % Step 11.5
    gamma(ell+1, 1) = gamma(ell+1, 2);
    omega = gamma(ell+1, 1);
    % Step 11.6
    for jj = ell:-1:2
        % Step 11.6.1
        gamma(jj, 1) = gamma(jj, 2) - tau(jj, jj+1:ell+1)*gamma(jj+1:ell+1, 1);
    end
    % Step 11.7
    for jj = 2:ell
        % Step 11.7.1
        gamma(jj, 3) = gamma(jj+1, 1) + tau(jj, jj+1:ell)*gamma(jj+2:ell+1, 1);
    end
    norms = distributedutil.vector2NormFromLocalParts([{rt(:, 1)} {xt}]);
    stag = iterCheckStag(stag, norms, gamma(2, 1));
    % Step 11.8
    xt = xt + gamma(2, 1)*rt(:, 1);
    % Step 11.9
    rt(:, 1) = rt(:, 1) - gamma(ell+1, 2)*rt(:, ell+1);
    % Step 11.10
    ut(:, 1) = ut(:, 1) - gamma(ell+1, 1)*ut(:, ell+1);
    % Step 11.11
    normRes = distributedutil.vector2NormFromLocalParts(rt(:, 1));
    resvec((kk-1)*ell*2+ell+2) = normRes;
    % Step 11.12
    [resvec, normResAct, flag, stag, moreSteps, imin, xmin, resMin,...
        rt(:, 1), pxt] = iCheckConv(A, localB, xt, x0, resvec, normRes,...
        (kk-1)*ell*2+ell+1, ell, tolb, stag, moreSteps, M1, M2, maxMoreSteps,...
        imin, xmin, resMin, rt(:, 1), flag, partition, codistr);
    if flag ~= 1
        iter = ((kk-1)*ell*2+ell+1)/(ell*2);
        resvec = resvec(1:2*ell*(kk-1)+ell+2);
        break
    end
    % Step 11.13
    for jj = 2:ell
        % Step 11.13.1
        ut(:, 1) = ut(:, 1) - gamma(jj, 1)*ut(:, jj);
        norms = distributedutil.vector2NormFromLocalParts([{rt(:, jj)} {xt}]);
        stag = iterCheckStag(stag, norms, gamma(jj, 3));
        % Step 11.13.2
        xt = xt + gamma(jj, 3)*rt(:, jj);
        % Step 11.13.3
        rt(:, 1) = rt(:, 1) - gamma(jj, 2)*rt(:, jj);
        % Step 11.13.4
        normRes = distributedutil.vector2NormFromLocalParts(rt(:, 1));
        resvec((kk-1)*ell*2+ell+1+jj) = normRes;
        % Step 11.13.5
        [resvec, normResAct, flag, stag, moreSteps, imin, xmin, resMin,...
            rt(:, 1), pxt] = iCheckConv(A, localB, xt, x0, resvec, normRes,...
            (kk-1)*ell*2+jj+ell, ell, tolb, stag, moreSteps, M1, M2, maxMoreSteps,...
            imin, xmin, resMin, rt(:, 1), flag, partition, codistr);
        if flag ~= 1
            iter = ((kk-1)*ell*2+ell+jj)/(ell*2);
            resvec = resvec(1:2*ell*(kk-1)+ell+jj+1);
            break
        end
    end
    if flag ~= 1
        break
    end
    % Step 11.14
    ukm1 = ut(:, 1);
    rk = rt(:, 1);
    xk = xt;
end
if isempty(kk)
    kk = 0;
    xt = xmin; % don't have xt from last iteration, set equal xmin
end

% Returned solution is first with minimal residual
% Step 12
if flag == 0 
    relres = normResAct/normB;
    % Step 12.1
    localX = x0 + pxt;
% Step 13
else
    % Step 13.1
    pxt = iPrecondition('mldivide', xmin, M1, M2, codistr);
    % Step 13.2
    pxt2 = iPrecondition('mldivide', xt, M1, M2, codistr);
    % Step 13.3
    resComp = localB - iterApplyA(A, x0+pxt, x0, 'notransp', partition);
    % Step 13.4
    resComp2 = localB - iterApplyA(A, x0+pxt2, x0, 'notransp', partition);
    normComp = distributedutil.vector2NormFromLocalParts([{resComp} {resComp2}]);
    % Step 13.5
    if normComp(1) < normComp(2)
        % Step 13.5.1
        localX = x0 + pxt;
        iter = imin;
        relres = normComp(1)/normB;
    % Step 13.6
    else
        % Step 13.6.1
        localX = x0 + pxt2;
        iter = kk;
        relres = normComp(2)/normB;
    end 
end
end

function [resvec, normResAct, flag, stag, moreSteps, imin, xmin, normResMin, r,...
    pxt] = iCheckConv(A, b, x, x0, resvec, normRes, mm, ell, tolb, stag,...
    moreSteps, M1, M2, maxMoreSteps, imin, xmin, normResMin, r, flag,...
    partition, codistr)
% ICHECKCONV checks a partial-iteration step of the BICGSTABL 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;
pxt = 0;
if normRes <= tolb || stag >= maxStagSteps || moreSteps
    % Double-check residual norm is less than tolerance
    pxt = iPrecondition('mldivide', x, M1, M2, codistr);
    r = b - iterApplyA(A, pxt+x0, x0, 'notransp', partition);
    normResAct = distributedutil.vector2NormFromLocalParts(r);
    resvec(mm+1) = normResAct;
    if (normResAct <= tolb)
        flag = 0;
    else
        if stag >= maxStagSteps && moreSteps == 0
            stag = 0;
        end
        moreSteps = moreSteps + 1;
        if moreSteps >= maxMoreSteps
            flag = 3;
            return
        end
    end
end
if normResAct < normResMin   % update minimal norm quantities
    normResMin = normResAct;
    xmin = x;
    imin = mm/(ell*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
