function [localX, flag, relres, iter, resvec] = hGmresImpl(codistr, A, b, restart, tol, maxit, M1, M2, x)
% HGMRESIMPL Implementation of gmres 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);
m = codistr.Cached.GlobalSize(1);
localB = getLocalPart(b);
localX = getLocalPart(x);

normB = sqrt(sum(spmdCat(localB'*localB)));        % 2-norm  of rhs (b)
flag = 1;

resvec = zeros(restart*maxit+1, 1); % Preallocate vector for norm of residuals
stag = 0;

% Compute with initial guess
solution = iterApplyA(A, localX, localX, 'notransp', partition);
% Step 2
r = localB - solution;
% Check if residual is small enough
buffer = distributedutil.vector2NormFromLocalParts([{r} {localX}]); 
resvec(1) = buffer(1);              % resvec(1) = norm(b-A*x0)
xnorm = buffer(2);
x0iszero = (xnorm == 0);
relresmin = resvec(1)/normB;
if relresmin <= tol
    resvec = resvec(1);
    flag = 0;
    relres = relresmin;
    iter = [0 0];
    return
end

% Step 3
[r, n2minv_r] = iPrecondition('mldivideWithRcondCheck', r, M1, M2, codistr);
if ~x0iszero
    [~, n2minv_b] = iPrecondition('mldivide', localB, M1, M2, codistr);
else
    n2minv_b = n2minv_r;
end
if ~all(isfinite([n2minv_r n2minv_b]))
    resvec = resvec(1);
    flag = 2;
    relres = relresmin;
    iter = [0 0];
    return
end

normRes = sqrt(n2minv_r);       % norm of the preconditioned residual
normResMin = normRes;           % norm of residual from xmin
resvec(1) = normRes;
n2minv_b = sqrt(n2minv_b);      % norm of the preconditioned rhs

% Step 4
tolb = tol*n2minv_b;
if (normRes <= tolb)            % Initial guess is a good enough solution
    resvec = n2minv_b;
    flag = 0;
    relres = normRes/n2minv_b;
    iter = [0 0];
    return
end

% Preallocate J to hold the Givens rotation constants
U = zeros(size(localX, 1), restart);
if spmdIndex == 1
    w = zeros(restart+1, 1);
else
    w = [];
end

% Step 5
% Set up struct to hold input parameters
inputs = struct('A', A, 'b', localB, 'M1', M1, 'M2', M2, 'partition',...
    partition, 'codistr', codistr, 'm', m, 'maxit', maxit, 'tolb', tolb,...
    'restart', restart);
for outiter = 1:maxit
    % Step 5.1
    u = r;
    sgn = scalarsign(r, 1, 1);
    sgn = spmdPlus(sgn);
    normRes = distributedutil.vector2NormFromLocalParts(r);
    beta = sgn*normRes;
    if (spmdIndex == 1)
        u(1) = u(1) + beta;
    end
    unorm = distributedutil.vector2NormFromLocalParts(u);
    u = u/unorm;
    U(:, 1) = u;
    
    %Step 5.2
    w(1) = -beta;
    
    % Step 5.3
    [w, U, R, flag, resvec, localX, iter, mins, normResAct, initer] =...
        iInnerLoop(u, w, U, resvec, localX, outiter, inputs, normResMin);

    evalxm = 0;
    
    if flag ~= 0
        % Step 5.4
        if mins.minupdated
            idx = mins.jmin;
        else
            idx = initer;
        end
        % Steps 5.5-5.8
        localX = iEvalXm(R, w, U, localX, idx, evalxm, stag);
        mins.xmin = localX;
        % Step 5.9
        r = localB - iterApplyA(A, localX, localX, 'notransp', partition);
        % Step 5.10
        [minv_r, n2minv_r] = iPrecondition('mldivide', r, M1, M2, codistr);
        if ~isfinite(n2minv_r)
            flag = 2;
            break
        end
        normResAct = sqrt(n2minv_r);
        % Step 5.11
        r = minv_r;
    end
    if normResAct <= mins.normrmin
        mins.xmin = localX;
        mins.normrmin = normResAct;
        mins.imin = outiter;
        mins.jmin = initer;
    end
    if flag == 3
        break
    end
    if normResAct <= tolb
        flag = 0;
        iter = [outiter, initer];
        break
    end
    normResMin = mins.normrmin;
end
if isempty(outiter)
    outiter = 0;
    initer = 0;
    iter = [outiter, initer];
    normResAct = normResMin;
    mins.xmin = localX;
    mins.normrmin = normResAct;
    mins.imin = outiter;
    mins.jmin = initer;
end

% returned solution is that with minimum residual
if flag ~= 0
    localX = mins.xmin;
    iter = [mins.imin mins.jmin];
end
relres = normResAct/n2minv_b;
resvec = resvec(1:max(outiter-1,0)*restart+initer+~(flag == 2 && initer ~= 0));
end

function [w, U, R, flag, resvec, x, iter, mins, normResAct, initer] =...
        iInnerLoop(u, w, U, resvec, x, outiter, inputs, normResMin)
% iInnerLoop does an Arnoldi iteration to minimize the residual.

moreSteps = 0;
evalxm = 0;
stag = 0;
flag = 1;
iter = [0 0];
mins = struct('imin', 0, 'jmin', 0, 'xmin', x, 'minupdated', 0,...
    'normrmin', normResMin);
restart = inputs.restart;
R = [];
if spmdIndex == 1
    J = zeros(2, restart);
    R = zeros(restart, restart);
end

for initer = 1:restart
    % Step 5.3.1
    if spmdIndex == 1
        uiniter = spmdBroadcast(1, u(initer));
    else
        uiniter = spmdBroadcast(1);
    end
    v = -2*(uiniter')*u;
    % Step 5.3.2
    if spmdIndex == 1
        v(initer) = v(initer) + 1;
    end
    % Step 5.3.3
    for k = (initer-1):-1:1
        % Step 5.3.3.1
        Uxv = distributedutil.vectorProdFromLocalParts(U(:, k), v);
        v = v - U(:, k)*2*Uxv;
    end
    % Step 5.3.4
    normv = distributedutil.vector2NormFromLocalParts(v);
    v = v/normv;
    % Step 5.3.5
    v = iterApplyA(inputs.A, v, v, 'notransp', inputs.partition);
    % Step 5.3.6
    [v, n2minv_v] = iPrecondition('mldivide', v, inputs.M1, inputs.M2, inputs.codistr);
    if ~isfinite(n2minv_v)
        flag = 2;
        break
    end
    % Step 5.3.7
    for k = 1:initer
        % Step 5.3.7.1
        Uxv = distributedutil.vectorProdFromLocalParts(U(:, k), v);
        v = v - U(:, k)*2*Uxv;
    end
    if (initer ~= inputs.m)
        % Step 5.3.8
        if spmdIndex == 1
            u = [zeros(initer, 1); v((initer + 1):end)];
        else
            u = v;
        end
        % Step 5.3.9
        alpha = distributedutil.vector2NormFromLocalParts(u);
        if (alpha ~= 0)
            sgn = spmdPlus(scalarsign(v, 1, initer+1));
            alpha = sgn*alpha;
        end
        % Step 5.3.10
        if spmdIndex == 1
            u(initer+1) = u(initer+1) + alpha;
        end
        % Step 5.3.11
        normu = distributedutil.vector2NormFromLocalParts(u);
        u = u/normu;
        U(:, initer+1) = u;
        % Step 5.3.12
        if spmdIndex == 1
            v(initer+1) = -alpha;
            v((initer+2):end) = 0;
        else
            v = 0;
        end
    end
    if spmdIndex == 1
        [w, J, R] = iGivensRotation(initer, inputs.m, v, w, J, R, restart);
        % Step 5.3.20
        normRes = abs(w(initer+1));
        spmdBroadcast(1, normRes);
    else
        normRes = spmdBroadcast(1);
    end
    resvec((outiter-1)*restart+initer+1) = normRes;
    normResAct = normRes;
    % Step 5.3.21
    [flag, x, mins, iter, evalxm, resvec] = iCheckConvergence(R, w, U, ...
        x, initer, outiter, evalxm, resvec, normRes, normResMin, inputs, ...
        stag, moreSteps);
    if flag ~= 1
        break
    end
end
if isempty(initer)
    initer = 0;
    normResAct = normResMin;
end

end

function [flag, x, mins, iter, evalxm, resvec] = iCheckConvergence(R, w, U,...
        x, initer, outiter, evalxm, resvec, normRes, normResMin, inputs, stag,...
        moreSteps)
flag = 1;
iter = [0, 0];
mins = struct('imin', 0, 'jmin', 0, 'xmin', x, 'minupdated', 0,...
    'normrmin', normResMin);
maxStagSteps = 3;
maxMoreSteps = min([floor(inputs.m/50), 5, inputs.m-inputs.maxit]);
normResAct = normRes;    
if (normRes <= inputs.tolb || stag >= maxStagSteps || moreSteps)
    [xm, evalxm, stag] = iEvalXm(R, w, U, x, initer, evalxm, stag);
    r = inputs.b - iterApplyA(inputs.A, xm, xm, 'notransp', ...
        inputs.partition);
    normRes = distributedutil.vector2NormFromLocalParts(r);
    if normRes <= inputs.tolb
        x = xm;
        flag = 0;
        iter = [outiter, initer];
        return
    end
    [~, n2minv_r] = iPrecondition('mldivide', r, inputs.M1, inputs.M2, inputs.codistr);
    if ~isfinite(n2minv_r)
        flag = 2;
        return
    end
    normResAct = sqrt(n2minv_r);
    resvec((outiter-1)*inputs.restart+initer+1) = normResAct;
    if normResAct <= mins.normrmin
        mins.normrmin = normResAct;
        mins.imin = outiter;
        mins.jmin = initer;
        mins.xmin = xm;
        mins.minupdated = 1;
    end

    if normResAct < inputs.tolb
        x = xm;
        flag = 0;
        iter = [outiter, initer];
        return
    else
        if stag >= maxStagSteps && moreSteps == 0
            stag = 0;
        end
        moreSteps = moreSteps + 1;
        if moreSteps >= maxMoreSteps
            flag = 3;
            return
        end
    end
end

if normResAct <= mins.normrmin
    mins.normrmin = normResAct;
    mins.imin = outiter;
    mins.jmin = initer;
    mins.minupdated = 1;
end

if stag >= maxStagSteps
    flag = 3;
    return
end
end

function [localZ, res] = 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
res = distributedutil.vectorProdFromLocalParts(localZ, localZ);
end

function [xm, evalxm, stag] = iEvalXm(R, w, U, xm, initer, evalxm, stag)
% Evaluate correction to x
if evalxm == 0
    if spmdIndex == 1
        ytmp = R(1:initer, 1:initer) \ w(1:initer);
        const = -2*ytmp(initer)*conj(U(initer, initer));
        ytmp = [ytmp; const];
        spmdBroadcast(1, ytmp);
    else
        ytmp = spmdBroadcast(1);
    end
    additive = U(:, initer)*ytmp(initer+1);
    if spmdIndex == 1
        additive(initer) = additive(initer) + ytmp(initer);
    end
    for k = initer-1:-1:1
        if spmdIndex == 1
            additive(k) = additive(k) + ytmp(k);
        end
        additive = additive - U(:, k)*(2*...
            distributedutil.vectorProdFromLocalParts(U(:, k), additive));
    end
    norms = distributedutil.vector2NormFromLocalParts([{additive} {xm}]);
    stag = iterCheckStag(stag, norms);
    xm = xm + additive;
    evalxm = 1;
else
    if spmdIndex == 1
        addvc = [-(R(1:initer-1, 1:initer-1)\R(1:initer-1, initer))*...
            (w(initer)/R(initer, initer)); w(initer)/R(initer, initer)];
        const = -2*addvc(initer)*conj(U(initer, initer));
        addvc = [addvc; const];
        spmdBroadcast(1, addvc);
    else
        addvc = spmdBroadcast(1);
    end
    norms(1) = norm(addvc(1:initer));
    norms(2) = distributedutil.vector2NormFromLocalParts(xm);
    stag = iterCheckStag(stag, norms);
    additive = U(:, initer)*addvc(initer+1);
    if spmdIndex == 1
        additive(initer) = additive(initer) + addvc(initer);
    end
    for k = initer-1:-1:1
        if spmdIndex == 1
            additive(k) = additive(k) + addvc(k);
        end
        additive = additive - U(:, k)*(2*...
            distributedutil.vectorProdFromLocalParts(U(:, k), additive));
    end
    xm = xm+additive;
end
end

function [w, J, R] = iGivensRotation(initer, m, v, w, J, R, restart)
% Do Givens rotation on worker 1
% Step 5.3.13
for colJ = 1:initer-1
    tmpv = v(colJ);
    v(colJ) = conj(J(1, colJ))*v(colJ) + conj(J(2, colJ))*v(colJ+1);
    v(colJ+1) = -J(2, colJ)*tmpv + J(1, colJ)*v(colJ+1);
end
% Step 5.3.14
if ~(initer==m)
    rho = norm(v(initer:initer+1));
    J(:, initer) = v(initer:initer+1)./rho;
    % Step 5.3.15
    w(initer+1) = -J(2, initer)*w(initer);
    % Step 5.3.16
    w(initer) = conj(J(1, initer))*w(initer);
    % Step 5.3.17
    v(initer) = rho;
    % Step 5.3.18
    v(initer+1) = 0;
end
% Step 5.3.19
R(:, initer) = v(1:restart);
end

function sgn = scalarsign(d, workerIdx, index)
% Function to determine the sign of element index of array d on a specified
% worker.  On the worker whose spmdIndex is workerIdx, it returns 1 if the
% element is non-negative, and -1 if it is negative.  On all other workers,
% it returns 0.  spmdPlus(sgn) returns the same result: (1 or -1) on all workers.
if spmdIndex ~= workerIdx
    sgn = 0;
else
    sgn = sign(d(index));
    if (sgn == 0)
        sgn = 1;
    end
end

end
