function [localX, flag, relres, iter, resvec] = hQmrImpl(codistr, A, b, tol, maxit, M1, M2, x)
% HQMRIMPL Implementation of qmr function with minimum communication using
% algorithm described in Codistributed_sparse_iterative_solvers_RFA on the wiki.
% b is codistributed 1d by rows; A.fun and A.tfun are the local parts of A and
% A' distributed 1d by rows, resp. columns, using the same partition scheme
% as b; and 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 2016-2022 The MathWorks, Inc.

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

flag = 1;
n = codistr.Cached.GlobalSize(1);

resvec = zeros(maxit+1,1);         % Preallocate vector for norms of residuals
% Step 2
% Compute with initial guess
solution = iterApplyA(A, localX, localX, 'notransp', partition);
r = localB - solution;
% 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
[y, n2minv_y] = iPrecondition('mldivideWithRcondCheck', r, M1, 'notransp', codistr);
% Step 7
[z, n2minv_z] = iPrecondition('mldivide', r, M2, 'transp', codistr);
if ~all(isfinite([n2minv_y n2minv_z]))
    flag = 2;
    resvec = normRes;
    return
end
% Step 8
rho = sqrt(n2minv_y);
psi = sqrt(n2minv_z);
% Steps 9-11
[localX, flag, normResAct, ii, resvec, imin, xmin] = iIterate(A, localB, r, localX, y,...
    z, tolb, rho, psi, partition, n, maxit, M1, M2, resvec, normRes, flag,...
    codistr);

% Returned solution is first with minimal residual
if flag == 0
    iter = ii;
    relres = normResAct/normB;
else
% Step 12
    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
% Step 13 truncate the zeros from resvec
if flag <= 1 || flag == 3
    resvec = resvec(1:ii+1);
else
    resvec = resvec(1:ii);
end
end

function [x, flag, normResAct, ii, resvec, imin, xmin] = iIterate(A, b,...
    r, x, y, z, tolb, rho, psi, partition, n, maxit, M1, M2, resvec,...
    normRes, flag, codistr)
% IITERATE is the iteration loop of the QMR algorithm.  This function is
% executed only once.  When we leave this function, either the iteration
% has converged, the specified number of iterations (maxit) has been
% reached without convergence, or an error has occurred.

xmin = x;                          % Iterate which has minimal residual so far
imin = 0;                          % Iteration at which xmin was computed
normResAct = normRes;
normResMin = normRes;              % Norm of residual from xmin

% Step 9
v = r/rho;
y = y/rho;
w = r/psi;
z = z/psi;
% Step 10
gamm = 1;
eta = -1;
thet = 0;
stag = 0;
moreSteps = 0;
maxMoreSteps = min([floor(n/50), 5, n-maxit]);
% Step 11
for ii = 1:maxit
    % Step 11.1
    delta = distributedutil.vectorProdFromLocalParts(z, y);
    flag = iterCheckParam(delta, flag);
    if flag>1
        return
    end
    gamm1 = gamm;
    thet1 = thet;
    rho1 = rho;
    % Step 11.2
    [y, n2minv_y] = iPrecondition('mldivide', y, M2, 'notransp', codistr);
    % Step 11.3
    [z, n2minv_z] = iPrecondition('mldivide', z, M1, 'transp', codistr);
    if ~all(isfinite([n2minv_y n2minv_z]))
        flag = 2;
        return
    end
    % Step 11.4
    if ii == 1
        p = y;
        q = z;
    % Step 11.5
    else
        % Step 11.5.1
        pde = psi*delta/epsilon;
        flag = iterCheckParam(pde, flag);
        % Step 11.5.2
        rde = rho*conj(delta/epsilon);
        flag = iterCheckParam(rde, flag);
        if flag>1
            return
        end
        % Step 11.5.3
        p = y - pde*p;
        % Step 11.5.4
        q = z - rde*q;
    end
    % Steps 11.6, 11.7
    [pt, epsilon] = iterApplyA(A, p, q, 'notransp', partition);
    flag = iterCheckParam(epsilon, flag);
    % Step 11.8
    beta = epsilon/delta;
    flag = iterCheckParam(beta, flag);
    if flag>1
        return
    end
    % Step 11.9
    v = pt - beta*v;
    % Step 11.10
    [y, n2minv_v] = iPrecondition('mldivide', v, M1, 'notransp', codistr);
    % Step 11.11
    w = iterApplyA(A, q, r, 'transp', partition, n) - conj(beta)*w;
    % Step 11.12
    [z, n2minv_w] = iPrecondition('mldivide', w, M2, 'transp', codistr);
    if ~all(isfinite([n2minv_v n2minv_w]))
        flag = 2;
        return
    end
    % Step 11.13
    rho = sqrt(n2minv_v);
    % Step 11.14
    psi = sqrt(n2minv_w);
    % Step 11.15
    thet = rho/(gamm*abs(beta));
    % Step 11.16
    gamm = 1/sqrt(1 + thet^2);
    flag = iterCheckParam(gamm, flag, 1);
    if flag>1
        return
    end
    % Step 11.17
    eta = -eta*rho1*gamm^2/(beta*gamm1^2);
    flag = iterCheckParam(eta, flag, 2);
    if flag>1
        return
    end
    if ii == 1
        d = eta*p;
        s = eta*pt;
    else
        % Step 11.18
        d = eta*p + (thet1*gamm)^2*d;
        % Step 11.19
        s = eta*pt + (thet1*gamm)^2*s;
    end
    % Step 11.20 Check for stagnation of the method
    norms = distributedutil.vector2NormFromLocalParts([{d} {x}]);  %[norm(d) norm(x)]
    stag = iterCheckStag(stag, norms);
    % Step 11.21 Form the new iterate
    x = x + d;
    % Step 11.22 Update the residual
    r = r - s;
    % Step 11.23
    normRes = distributedutil.vector2NormFromLocalParts(r);
    resvec(ii+1) = normRes;
    % Step 11.24 Check for convergence
    [resvec, flag, stag, moreSteps, imin, xmin, normResMin, normResAct] =...
        iterCheckConv(A, b, x, resvec, normRes, ii, tolb, stag, moreSteps, ...
        maxMoreSteps, imin, xmin, normResMin, flag, partition);
    if flag ~= 1
        return
    end
    % Step 11.25
    flag = iterCheckParam(rho, flag);
    flag = iterCheckParam(psi, flag);
    if flag>1
        return
    end
    v = v/rho;
    y = y/rho;
    w = w/psi;
    z = z/psi;
end
if isempty(ii)
    ii = 0;
end
end

function [localR, res] = iPrecondition(op, localR, M, tran, codistr)
% IPRECONDITION applies the preconditioner (if supplied) to r.  It
% uses local parts for input and output.
if ~iterIsEmptyArg(M)
    if ~iterIsEmptyArg(M)
        r = codistributed.build(localR, codistr, 'noCommunication');
    end
    switch tran
        case 'notransp'
            r = distributedutil.iterapp(op, M.fun, M.type, ...
                r, 'notransp');
        case 'transp'
            r = distributedutil.iterapp(op, M.tfun, M.type, ...
                r, 'transp');
    end
    if ~iterIsEmptyArg(M)
        localR = getLocalPart(r);
    end
end
res = distributedutil.vectorProdFromLocalParts(localR, localR);
end
