function [localX, flag, relres, iter, resvec, resveccg] = hSymmlqImpl(codistr, A, b, tol, maxit, M1, M2, x)
% HSYMMLQIMPL Implementation of symmlq 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
% 5  Preconditioner M was not symmetric positive definite

%   Copyright 2017-2022 The MathWorks, Inc.

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

resvec = zeros(maxit+1,1);         % Preallocate vector for norms of residuals
if nargout >= 6
    resveccg = zeros(maxit+2, 1);  % Preallocate vector for CG residuals
end
tolb = tol*normB;                    % Relative tolerance

% Compute with initial guess
solution = iterApplyA(A, localX, localX, 'notransp', partition);
% Step 2
r = localB - solution;
% Check if we have a solution
normRes = distributedutil.vector2NormFromLocalParts(r);
resvec(1) = normRes;                  % resvec(1) = norm(b-A*x0)
resveccg(1) = normRes;
relresmin = normRes/normB;
if relresmin <= tol                 % Initial guess is a good enough solution
    flag = 0;
    relres = relresmin;
    resvec = resvec(1);
    resveccg = resveccg(1);
    return
end

% Allocate local array
v = repmat(r, 1, 3);
index = 1:3;
[new, old] = deal(index(1), index(2));
% Step 3
[v(:, new), beta] = iPrecondition('mldivideWithRcondCheck', r, M1, M2, codistr);
if beta <= 0
    flag = 5;
    relres = relresmin;
    resvec = resvec(1);
    resveccg = resveccg(1);
    return
end
% Step 4
beta = sqrt(beta);
% Step 5
vv = v(:, new)/beta;
% Steps 6-7
[v(:, new), alpha] = iterApplyA(A, vv, vv, 'notransp', partition);
% Step 8
v(:, new) = v(:, new) - (alpha/beta)*r;
% Step 9
% Local reorthogonalization
numer = distributedutil.vectorProdFromLocalParts(vv, v(:, new));
denom = distributedutil.vectorProdFromLocalParts(vv, vv);
v(:, new) = v(:, new) + (numer/denom)*vv;
% Step 10
v(:, old) = r;
index = mod(index+1, 3)+1;
[new, old] = deal(index(1), index(2));
% Step 11
betaold = beta;
[v(:, new), beta] = iPrecondition('mldivide', v(:, old), M1, M2, codistr);
% Step 12
if beta < 0
    flag = 5;
    relres = relresmin;
    resvec = resvec(1);
    resveccg = resveccg(1);
    return
end
beta = sqrt(beta);

[flag, ii, imin, iter, resvec, resveccg, localX, xmin] = doComp(A, localB, alpha, beta,...
    betaold, partition, codistr, M1, M2, maxit, n, resvec, resveccg, tolb,...
    v, vv, localX);

% Returned solution is first with minimal residual
r = localB - iterApplyA(A, localX, localX, 'notransp', partition);
normRes = distributedutil.vector2NormFromLocalParts(r);
if isempty(ii)
    ii = 1;
end
if(flag == 0)
    relres = normRes/normB;
else
    resComp = localB - iterApplyA(A, xmin, xmin, 'notransp', partition);
    normResComp = distributedutil.vector2NormFromLocalParts(resComp);
    if normResComp <= normRes
        localX = xmin;
        iter = imin;
        relres = normResComp/normB;
        resvec(imin+1) = normResComp;
    else
        iter = ii;
        relres = normRes/normB;
    end
end

% Truncate the zeros from resvec
if(flag <= 1 || flag == 3)
    resvec = resvec(1:ii+1);
    resveccg = resveccg(1:ii+2);
else
    resvec = resvec(1:ii);
    resveccg = resveccg(1:ii+1);
end
end

function [flag, ii, imin, iter, resvec, resveccg, x, xmin] = doComp(A, b, alpha,...
    beta, betaold, partition, codistr, M1, M2, maxit, n, resvec, resveccg,...
    tolb, v, wbar, x)
flag = 1;
xmin = x;
imin = 1;                          % Iteration at which xmin was computed
index = [3, 1, 2];
[new, old, older] = deal(index(1), index(2), index(3));
iter = 0;
normResMin = resvec(1);
% Step 13
gammabar = alpha;
% Step 14
gamma = sqrt(gammabar^2 + beta^2);
% Step 15
cs = gammabar/gamma;
sn = beta/gamma;
% Step 16
zeta = betaold/gamma;
% Step 17
% Estimate norm of 1st CG iterate.
if ~iterIsEmptyArg(M1)
% Step 17.1
    resveccg(2, 1) = (betaold/gammabar)*...
        distributedutil.vector2NormFromLocalParts(v(:, old));
% Step 18
else
    % Step 18.1
    normRescgcs = abs(betaold*sn);
    % Step 18.2
    % It's possible that this cs value is zero (CG iterate does not exist)
    if cs == 0
        normRescg = Inf;
    else
        normRescg = normRescgcs/abs(cs);
    end
    resveccg(2, 1) = normRescg;
end
% Step 19
if resveccg(2, 1) <= tolb
    % Step 19.1
    x = x + (zeta/cs)*wbar;
    % Step 19.2
    nits = 0;
    iter = 1;
    flag = 0;
% Step 20
else
    % Step 20.1
    nits = maxit;
end
% Step 21
deltabar = beta;
epsilonzeta = 0;
stag = 0;
moreSteps = 0;
maxMoreSteps = min([floor(n/50),5,n-maxit]);
% Step 22
for ii = 1:nits
    % Step 22.1
    vv = v(:, new)/beta;
    % Step 22.2
    w = cs*wbar + sn*vv;
    norms = distributedutil.vector2NormFromLocalParts([{w} {x}]); 
    stag = iterCheckStag(stag, norms, zeta);
    % Step 22.3
    x = x + zeta*w;
    % Step 22.4
    wbar = sn*wbar - cs*vv;
    % Steps 22.5, 22.7
    [v(:, new), alpha] = iterApplyA(A, vv, vv, 'notransp', partition);
    % Step 22.6
    v(:, new) = v(:, new) - (beta/betaold)*v(:, older);
    % Step 22.8
    v(:, new) = v(:, new) - (alpha/beta)*v(:, old);
    % Step 22.9
    index = mod(index+1, 3)+1;
    [new, old, older] = deal(index(1), index(2), index(3));
    % Step 22.11
    betaold = beta;
    % Step 22.10
    [v(:, new), beta] = iPrecondition('mldivide', v(:, old), M1, M2, codistr);
    if beta < 0
        flag = 5;
        break
    end
    beta = sqrt(beta);
    % Step 22.12
    delta = cs*deltabar + sn*alpha;
    % Step 22.13
    deltazeta = -delta*zeta;
    % Step 22.14
    gammabar = sn*deltabar - cs*alpha;
    % Step 22.15
    epsilon = sn*beta;
    % Step 22.16
    deltabar = -cs*beta;
    % Step 22.17
    gamma = sqrt(gammabar^2 + beta^2);
    % Step 22.18
    csold = cs;
    snzeta = sn*zeta;
    % Step 22.19
    cs = gammabar/gamma;
    sn = beta/gamma;
    % Step 22.20
    epszdelz = epsilonzeta + deltazeta;
    % Step 22.21
    epsilonzeta = -epsilon*zeta;
    % Step 22.22
    zeta = epszdelz/gamma;
    % Step 22.23
    if ~iterIsEmptyArg(M1)
        % Step 22.23.1
        Res = (zeta*gamma/betaold)*v(:, older) - snzeta*v(:, old);
        normRes = distributedutil.vector2NormFromLocalParts(Res);
        % Step 22.23.2
        Res = (csold*epszdelz/gammabar - snzeta)*v(:, old);
        normRescg = distributedutil.vector2NormFromLocalParts(Res);
    % Step 22.24
    else
        % Step 22.24.1
        normRes = sqrt(epszdelz^2 + epsilonzeta^2);
        % Step 22.24.2
        normRescgcs = normRescgcs*abs(sn);
        % Step 22.24.3
        if cs == 0
            normRescg = Inf;
        else
            normRescg = normRescgcs/abs(cs);
        end
    end
    % Step 22.25
    resvec(ii+1,1) = normRes;
    resveccg(ii+2,1) = normRescg;
    % Step 22.26
    [x, flag, stag, moreSteps, imin, xmin, normResMin] = ...
        iCheckConv(A, b, x, normRes, ii, tolb, stag, moreSteps, ...
        maxMoreSteps, imin, xmin, normResMin, flag, partition);
    if flag ~= 1
        iter = ii;
        break
    end
    [x, flag, stag, moreSteps, imin, xmin, normResMin] = ...
        iCheckConv(A, b, x, normRescg, ii, tolb, stag, moreSteps, ...
        maxMoreSteps, imin, xmin, normResMin, flag, partition, ...
        epszdelz, gammabar, wbar);
    if flag ~= 1
        iter = ii;
        break
    end
end
if isempty(ii)
    ii = 0;
end
end

function [x, flag, stag, moreSteps, imin, xmin, normResMin] =...
    iCheckConv(A, b, x, normRes, ii, tolb, stag, moreSteps, ...
    maxMoreSteps, imin, xmin, normResMin, flag, partition, ...
    epszdelz, gammabar, wbar)
% ICHECKCONV checks an iteration step of the SYMMLQ algorithm for convergence.
% It is called twice, once to check the residual norm and once for the cg norm
% If the iteration has not converged, it also checks for stagnation and
% updates the best result so far.

maxStagSteps = 3;
cgcheck = false;
if nargin > 15
    cgcheck = true;
end
if normRes <= tolb || stag >= maxStagSteps || moreSteps
    if cgcheck
        xcg = x + (epszdelz/gammabar)*wbar;
    end
    r = b - iterApplyA(A, xcg, xcg, 'notransp', partition);
    normRes = distributedutil.vector2NormFromLocalParts(r);
    if (normRes <= tolb)
        x = xcg;
        flag = 0;
        return
    else
        if stag >= maxStagSteps && moreSteps == 0
            stag = 0;
        end
        moreSteps = moreSteps + 1;
        if moreSteps >= maxMoreSteps
            flag = 3;
            return
        end
    end
end
if normRes < normResMin
    normResMin = normRes;
    if cgcheck
        xmin = x + (epszdelz/gammabar)*wbar;
    else
        xmin = x;
    end
    imin = ii;
end
if stag >= maxStagSteps
    flag = 3;
    return
end
end

function [localZ, rho] = iPrecondition(op, localR, M1, M2, 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');
    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);
    localR = getLocalPart(r);
else
    localZ = localR;
end
rho = distributedutil.vectorProdFromLocalParts(localR, localZ);
end
