function varargout = gmres(varargin)
%GMRES   Generalized Minimum Residual Method for gpuArray.
%   X = GMRES(A,B)
%   X = GMRES(A,B,RESTART)
%   X = GMRES(A,B,RESTART,TOL)
%   X = GMRES(A,B,RESTART,TOL,MAXIT)
%   X = GMRES(A,B,RESTART,TOL,MAXIT,M)
%   X = GMRES(A,B,RESTART,TOL,MAXIT,M1,M2)
%   X = GMRES(A,B,RESTART,TOL,MAXIT,M1,M2,X0)
%   [X,FLAG] = GMRES(A,B,...)
%   [X,FLAG,RELRES] = GMRES(A,B,...)
%   [X,FLAG,RELRES,ITER] = GMRES(A,B,...)
%   [X,FLAG,RELRES,ITER,RESVEC] = GMRES(A,B,...)
%   
%   For sparse matrix input A, the following apply:
%   * Only one sparse matrix preconditioner M1 is supported.
%   * For a sparse matrix preconditioner, RESVEC returns the residual
%     norms NORM(B-A*X) without applying the preconditioner.
%   * To use two preconditioners, M1 and M2, they must both be functions.
%   * Stagnation (FLAG=3) is not detected and will be reported as failure
%     to converge (FLAG=1).
%   
%   Example:
%   Using GMRES to solve a large system of gpuArray data:
%    
%       n = 1e6; 
%       dOnes = gpuArray.ones(n,1); 
%       dA = spdiags([-2*dOnes 4*dOnes -dOnes],-1:1,n,n);
%       db = full(sum(dA,2)); 
%       tol = 1e-8; 
%       maxit = 15; 
%       dM = spdiags(4*dOnes, 0, n, n);
%   
%       dX = gmres(dA,db,10,tol,maxit,dM)
%   
%   See also GMRES, GPUARRAY.


%   Copyright 1984-2020 The MathWorks, Inc.

% Override iterchkBasic
import parallel.internal.gpu.iterchkBasic;

if (nargin < 2)
    error(message('MATLAB:gmres:NumInputs'));
end

% Ensure that only the data inputs (args 1,2,6,7,8) are on the GPU
varargin = prepareSolverInputs(varargin, [1 1 0 0 0 1 1 1]);

% Determine properties of the inputs
[~, ~, ~, issparseA] = iterchkBasic(1, varargin);
[~, existsM1, ~, issparseM1] = iterchkBasic(6, varargin);
[~, existsM2, ismatrixM2, ~] = iterchkBasic(7, varargin);

% If M1 is a sparse matrix, only the built-in can use it, so A has to be a
% sparse matrix too
if issparseM1 && ~issparseA
    error(message('parallel:gpu:array:SparseSolverPreconditionersNotSupported'));
end
% M2 being a matrix cannot work with a sparse matrix A
if ismatrixM2 && issparseA
    error(message('parallel:gpu:array:SparseSolverM2NotSupported'));
end

% Execute built-in where possible
if issparseA && (~existsM1 || issparseM1) && ~existsM2
    [varargout{1:nargout}] = iGmresBuiltinWrapper(varargin{:});
else % Go to flowthrough instead. This covers dense matrices and functions
    [varargout{1:nargout}] = parallel.internal.flowthrough.gmres(varargin{:});
end

end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function [x,flag,relres,iter,resvec] = iGmresBuiltinWrapper(A,b,restart,tol,maxit,M1,~,x,varargin)
% Wrap built-in to do validation and to support variable inputs and outputs

% Override the standard itermsg
import parallel.internal.flowthrough.itermsg;

% Check matrix and right hand side vector inputs have appropriate sizes
[m,n] = size(A);
if (m ~= n)
    error(message('MATLAB:gmres:SquareMatrix'));
end
if ~isequal(size(b),[m,1])
    error(message('MATLAB:gmres:VectorSize', m));
end
if nargin >= 6
    if ~isempty(M1) && ~isequal(size(M1),[m,m])
        error(message('MATLAB:gmres:PreConditioner1Size', m));
    end
else
    M1 = [];
end

% Assign default values to unspecified parameters
if (nargin < 3) || isempty(restart) || (restart == n)
    restarted = false;
else
    restarted = true;
end
if (nargin < 4) || isempty(tol)
    tol = 1e-6;
end
if tol < eps
    warning(message('MATLAB:gmres:tooSmallTolerance'));
    tol = eps;
elseif tol >= 1
    warning(message('MATLAB:gmres:tooBigTolerance'));
    tol = 1-eps;
end
if (nargin < 5) || isempty(maxit)
    if restarted
        maxit = min(ceil(n/restart),10);
    else
        maxit = min(n,10);
    end
end

if restarted
    outer = maxit;
    if restart > n
        warning(message('MATLAB:gmres:tooManyInnerItsRestart',restart, n));
        restart = n;
    end
    inner = restart;
else
    outer = 1;
    if maxit > n
        warning(message('MATLAB:gmres:tooManyInnerItsMaxit',maxit, n));
        maxit = n;
    end
    inner = maxit;
end

% Deal with sparse b by converting to dense - for a column vector this
% actually takes up less space
b = full(b);

% Check for all zero right hand side vector => all zero solution
n2b = norm(b);                   % Norm of rhs vector, b
if (n2b == 0)                    % if    rhs vector is all zeros
    x = zeros(n,1,'like',b);     % then  solution is all zeros
    flag = 0;                    % a valid solution has been obtained
    relres = 0;                  % the relative residual is actually 0/0
    iter = [0 0];                % no iterations need be performed
    resvec = 0;                  % resvec(1) = norm(b-A*x) = norm(0)
    if (nargout < 2)
        parallel.internal.flowthrough.itermsg('gmres',tol,maxit,0,flag,iter,NaN);
    end
    return
end

% Validate x0 input
if ((nargin >= 8) && ~isempty(x))
    if ~isequal(size(x),[n,1])
        error(message('MATLAB:gmres:XoSize', n));
    end
else
    x = zeros(n,1,'like',b);
end

% Trap all-zero preconditioner - if passed to built-in it will assert
if ~isempty(M1) && nnz(M1) == 0
    warning(message('MATLAB:singularMatrix')); % Same warning as core MATLAB
    flag = 2;             % Ill-conditioned preconditioner
    resvec = gather(norm(b-A*x)); % Final residual same as starting residual
    relres = resvec/gather(n2b);  % Compute relative residual
    iter = [0 0];         % No iterations were performed
    return;
end

% Additional inputs are allowed for flowthrough (for functions), but
% not in this code path
if (nargin > 8)
    error(message('MATLAB:gmres:TooManyInputs'));
end

%**** CALL GPUARRAY METHOD *****
% Avoid always asking for the residual vector since it causes
% unnecessary memory accesses
resvecOut = cell(nargout-4); % Empty if no resvec required
[x, flag, finalRes, iter, resvecOut{:}] = ...
    hGmresCore(A, gpuArray(b), inner, tol, outer, ...
               gpuArray(M1), [], gpuArray(full(x)));

% Convert final residual into a relative residual
relres = finalRes / gather(n2b);

% If the preconditioner is ill-conditioned we can sometimes detect it
% because the iterations may terminate early thinking they've converged,
% but will then discover they haven't converged when checked against the
% unpreconditioned input, and so return the failed to converge flag (1).
if ~isempty(M1) && flag == 1 && iter(1) < outer
    flag = 2;
end

% Copy out the residuals, and truncate zeros
if nargout >= 5
    resvec = resvecOut{1};
    initer = iter(2);
    % If the initial solution was good, iter(1)==0, but outiter is
    % officially still 1
    outiter = max(1, iter(1));
    resvec = resvec(1:(outiter-1)*inner+initer+1);
    if flag == 2 && initer ~= 0
        resvec(end) = [];
    end
end

% Only display a message if the output flag is not used
if nargout < 2
    if restarted
        itermsg(sprintf('gmres(%d)',restart),tol,maxit,iter,flag,iter,relres);
    else
        itermsg(sprintf('gmres'),tol,maxit,iter(2),flag,iter(2),relres);
    end
end

    
end
