function [x,Arrnorm] = irlsqr(A,b,x,OPTS)
% IRLSQR is an implictly restarted LSQR method. The solution is computed via
% an implictly restarted Golub-Kahan bidiagonalization method with
% harmonic Ritz values as shifts. The program then runs LSQR on the updated
% system. The IRLSQR method solves the least squares problem: 
%
%                  min || b - A*x || 
%                   x
%
% where A is an (m x n) matrix, b an (m x 1) vector, and x is an (n x 1) 
% vector. b is not assumed to be in span(Col(A)). ||.|| always denotes
% 2 norm.
%-----------------------------------------------------------------------

% [x,Arrnorm] = IRLSQR(A,b,x,OPTS)
%
% The first input argument into ... = irlsqr(A,b,x) is a numeric 
% matrix A. The second input is the right hand side vector b and third
% input is the vector x, an initial approximation to min || b - A*x ||,
% typically chosen to be 0.

%-----------------------------------------------------------------------
%
% OUTPUT OPTIONS:
% ---------------
%
% I.)   x = IRLSQR(A,b,x,OPTS) 
%       Returns only an approximate solution vector x of min || b - A*x ||.
%
% II.) [x,Arrnorm] = IRLSQR(A,b,x,OPTS)
%            x - An approximate solution of min || b - A*x ||.
%
%      Arrnorm - Three dimensional array such that 
%                Arrnorm(j,1) = ||A'*(b - A*x_j)|| = ||A'*r_j||
%                Arrnorm(j,2) = ||b - A*x_j|| = ||r_j||
%                and Arrnorm(j,3) is the total number of matrix-vector 
%                products with A and A' required to compute the solution x(j).
%
%-------------------------------------------------------------------------
%
% INPUT OPTIONS:
% --------------
%                                   
%       [x,Arrnorm] = IRLSQR(A,b,x,OPTS)
%       OPTS is a structure containing input parameters. The input 
%       parameters can be given in any order. The structure OPTS may 
%       contain some or all of the following input parameters. The string 
%       for the input parameters can contain upper or lower case characters.
%       irlsqr(A,b,x) without the structure OPTS will use all of the default
%       values for OPTS.
%
%  INPUT PARAMETER      DESCRIPTION    
%
%
%  OPTS.LTOL         Tolerance used for stopping rule.
%                    ||A'*r|| <= LTOL*||A'*r_0||
%                    DEFAULT VALUE    LTOL = 1d-10 
%
%  OPTS.MAXITL       Maximum number of restarts for IRLSQR.
%                    DEFAULT VALUE   MAXITL = min(m,n). 
%
%  OPTS.M_B          Size of the GK bidiagonal matrix. The program may 
%                    increase M_B to ensure certain requirements are 
%                    satisfied. A warning message will be displayed if M_B 
%                    increased.                          
%                    DEFAULT VALUE   M_B = 90
%
%  OPTS.REORTH012    Perform 0, 1, or 2-sided reorthogonalization within
%                    GK vectors.
%                    DEFAULT VALUE   REORTH012 = 1
%
%  OPTS.P            Number of shifts to apply.
%                    DEFAULT VALUE   P = 20
%
%  OPTS.GAP_VAL      apply p \pm GAP_VAL number of shifts based on largest
%                    gap between harmonic Ritz values.
%                    DEFAULT VALUE   GAP_VAL = 0
%
%  OPTS.BIDIAG_RED   String for returning the matrix B to bidiagonal form
%                    after the application of the shifts. A value of 0 does
%                    not return B to bidiagonal form, and a value of 1
%                    does.
%                    DEFAULT VALUE   GAP = 0
%
%-------------------------------------------------------------------------

%  DATE:  10/1/13
%  VER:  1.0
%
%  AUTHORS:
%  James Baglama     
%  University of Rhode Island 
%  EMAIL: jbaglama@math.uri.edu
% 
%  Daniel Richmond
%  University of Rhode Island 
%  EMAIL: dan@math.uri.edu 

% REFERENCES:
% 1.) "Implictly Restarting the LSQR Algorithm", 
%      J. Baglama and D. Richmond, ETNA, submitted.

%----------------------------%
% BEGIN: PARSE INPUT VALUES. %
%----------------------------%

% Wrong number of input or output values.
if (nargin < 3 || nargin > 4 || nargout < 1 || nargout > 3 )
    error('irlsqr: Incorrect number of input or output arguments'); 
end

% Matrix A must be input as numeric matrix.
if isa(A,'numeric')
    [m,n] = size(A);
else
   error('irlsqr: Unknown type for matrix A.');
end

%-------------------------------%
% BEGIN: PARSE INPUT VARIABLES. %
%-------------------------------%
    
% Set all input options to default values.
 ltol = 1d-10; maxitl = min(m,n); m_b = 100;
 reorth012 = 1; p = 20; gap_val = 0; bidiag_red = 0;

if nargin == 4
   if ~isa(OPTS,'struct')
        error('irlsqr:OptionsNotStructure',...
        'Options argument must be a structure.')
   end
   names = fieldnames(OPTS);
   I = strcmpi('LTOL',names); I=find(I>0);
   if ~isempty(I), ltol = OPTS.(names{I}); end
   I = strcmpi('MAXITL',names); I=find(I>0);
   if ~isempty(I), maxitl = OPTS.(names{I}); end
   I = strcmpi('M_B',names); I=find(I>0);
   if ~isempty(I), m_b = OPTS.(names{I}); end
   I = strcmpi('REORTH012',names); I=find(I>0);
   if ~isempty(I), reorth012 = OPTS.(names{I}); end 
   I = strcmpi('P',names); I=find(I>0);
   if ~isempty(I), p = OPTS.(names{I}); end
   I = strcmpi('GAP_VAL',names); I=find(I>0);
   if ~isempty(I), gap_val = OPTS.(names{I}); end
   I = strcmpi('BIDIAG_RED',names); I=find(I>0);
   if ~isempty(I), bidiag_red = OPTS.(names{I}); end

end

%------------------------------------------------------%
% BEGIN: DESCRIPTION AND INITIALIZATION OF VARIABLES.  %
%------------------------------------------------------%

rlsqr = [];            % Residual of LSQR
alpha = [];            % alpha_{m+1} in the GK decomposition.
B = [];                % Bidiagonal matrix from GK decomposition.    
f = [];                % Initial right hand side. Small vector, maximum size m_b.
S_B = [];              % Singular values of B.
U_B = [];              % Left singular vectors of B.
V_B = [];              % Right singular vectors of B.
Q(:,1) = b;            % Starting vector
normAr1 = norm(A'*b);  % Used in plotting relative residual
normr1 = norm(b);      % Used in plotting relative residual
nrcount = 0;           % Variable used to keep track of current iteration
w = [];                % Work vector for LSQR solution procedure
f1 = [];               % Used in computation of vector rlsqr
qres = 0;              % q_{m_b+1,k+1} element of Q_L used in restart vector
iter = 1;              % counter for main iteration loop
flag = 0;              % used for when method has converged
ptemp = 0;             % holds the number of shifts after gap strategy 

%----------------------------------------------------%
% END: DESCRIPTION AND INITIALIZATION OF VARIABLES.  %
%----------------------------------------------------%

%----------------------------------------------------%
% BEGIN: MAIN METHOD                                 %
%----------------------------------------------------%

% Starting iteration, set all initial values
beta = norm(Q(:,1)); % Compute beta_1 (not used in bidiagonal matrix)
f1 = zeros(m_b+1,1);
f(1,1) = beta;
f1(1,1) = beta;
Q(:,1) = Q(:,1)/beta; 

P(:,1) = A'*Q(:,1); % Compute first vector in second Krylov space
alpha = norm(P(:,1)); % First entry in Bidiagonal matrix B
P(:,1) = P(:,1)/alpha; 
B(1,1) = alpha;
Arrnorm(1,1) = alpha*beta/normAr1; %||A'r_0|| = ||beta*A'*Q(:,1)||
Arrnorm(1,2) = beta/normr1;   %||r_0||
Arrnorm(1,3) = 1;    
nrcount = 1;

% Initial call to GK bidiag with k=0;
    [alpha,Q,P,B,nrcount,x,normr1,normAr1,Arrnorm,flag] = gkbd(A,...
        Q,P,B,0,m_b,nrcount,f,x,normr1,normAr1,Arrnorm,ltol,flag,reorth012,m,n);

while (iter <= maxitl && ~flag)
 
    Pres = P(:,m_b+1); % Used for restart vector
    P = P(:,1:m_b); % Re-size matrix P
    
    % The following GK equations must be satisfied on exit:
    %   A * P(:,1:m_b)     =   Q(:,1:m_b+1)*B(1:m_b+1,1:m_b)
    %   A'* Q(:,1:m_b+1)   =   P(:,1:m_b)*B(1:k+1,1:k)'+alpha*P(:,m_b+1)*flipud(eye(m_b+1,1))

    % Compute the residual vector, rlsqr
    [q,~]=qr(B);      
    phi_bar = flipud(eye(m_b+1,1))'*q'*f1;
    rlsqr = Q*phi_bar*q*flipud(eye(m_b+1,1)); %residual of LSQR

    % Compute SVD of B to get updated U_B, S_B, and V_B matrices
    [U_B,S_B,V_B] = svd(B);
 
    % Implement harmonic gap strategy
    gap_old = 0;
    ptemp = p;
    for i = (p-gap_val):(p+gap_val)
       gap = abs(S_B(i,i)-S_B(i+1,i+1));
        if gap > gap_old
            gap_old = gap; % Replace with the larger gap
            p=i; % Update the number of shifts to apply
        end
    end
    k = m_b-p;         
    
    % Application of the hRv's as shifts
    [Q_L, B, Q_R] = hbsQR(B,U_B,V_B,m_b,p); 
       
    % Update P, Q, B matrices and truncate
    Q = Q*Q_L(:,1:k+1);  
    P = P*Q_R(:,1:k);  
    B = B(1:k+1,1:k); 

    % B may not be a bidiagonal matrix, use rowwise HH to get back
    % bidiagonal structure
    if bidiag_red == 1
        [B, Q_LB, Q_RB] = rowHH(B);
        % Update P, Q, and qres
        P=P*Q_RB;
        Q=Q*Q_LB';
        qres = Q_L(m_b+1,k+1)*Q_LB(k+1,k+1);
    else
       % Computation of Q_L_{m_b+1,k+1} element for restart equations
       qres = Q_L(m_b+1,k+1);
    end
 
    % The following GK equations must be satisfied after 'p' shifts applied
    %  A * P(:,1:k)    =   Q(:,1:k+1)*B(1:k+1,1:k)
    %  A'* Q(:,1:k+1)  =   P(:,1:k)*B(1:k+1,1:k)'+qres*alpha*P(:,m_b+1)*flipud(eye(k+1,1))'
        
    % Compute restart vector and next element of B matrix
        rk = alpha*qres*Pres; % alpha = ||p(:,m_b+1)||
        rknorm = norm(rk);
        P = [P(:,1:k) rk/rknorm]; 
        B(k+1,k+1) = rknorm;       

    % Create f vector needed for restart for solving for LSQR solution
    f = Q'*rlsqr; % rlsqr must be in col(Q) to restart correctly
    
    % Need to update f1 for computation of rlsqr after restart
    f1 = [f; zeros(p,1)];
    
    % Increase iteration counter for print out of iteration number
    iter=iter+1 
    p = ptemp; % Reset value of p to pre-gap value
    
    % Call GK bidiagonalization to build up space to size m_b
    [alpha, Q,P,B,nrcount,x,normr1,normAr1,Arrnorm,flag] = gkbd(A,...
        Q,P,B,k,m_b,nrcount,f,x,normr1,normAr1,Arrnorm,ltol,flag,reorth012,m,n);
end

%----------------------------------------------------%
% END: MAIN METHOD                                   %
%----------------------------------------------------%

%----------------------------------------------------%
% BEGIN: GK BIDIAGONALIZATION / LSQR                 %
%----------------------------------------------------%
function [alpha,Q,P,B,nrcount,x,normr1,normAr1,Arrnorm,flag] = gkbd(A,...
    Q,P,B,k,m_b,nrcount,f,x,normr1,normAr1,Arrnorm,ltol,flag,reorth012,m,n)

alpha = B(k+1,k+1); 

% Compute next vector in Q
Q(:,k+2) = A*P(:,k+1)- Q(:,k+1)*alpha;

% Reorthogonalization of Q vectors
if (reorth012 == 2 || (reorth012 == 1 && m < n ))
    for j=1:k+1
        Q(:,k+2) = Q(:,k+2) - Q(:,j)*(Q(:,j)' * Q(:,k+2));
    end
end

% Computes beta_{k+2} of the B(k+2,k+1) matrix
beta = norm(Q(:,k+2)); 
Q(:,k+2) = Q(:,k+2)/beta; 
B(k+2,k+1) = beta;

% Compute A'*Q(:,k+2) - beta_{k+2}*P(:,k+1)
P(:,k+2) = (Q(:,k+2)'*A)' - P(:,k+1)*beta;

% Reorthogonalization of P vectors
if (reorth012 == 2 || (reorth012 == 1 && m >= n ))
    for j=1:k+1
        P(:,k+2) = P(:,k+2) - P(:,j)*(P(:,j)'*P(:,k+2));
    end
end

% B(k+2,k+2) element.
alpha = norm(P(:,k+2));  
if k~=m_b-1
    B(k+2,k+2) = alpha; %needed for applying only 1 shift
end
P(:,k+2) = P(:,k+2)/alpha;

% Update vector f
f(k+2,1) = 0; 

% Use MATLAB's internal qr. This gives +- values on the diagonal 
[Q_B,R_B] = qr(B(1:k+2,1:k+1));
 
% Setup of transition values for LSQR
rho_bar = alpha*Q_B(k+2,k+2);
theta =   alpha*Q_B(k+2,k+1); 
    
% Solve ||f(k+2,1) - B(k+2,k+1)*y||_2
% Replace B(k+2,k+1) with Q(1:k+2,1:k+2)*R(1:k+2,1:k)
% Replace f(k+2,1) with Q'*f
f = Q_B'*f; 
phi_bar = f(k+2,1); % Intermediate right hand side value needed for LSQR

% Compute the solution y of ||f(k+2,1) - B(k+2,k+1)*y||_2
y = R_B(1:k+1,1:k+1)\f(1:k+1);
  
% Update the solution vector
x = x + P(:,1:k+1)*y;

% Check the stopping criteria/compute residual norms
nrcount = nrcount + 1;
Arrnorm(nrcount,1) = abs(alpha*beta*y(k+1))/normAr1; 
Arrnorm(nrcount,2) = abs(phi_bar)/normr1;
Arrnorm(nrcount,3) = Arrnorm(nrcount-1,3) + 2;
if Arrnorm(nrcount,1) <= ltol 
    Arrnorm = Arrnorm(1:nrcount,1:3); flag=1;
    return;
end
     
% Update w vector
w = P(:,k+2) - P(:,1:k+1)*(y/f(k+1,1)*theta);

for i = k+2:m_b %start with the second iteration
    Q(:,i+1) = A*P(:,i) - alpha*Q(:,i); 
   
    % Reorthogonalization of Q vectors
    if ((reorth012 == 2) || (reorth012 == 1 && m < n ))
        for j=1:i
            Q(:,i+1) = Q(:,i+1) - Q(:,j)*(Q(:,j)' * Q(:,i+1));
        end
    end
    
    % Normalize the q_i vector
    beta = norm(Q(:,i+1));
    Q(:,i+1) = Q(:,i+1)/beta;
    
    % Update the B matrix
    B(i+1,i) = beta;
    
    % Continue the bidiagonalization
    P(:,i+1) = A'*Q(:,i+1) - beta*P(:,i);
        
    % Reorthogonalization of P vectors
    if ((reorth012 == 2) || (reorth012 == 1 && m >= n ))
        for j=1:i
            P(:,i+1) = P(:,i+1) - P(:,j)*(P(:,j)' * P(:,i+1));
        end
    end
    
    % Update the B matrix
    alpha = norm(P(:,i+1));
    P(:,i+1) = P(:,i+1)/alpha;
    B(i+1,i+1) = alpha;
    if i==m_b
        B = B(1:i+1,1:i);
    end
    
    % Givens values. 
    rho = pythag(rho_bar,beta); c1 = rho_bar/rho; s1 = beta/rho; 
    
    % Translate B from lower bidiagonal to upper bidiagonal matrix R.
    % Super diagonal element in R is theta_{i} and is computed
    % from applying Givens to B. The diagonal value of R is rho_{i}
    % and is given above as two norm of [rho_bar,beta].
    theta = s1*alpha; 
    
    % Updating the last diagonal element of R.
    rho_bar = -c1*alpha;
    
    % Updating elements in f (two elements up).
    phi = c1*phi_bar; phi_bar = s1*phi_bar;
    
    % Update the solution.
    x = x + (phi/rho)*w; w = P(:,i+1) - (theta/rho)*w;
    
    % Compute norms of residual
    nrcount = nrcount + 1;
    Arrnorm(nrcount,1) = abs(phi_bar*rho_bar)/normAr1; 
    Arrnorm(nrcount,2) = abs(phi_bar)/normr1;
    Arrnorm(nrcount,3) = Arrnorm(nrcount-1,3) + 2;
    if Arrnorm(nrcount,1) <= ltol
        Arrnorm = Arrnorm(1:nrcount,1:3); flag=1;
        return;
    end
end       
%----------------------------------------------------%
% END: GK BIDIAGONALIZATION / LSQR                   %
%----------------------------------------------------%

%----------------------------------------------------%
% BEGIN: HARMONIC BIDIAGONAL SHIFTED QR              %
%----------------------------------------------------%
function [Q_L, B_plus, Q_R] = hbsQR(B,U_B,V_B,m_b,p)
% This function computes the orthogonal matrices Q_L and Q_R such that B_plus =
% Q_L'*B*Q_R is the matrix obtained after p implicit shifts have been
% performed. B_plus is such that it is no longer bidiagonal, zero elements
% begin to fill in due to numerical error.
%
% INPUT:
%          B -  m_b+1 x m_b lower bidiagonal matrix from GK bidiagonalization
%        U_B -  matrix of left singular vectors of B
%        V_B -  matrix of right singular vectors of B
%        m_b -  size of the matrix B 
%          p -  number of shifts to be applied
%
% OUTPUT:
%        Q_L - (m_b+1) x (m_b+1) orthogonal upper Hessenberg matrix
%        B_plus - updated bidiagonal matrix
%        Q_R - m_b x m_b orthogonal upper Hessenberg matrix
%

% Initialize values
sL = m_b+1;
sR = m_b;
Q_L = zeros(m_b+1); Q_R = zeros(m_b); Is = p:-1:1;

% Take QR of the p+1 x p principal submatrix of U_B' to get first column of Q_L
[Q,~] = qr(U_B(1:p+1,1:p)); Q_L(1:p+1,1) = Q(:,p+1);

% Continue process to get first m_b+1-p columns of Q_L

for i = 2:sL-p
    [Q,~] = qr([Q_L(1:i+p,1:i-1) U_B(1:i+p,1:p)]);
    Q_L(1:i+p,i) = Q(:,p+i);
end

% Last p columns of Q_L are the first p columns of U_B (when ordered largest to smallest) 
 Q_L(:,sL-p+1:sL) = U_B(:,Is);

% Take QR of p+1 x p principal submatrix of V_B' to get first column of Q_R
[Q,~] = qr(V_B(1:p+1,1:p)); Q_R(1:p+1,1) = Q(:,p+1);

% Continue process to get first m_b-p columns of Q_R
for i = 2:sR-p
  [Q,~] = qr([Q_R(1:i+p,1:i-1) V_B(1:i+p,1:p)]);
  Q_R(1:i+p,i) = Q(:,i+p);
end  

% Last p columns of Q_R are the first p columns of V_B (when ordered largest
% to smallest)
Q_R(:,sR-p+1:sR) = V_B(:,Is);

% Update B
B_plus = Q_L'*B*Q_R;

%----------------------------------------------------%
% END: HARMONIC BIDIAGONAL SHIFTED QR                %
%----------------------------------------------------%

%----------------------------------------------------%
% BEGIN: RETURN B TO BIDIAGONAL FORM                 %
%----------------------------------------------------%
function [B,U,V] = rowHH(A)
% Produces lower bidiagonal matrix B
% such that B = U A V.
%
[m,n] = size(A);
B = A; QL=eye(m); QR=eye(n);
U=eye(m); V=eye(n);
for k = m:-1:3
    % Obtain Householder reflector to zero out rows of B,  b*G = alpha*e_k^T 
    [v,beta,~] = reflector(B(k,1:k-1)',2);
    QR(1:k-1,1:k-1)=eye(k-1)-beta*(v*v');
    V=V*QR; %update V matrix
    B=B*QR; %apply transformation on right to zero out rows
    
    % Obtain Householder reflector to zero out columns of B
    [v,beta,~] = reflector(B(1:k-1,k-1),2);
    QL(1:k-1,1:k-1) = eye(k-1)-beta*(v*v');
    U = QL*U; %update U matrix
    B = QL*B; %Apply transformation on left to zero out columns
    QR=eye(n); QL=eye(m);
end
%Here, U^T is created as just U, so B=UAV

%===========================================================

function [u,beta,alpha] = reflector(x,mode)
% Generate a reflector Q = eye - beta*u*u'
% such that Q*x = alpha*e_1 if mode == 1
% and       Q*x = alpha*e_n if mode == 2.

if mode == 2, m = size(x,1); else m = 1; end

% Rescale x so that x_1 is nonnegative and norm(x) = 1. 
scale = norm(x); 
if scale == 0
    u = x; beta = 0; alpha = 0;
else
    x = x/scale;
    if x(m) ~= 0
        phase = x(m)/abs(x(m));
        x = x*conj(phase); x(m) = real(x(m));
    else
        phase = 1;
    end

    % Build u and beta.
    u = x; u(m) = u(m) + 1;
    beta = 1/u(m);

    % Rescale
    alpha = -scale*phase;
end
%----------------------------------------------------%
% END: RETURN B TO BIDIAGONAL FORM                   %
%----------------------------------------------------%