function [localX, flag, relres, iter, resvec] = hTfqmrImpl(codistr, A, b, tol, maxit, M1, M2, x)
% HTFQMRIMPL Implementation of tfqmr 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

%   Copyright 2016-2022 The MathWorks, Inc.

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

resvec = zeros(maxit+1, 1);         % Preallocate vector for norms of residuals
% Step 2
% Compute with initial guess
r = localB - iterApplyA(A, localX, localX, 'notransp', partition);
r0 = r;
% Step 3
buffer = distributedutil.vector2NormFromLocalParts([{r} {localB}]);
normB = buffer(2);
tolb = tol*normB;
% Step 4
normRes = buffer(1);
resvec(1) = normRes;                 % resvec(1) = norm(b-A*x0)
iter = 0;
relres = normRes/normB;
% Step 5
if (normRes <=  tolb)                % Initial guess is a good enough solution
    flag = 0;
    resvec = normRes;
    return
end
% Step 6
xmin = localX;                          % Iterate which has minimal residual so far
imin = 0;                          % Iteration at which xmin was computed
normResMin = normRes;              % Norm of residual from xmin
u_m = r;
w = r;
% Step 7
pu_m = iPrecondition('mldivideWithRcondCheck', r, M1, M2, codistr);
res = distributedutil.vectorProdFromLocalParts(pu_m, pu_m);
if ~isfinite(res)
    flag = 2;
    return
end
% Step 8
Au = iterApplyA(A, pu_m, pu_m, 'notransp', partition);
v = Au;
% Step 9
d = zeros(size(r), 'like', r);
Ad = d;
% Step 10
tau = normRes;
theta = 0;
eta = 0;
% Step 11
rho = tau*tau;
% Step 12
stag = 0;
moreSteps = 0;
maxMoreSteps = min([floor(n/50), 5, n-maxit]);
even = false;
for mm = 1:maxit*2
    % Step 12.1
    if ~even
        % Step 12.1.1
        alpha = rho/distributedutil.vectorProdFromLocalParts(r0, v);
        % Step 12.1.2
        u_mp1 = u_m - alpha*v;
    end
    % Step 12.2
    w = w - alpha*Au;
    % Step 12.3
    sigma = (theta^2/alpha)*eta;
    % Step 12.4
    d = pu_m + sigma*d;
    % Step 12.5
    Ad = Au + sigma*Ad;
    % Step 12.6
    norms = distributedutil.vector2NormFromLocalParts([{d} {localX} {w}]); %[norm(d) norm(x) norm(w)]
    theta = norms(3)/tau;
    % Step 12.7
    c_mp1 = 1/sqrt(1 + theta^2);
    % Step 12.8
    tau = tau*theta*c_mp1;
    % Step 12.9
    eta = c_mp1^2*alpha;
    % Step 12.10
    stag = iterCheckStag(stag, norms, eta);
    % Step 12.11
    localX = localX + eta*d;
    % Step 12.12
    r = r - eta*Ad;
    % Step 12.13
    normRes = distributedutil.vector2NormFromLocalParts(r);
    resvec(mm+1) = normRes;
    % Step 12.14
    [resvec, flag, stag, moreSteps, imin, xmin, normResMin, normResAct] =...
         iterCheckConv(A, localB, localX, resvec, normRes, mm, tolb, stag, moreSteps, ...
         maxMoreSteps, imin, xmin, normResMin, flag, partition);
    if flag ~= 1
        break
    end
    % Step 12.15
    if even
        % Step 12.15.1
        beta = distributedutil.vectorProdFromLocalParts(r0, w)/rho;
        % Step 12.15.2
        rho = beta*rho;
        % Step 12.15.3
        u_mp1 = w + beta*u_m;
    end
    % Step 12.16
    pu_m = iPrecondition('mldivide', u_mp1, M1, M2, codistr);
    res = distributedutil.vectorProdFromLocalParts(pu_m, pu_m);
    if ~isfinite(res)
        flag = 2;
        return
    end
    % Step 12.17
    Au_new = iterApplyA(A, pu_m, pu_m, 'notransp', partition);
    % Step 12.18
    if even
        % Step 12.18.1
        v = Au_new + beta*(Au+beta*v);
    end
    % Step 12.19
    Au = Au_new;
    u_m = u_mp1;
    even = ~even;
end
if isempty(mm)
    mm = 0;
end

% Returned solution is first with minimal residual
if flag == 0
    iter = mm;
    relres = normResAct/normB;
else
    % Step 13
    r = localB - iterApplyA(A, xmin, xmin, 'notransp', partition);
    normRes = distributedutil.vector2NormFromLocalParts(r);
    if flag == 1 || flag == 3
        r = localB - iterApplyA(A, localX, localX, 'notransp', partition);
        normResAct = distributedutil.vector2NormFromLocalParts(r);
    end
    if normRes <= normResAct
        localX = xmin;
        iter = imin;
        relres = normRes/normB;
    else
        iter = mm;
        relres = normResAct/normB;
    end
end
% Step 14
resvec = resvec(1:mm+1);
iter = floor(iter/2);
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
