function [X,info] = IRrrgmres(A,b,varargin)
% IRrrgmres Range Restricted GMRES for square systems
%
% options  = IRrrgmres('defaults')
% [X,info] = IRrrgmres(A,b);
% [X,info] = IRrrgmres(A,b,K);
% [X,info] = IRrrgmres(A,b,options);
% [X,info] = IRrrgmres(A,b,K,options);
%
% This function implements the Range Restricted GMRES method for solving
% a square system.
%
% With 'defaults' as input returns the default options.  Otherwise outputs
% the iterates specified in K, using max(K) as MaxIter, and using all other
% default options.  With options as input: uses the user-specified options
% and all the other default options.
%
% Inputs:
%  A : either (a) a full or sparse matrix
%             (b) a matrix object that performs the matrix*vector operation
%             (c) user-defined function handle
%  b : right-hand side vector
%  K : (optional) integer vector that specifies which iterates are returned
%      in X; the maximum number of iterations is assumed to be max(K)
%      [ positive integer | vector of positive components ]
%  options : structure with the following fields (optional)
%      x0         - initial guess for the iterations; default = zero vector
%                   [ array | {'none'} ]
%      MaxIter    - maximum allowed number of iterations
%                   [ {100} | positive integer ]
%                   NOTE: K overrules MaxIter if both are assigned
%      x_true     - true solution; allows us to returns error norms with
%                   respect to x_true at each iteration
%                   [ array | {'none'} ]
%      NoiseLevel - norm of noise in rhs divided by norm of rhs 
%                   [ {'none'} | nonnegative scalar ]
%      eta        - safety factor for the discrepancy principle
%                   [ {1.01} | scalar greater than (and close to) 1 ]
%      DecompOut  - return the decomposition to the user
%                   [ 'on' | {'off'} ]
%      IterBar    - shows the progress of the outer iterations
%                   [ {'on'} | 'off' ]
%      NoStop     - specifies weather the iterations should proceed
%                   after a stopping criterion has been satisfied
%                   [ 'on' | {'off'} ]
% Note: the options structure can be created using the function IRset.
%
% Outputs:
%   X : computed solutions, stored column-wise (at the iterations listed in K)
%   info: structure with the following fields:
%      its     - number of the last computed iteration 
%      saved_iterations - iteration numbers of iterates stored in X 
%      StopFlag - stringg that describes the output/stopping condition:
%                   * Performed max number of iterations
%                   * Residual tolerance satisfied (discrepancy principle) 
%                   * Normal equationa residual tolerance satisfied
%      Rnrm     - relative residual norms at each iteration
%      Xnrm     - solution norms at each iteration
%      Enrm     - error norms (requires x_true) at each iteration
%      V        - matrix generated by the Arnoldi algorithm: the orthonormal
%                 columns of V are a basis for the RRGMRES solution
%      Q        - orthogonal matrix in QR factorization of Hessenberg matrix
%      T        - triangular matrix in QR factorization of Hessenberg matrix
%      StopReg  - struct containing information about the solution that
%                 satisfies the stopping criterion.  Fields:
%                   It   : iteration where stopping criterion is satisfied
%                   X    : solution satisfying the stopping criterion
%                   Enrm : the corresponding relative error (requires x_true)
%      BestReg  - struct containing information about the solution that
%                 minimizes Enrm (requires x_true). Fields:
%                   It   : iteration where the minimum is attained
%                   X    : best solution
%                   Enrm : best relative error
% 
% See also: IRcgls, IRhybrid_gmres, IRget, IRset

% Silvia Gazzola, University of Bath
% Per Christian Hansen, Technical University of Denmark
% James G. Nagy, Emory University
% April, 2018.

% This file is part of the IR Tools package and is distributed under the 
% 3-Clause BSD License. A separate license file should be provided as part 
% of the package.

% Set default values for options.
defaultopt = struct('x0','none', 'MaxIter',100, 'x_true','none', ...
    'NoiseLevel','none', 'eta',1.01, 'IterBar','on', ...
    'NoStop','off', 'DecompOut','off');
  
% If input is 'defaults,' return the default options in X.
if nargin==1 && nargout <= 1 && isequal(A,'defaults')
    X = defaultopt;
    return;
end

defaultopt.restart = 'off';
defaultopt.verbosity = 'on';

% Check for acceptable number of optional input arguments.
switch length(varargin)
    case 0
        K = []; options = [];
    case 1
        if isa(varargin{1}, 'double')
            K = varargin{1}; options = [];
        else
            K = []; options = varargin{1};
        end
    case 2
        if isa(varargin{1}, 'double')
            K = varargin{1}; options = varargin{2};
        else
            K = varargin{2}; options = varargin{1};
        end
        if isfield(options, 'MaxIter') && ~isempty(options.MaxIter) && (~isempty(K) && options.MaxIter ~= max(K))
            warning('The value of MaxIter is discarded; the maximum value in K is taken as MaxIter')
        end 
    otherwise
        error('Too many input parameters')
end

if isempty(options)
    options = defaultopt;
end

options = IRset(defaultopt, options);

MaxIter    = IRget(options, 'MaxIter',    [], 'fast');
x_true     = IRget(options, 'x_true',     [], 'fast');
NoiseLevel = IRget(options, 'NoiseLevel', [], 'fast');
eta        = IRget(options, 'eta',        [], 'fast');
IterBar    = IRget(options, 'IterBar',    [], 'fast');
NoStop     = IRget(options, 'NoStop',     [], 'fast');
restart    = IRget(options, 'restart',    [], 'fast');
verbose    = IRget(options, 'verbosity',  [], 'fast');
DecompOut  = IRget(options, 'DecompOut',  [], 'fast');

restart = strcmp(restart, 'on');
verbose = strcmp(verbose, 'on');

% Setting K.
if isempty(K)
    K = MaxIter;
end
% Sorting the iterations (in case they are shuffled in input).
K = K(:); K = sort(K,'ascend'); K = unique(K);
if ~((isreal(K) && (all(K > 0)) && all(K == floor(K))))
    error('K must be a vector of positive real integers')
end
if K(end) ~= MaxIter
    MaxIter = K(end);  
end

StopIt = MaxIter;

if isempty(NoiseLevel) || strcmp(NoiseLevel,'none')
    Rtol = 0;
else
    Rtol = eta*NoiseLevel;
end

n = length(b);
test_sq = ones(n,1);
try
    test_sq = A_times_vec(A, test_sq);
    if (length(test_sq)~=n)
        error('The matrix A shuold be square; check the length of b.')
    end
catch
    error('The matrix A must be square; check the length of b.')
end

% See if an initial guess is given; if not then use 0 as initial guess.  
x0 = IRget(options, 'x0', [], 'fast');
Ab = A_times_vec(A, b);
nrmb = norm(b(:));
if strcmp(x0,'none')
    x0 = zeros(n,1);
    r = b;
    rr = Ab;
else
    try
        Ax = A_times_vec(A, x0);
        r = b - Ax;
        rr = Ab - Ax;
    catch
        error('Check the length of x')
    end
end
if restart
    ktotcount  = IRget(options, 'ktotcount', [],'fast');
    TotIterMax = IRget(options, 'TotIterMax',[],'fast');
    if strcmp(TotIterMax, 'none') || TotIterMax < MaxIter
        TotIterMax = MaxIter;
    end
    if strcmp(ktotcount, 'none')
        error('The total iteration counter must be assigned')
    end
    Ktot = IRget(options, 'Ktot', [], 'fast');
    % No checks on Ktot, it should be given from IRrestart.
end

% Declare matrices and prepare for the iterations.
X       = zeros(n,length(K));
Xnrm    = zeros(MaxIter,1);
Rnrm    = zeros(MaxIter,1);
V       = zeros(n,MaxIter); % Orthonormal vectors spanning Krylov subspace
h       = zeros(MaxIter,1); % New column of Hessenberg matrix H
Q       = zeros(MaxIter+1); % H = Q*T, Q orthogonal
T       = zeros(MaxIter);   % H = Q*T, T upper triangular
invT    = zeros(MaxIter);   % Inverse of the matrix T
coeff   = zeros(MaxIter,1); % Solution of the projected problem, i.e.,
                            % coefficient w.r.t. the Krylov basis.
rhs = zeros(MaxIter+1,1);   % Projected right-hand side

if strcmp(x_true,'none')
    errornorms = false;
else
    errornorms = true;
    Enrm = zeros(max(K),1);
    nrmtrue = norm(x_true(:));
    BestReg.It = [];
    BestReg.X = [];
    BestReg.Enrm = [];
    BestEnrm = 1e10;
    BestReg.Xnrm = [];
    BestReg.Rnrm = [];
end

NoStop = strcmp(NoStop,'on');

if restart
    saved_iterations = zeros(1, length(Ktot));
else
    saved_iterations = zeros(1, length(K));
end

% Prepare for the iterations.
alpharr  = norm(rr);
V(:,1) = rr/alpharr;
Q(1,1) = 1;
rhs(1) = V(:,1)'*r;
y = 0;

% Iterate.
noIterBar = strcmp(IterBar,{'off'});
if ~noIterBar
  h_wait = waitbar(0, 'Running iterations, please wait ...');
end
j = 0;
for k=1:MaxIter
    if restart, ktotcount = ktotcount + 1; end
    if ~noIterBar
        waitbar(k/MaxIter, h_wait)
    end 
    w = A_times_vec(A, V(:,k));
    % Modified Gram-Schmidt on the new vector.
    for ll=1:k
        h(ll) = V(:,ll)'*w;
        w = w - V(:,ll)*h(ll);
    end
    alpha = norm(w);
    % Store new Arnoldi vector and update projected rhs.
    V(:,k+1) = w/alpha;
    rhs(k+1) = V(:,k+1)'*r;
    beta = rhs(1:k+1);
    % Apply previous rotations to h.
    T(1:k,k) = Q(1:k,1:k)'*h(1:k);
    % Compute Givens rotation parameters.
    rc = T(k,k);
    if alpha == 0
        c = 1; s = 0;
    elseif abs(alpha) > abs(rc)
        tau = -rc/alpha;
        s = 1 / sqrt(1 + abs(tau)^2);
        c = s*tau;
    else
        tau = -alpha/rc;
        c = 1 / sqrt(1 + abs(tau)^2);
        s = c*tau;
    end
    % Apply givens rotations.
    T(k,k) = c'*rc - s'*alpha;
    Q(1:k,[k,k+1]) = Q(1:k,k)*[c s];
    Q(k+1,[k,k+1]) = [-s c];    
    if abs(T(k,k)) <= eps
        disp('Hessenberg matrix is (numerically) singular')
        X(:,j+1) = x;
        X = X(:,1:j+1);
        V = V(:,1:k+1);
        if restart
            saved_iterations(j+1) = ktotcount-1;
        else
            saved_iterations(j+1) = k-1;
        end
        saved_iterations = saved_iterations(1:j+1);
        % if k-1 < StopIt, StopIt = k-1; end
        if k>1
            Xnrm = Xnrm(1:k-1);
            Rnrm = Rnrm(1:k-1);
            if errornorms, Enrm = Enrm(1:k-1); end
        end
        % Stop because the Hessenberg matrix is (numerically) rank def.
        if StopIt == MaxIter
            StopFlag = 'Breakdown of the Arnoldi algorithm';
            StopReg.It = k-1;
            StopReg.X = x; 
            if errornorms, StopReg.Enrm = Enrm(k-1); end
        end
        k = k-1;
        break
    end
    % Update the inverse of T.
    invT(1:k-1,k) = -(invT(1:k-1,1:k-1)*T(1:k-1,k))/T(k,k);
    invT(k,k) = 1/T(k,k);
    % Update the previous projected solution.
    coeff(1:k-1) = y;
    y = coeff(1:k) + (Q(1:k+1,k)'*beta)*invT(1:k,k);    
    % Update solution.
    x = x0 + V(:,1:k)*y;
    % Compute norms.
    Xnrm(k)    = norm(x(:));
    if k==1
        Rnrm(k) = sqrt( abs( r'*r - abs(Q(1:k+1,k)'*beta)^2 ) )/nrmb;
    else
        Rnrm(k) = sqrt( abs( (Rnrm(k-1)*nrmb)^2 - abs(Q(1:k+1,k)'*beta)^2 ) )/nrmb;
    end
    if errornorms
        Enrm(k) = norm(x_true(:)-x(:))/nrmtrue;
        if Enrm(k)<BestEnrm
            BestReg.It = k;
            BestReg.X = x;
            BestEnrm = Enrm(k);
            BestReg.Enrm = BestEnrm;
            BestReg.Xnrm = Xnrm(k);
            BestReg.Rnrm = Rnrm(k);
        end
    end
    AlreadySaved = 0;
    if any(k==K)
        j = j+1;
        X(:,j) = x;
        saved_iterations(j) = k;
        if restart, saved_iterations(j) = ktotcount; end
        AlreadySaved = 1; 
    end 
    if restart
        if any(ktotcount == Ktot) && ~ AlreadySaved
            j = j+1;
            X(:,j) = x;
            saved_iterations(j) = ktotcount;
            AlreadySaved = 1;                
        end
        if ktotcount == TotIterMax
            if ~ AlreadySaved
                j = j+1;
                saved_iterations(j) = ktotcount;
                X(:,j) = x;
                AlreadySaved = 1;
            end
            StopIt = k;
            StopReg.It = k;
            StopReg.X = x;  
            if errornorms
                Enrm = Enrm(1:k);
                StopReg.Enrm = Enrm(k);
            end
            Xnrm     = Xnrm(1:k);
            Rnrm     = Rnrm(1:k);
            X = X(:,1:j);
            saved_iterations = saved_iterations(1:j);
            if verbose
                disp('Reached maximum number of iterations')
            end
            StopFlag = 'Reached maximum number of iterations';
            break
        end
    end
    if (Rnrm(k) <= Rtol) && (StopIt == MaxIter)
        if verbose
            disp('Residual tolerance satisfied')
        end
        StopFlag = 'Residual tolerance satisfied';
        if ~AlreadySaved && ~NoStop
            j = j+1;
            X(:,j) = x;
            if restart
                saved_iterations(j) = ktotcount;
            else
                saved_iterations(j) = k;
            end
            AlreadySaved = 1;
        end
        StopIt = k;
        StopReg.It = k;
        StopReg.X = x;
        if errornorms, StopReg.Enrm = Enrm(k); end
        if ~ NoStop
            Xnrm    = Xnrm(1:k);
            Rnrm    = Rnrm(1:k);
            if errornorms, Enrm = Enrm(1:k); end
            X = X(:,1:j);
            saved_iterations = saved_iterations(1:j);
            break
        end
    end
end
if k == MaxIter
  if StopIt == MaxIter
        % Stop because max number of iterations reached.
        if verbose
            disp('Reached maximum number of iterations')
        end
        StopFlag = 'Reached maximum number of iterations';
        if ~AlreadySaved
            j = j+1;
            X(:,j) = x;
            if restart
                saved_iterations(j) = ktotcount;
            else
                saved_iterations(j) = k;
            end
        end
        StopReg.It = k;
        StopReg.X = x;
        if errornorms, StopReg.Enrm = Enrm(k); end
        Xnrm    = Xnrm(1:k);
        Rnrm    = Rnrm(1:k);
        if errornorms, Enrm = Enrm(1:k); end
        X = X(:,1:j);
        saved_iterations = saved_iterations(1:j);
  end 
end
if ~noIterBar, close(h_wait), end
if nargout==2
  info.its = k;
  if restart
      info.ktotcount = ktotcount;
  end
  info.saved_iterations = saved_iterations(1:j);
  info.StopFlag = StopFlag;
  info.StopReg = StopReg;
  info.Rnrm = Rnrm(1:k);
  info.Xnrm = Xnrm(1:k);
  if errornorms
    info.Enrm = Enrm(1:k);
    info.BestReg = BestReg;
  end
  if strcmp(DecompOut,'on')
    info.V = V;
    info.Q = Q;
    info.T = T;
  end
end