function varargout = tfqmr(varargin)
%TFQMR   Transpose Free Quasi-Minimal Residual Method for gpuArray.
%   X = TFQMR(A,B)
%   X = TFQMR(A,B)
%   X = TFQMR(A,B,TOL)
%   X = TFQMR(A,B,TOL,MAXIT)
%   X = TFQMR(A,B,TOL,MAXIT,M)
%   X = TFQMR(A,B,TOL,MAXIT,M1,M2)
%   X = TFQMR(A,B,TOL,MAXIT,M1,M2,X0)
%   [X,FLAG] = TFQMR(A,B,...)
%   [X,FLAG,RELRES] = TFQMR(A,B,...)
%   [X,FLAG,RELRES,ITER] = TFQMR(A,B,...)
%   [X,FLAG,RELRES,ITER,RESVEC] = TFQMR(A,B,...)
%      
%   Note that B must be a non-sparse column vector.
%   
%   Example:
%   Using TFQMR 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; 
%       
%       dX = TFQMR(dA,db,tol,maxit)
%   
%   See also TFQMR, GPUARRAY.


%   Copyright 1984-2019 The MathWorks, Inc.

if (nargin < 2)
    error(message('MATLAB:tfqmr:NotEnoughInputs'));
end

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

% Determine properties of the inputs
import parallel.internal.gpu.iterchkBasic;
[~, ~, ~, issparseA] = iterchkBasic(1, varargin);
[~, existsM1, ~, issparseM1] = iterchkBasic(5, varargin);
[~, existsM2, ismatrixM2, ~] = iterchkBasic(6, 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}] = iTfqmrBuiltinWrapper(varargin{:});
else % Go to flowthrough instead. This covers dense matrices and functions
    [varargout{1:nargout}] = parallel.internal.flowthrough.tfqmr(varargin{:});
end

end

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

% Check matrix and right hand side vector inputs have appropriate sizes
[m,n] = size(A);
if (m ~= n)
    error(message('MATLAB:tfqmr:NonSquareMatrix'));
end
if ~iscolumn(b)
    error(message('MATLAB:tfqmr:RSHnotColumn'));
end
if ~isequal(size(b),[m,1])
    error(message('MATLAB:tfqmr:RSHsizeMatchCoeffMatrix', m));
end

% Assign default values to unspecified parameters
if (nargin < 3) || isempty(tol)
    tol = 1e-6;
end
if tol <= eps
    warning(message('MATLAB:tfqmr:tooSmallTolerance'));
    tol = eps;
elseif tol >= 1
    warning(message('MATLAB:tfqmr:tooBigTolerance'));
    tol = 1-eps;
end
% Assume that maxit was specified.
maxitSpecified = true;
if nargin < 4 || isempty(maxit)
    maxitSpecified = false;
    maxit = min(m,20);
end

if nargin >= 5
    if ~isempty(M1) && ~isequal(size(M1),[n,n])
        error(message('MATLAB:tfqmr:WrongPrecondSize', n));
    end
else
    M1 = [];
end

% Validate x0 input
if ((nargin >= 7) && ~isempty(x))
    if ~isequal(size(x),[n,1])
        error(message('MATLAB:tfqmr:WrongInitGuessSize', n));
    end
else
    x = zeros(n,1,'like',b);
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 (MOD 1)
    flag = 0;                      % a valid solution has been obtained
    relres = 0;                    % the relative residual is actually 0/0
    iter = 0;                      % no iterations need be performed
    resvec = 0;                    % resvec(1) = norm(b-A*x) = norm(0)
    if (nargout < 2)
        parallel.internal.flowthrough.itermsg('tfqmr',tol,maxit,0,flag,iter,NaN);
    end
    return
end

% Finally, make sure that maxit is not larger than n, if maxit was not
% specified by our caller.
if ~maxitSpecified
    maxit = min(n,maxit);
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;             % No iterations were performed
    return;
end

% Additional inputs are allowed for flowthrough (for functions), but
% not in this code path
if (nargin > 7)
    error(message('MATLAB:tfqmr: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
restart = 1; % Unused but a required argument for the built-ins
[x, flag, finalRes, iter, resvecOut{:}] = ...
    hTfqmrCore(A, gpuArray(b), restart, tol, 2 * maxit, ...
              gpuArray(M1), [], gpuArray(full(x)));
iter = iter(1); % Only GMRES uses both parts of output iter

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

% Check that the result really did converge. If it didn't we can put this
% down to an ill-conditioned preconditioner.
% A non-finite residual is a certain indicator of preconditioner problems,
% otherwise we can check against the tolerance, but only for square
% matrices.
if flag == 0 && ~isempty(M1) && ...
         ( ~isfinite(relres) || relres > tol )
    flag = 2;
end

% Copy out the residuals, and truncate zeros
if nargout >= 5
    resvec = resvecOut{1};
    resvec = resvec(1:iter+1);
    if flag == 2 && iter ~= 0
        resvec(end) = [];
    end
end
iter = floor((iter + 1) / 2);

% Only display a message if the output flag is not used
if nargout < 2
    parallel.internal.flowthrough.itermsg('tfqmr',tol,maxit,iter,flag,iter,relres);
end

end
