function [localX, flag, relres, iter, resvec, lsvec] = hLsqrImpl(codistrX, A, b, tol, maxit, M1, M2, x, emptyMaxit)
% HLSQRIMPL Implementation of lsqr 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. A.fun uses the partition scheme
% of x and A.tfun uses the partition scheme of b; M1 and M2, if they are matrices,
% are the local parts of M1 and M2 distributed by rows according to the partition
% scheme of x.  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(codistrX, 1);
partitionX = matlab.internal.ColonDescriptor(partitionStart, partitionEnd);
n = codistrX.Cached.GlobalSize(1);
localB = getLocalPart(b);
localX = getLocalPart(x);

flag = 1;

% Compute with initial guess
% Step 2
u = localB;
if any(spmdCat(nnz(localX) > 0))
    u = u - iterApplyA(A, localX, localX, 'notransp', partitionX);
end

% Step 3
norms = distributedutil.vector2NormFromLocalParts([{localB} {u}]);
normB = norms(1);              % 2-norm of rhs (b)
beta = norms(2);               % 2-norm of u
tolb = tol*normB;
normRes = beta;

% Step 4
if beta ~= 0
    u = u/beta;
end

relres = beta/normB;
iter = 0;
resvec = beta;
lsvec = zeros(0, 1);

% Step 5
v = iterApplyA(A, u, u, 'transp', partitionX, n);

% Step 6
[v, n2minv_v] = iPrecondition('mldivideWithRcondCheck', v, M1, M2, 'transp', codistrX);
if ~isfinite(n2minv_v)
    flag = 2;
    return
end

% Step 7
alpha = distributedutil.vector2NormFromLocalParts(v);

% Step 8
if alpha ~= 0
    v = v/alpha;
end

% Step 9 Check for exact solution
normar = alpha*beta;
if normar == 0
    flag = 0;
    return
end

% Step 10
c = 1;
s = 0;
norma = 0;
phibar = beta;
d = zeros(size(localX));
if emptyMaxit
    maxit = min(n, maxit);
end
resvec = zeros(maxit+1, 1);
resvec(1) = beta;
lsvec = zeros(maxit, 1);
stag = 0;
iter = maxit;
maxStagSteps = 3;

% Step 11 loop over maxit iterations (unless convergence or failure)

for ii = 1:maxit
    
    % Step 11.1
    [z, n2minv_z] = iPrecondition('mldivide', v, M1, M2, 'notransp', codistrX);
    if ~isfinite(n2minv_z)
        flag = 2;
        return
    end

    % Step 11.2
    u = iterApplyA(A, z, z, 'notransp', partitionX) - alpha*u;
    
    % Step 11.3
    beta = distributedutil.vector2NormFromLocalParts(u);
    
    % Step 11.4
    u = u/beta;
    
    % Step 11.5
    norma = sqrt([norma alpha beta]*[norma; alpha; beta]);
    
    % Step 11.6
    lsvec(ii) = normar/norma;
    
    % Step 11.7
    if normar/(norma*normRes) <= tol || normRes <= tolb
        flag = 0;
        iter = ii-1;
        resvec = resvec(1:ii);
        lsvec = lsvec(1:ii);
        break
    end
    
    % Step 11.8
    theta = -s*alpha;
    
    % Step 11.9
    rhot = c*alpha;
    
    % Step 11.10
    rho = sqrt(rhot^2 + beta^2);

    % Step 11.11
    d = (z - theta*d)/rho;
    
    % Step 11.12
    c = rhot/rho;
    
    % Step 11.13
    s = -beta/rho;
    
    % Step 11.14
    phi = c*phibar;
    
    % Step 11.15
    phibar = s*phibar;
    norms = distributedutil.vector2NormFromLocalParts([{d} {localX}]);  %[norm(d) norm(x)]
    stag = iterCheckStag(stag, norms, phi);
    if stag >= maxStagSteps
        flag = 3;
        iter = ii-1;
        resvec = resvec(1:ii);
        lsvec = lsvec(1:ii);
        break
    end

    % Step 11.16
    localX = localX + phi*d;
    
    % Step 11.17
    normRes = abs(s)*normRes;
    
    % Step 11.18
    resvec(ii+1) = normRes;
    
    % Step 11.19
    z = iterApplyA(A, u, u, 'transp', partitionX, n);
    
    % Step 11.20
    [z, n2minv_z] = iPrecondition('mldivide', z, M1, M2, 'transp', codistrX);
    if ~isfinite(n2minv_z)
        flag = 2;
        iter = ii;
        resvec = resvec(1:ii+1);
        lsvec = lsvec(1:ii);
        break
    end
    
    % Step 11.21
    v = z - beta*v;
    
    % Step 11.22
    alpha = distributedutil.vector2NormFromLocalParts(v);
    
    % Step 11.23
    v = v/alpha;
    
    % Step 11.24
    normar = alpha*abs(s*phi);
end
if isempty(ii)
    iter = 0;
end

% Step 12
if flag == 1
    if normar/(norma*normRes) <= tol || normRes <= tolb
        flag = 0;
        iter = maxit;
    end
end

relres = normRes/normB;
end

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