% %============================================================================%
% % Duke University                                                            %
% % K. P. Trofatter                                                            %
% % kpt2@duke.edu                                                              %
% %============================================================================%
% zgmres() - complex general minimum residual (GMRES)
% 
% USAGE:
%   [x, rr] = zgmres(A, b, x0, tol=1.0e-3, imax=numel(b))
%
% INPUT:
%   [n,n] complex  | A    | conjugate-symmetric positive-definite matrix
%   [1,1] function | A    | conjugate-symmetric positive-definite function
%   [n,1] complex  | b    | data vector
%   [n,1] complex  | x0   | initial guess vector
%   [1,1] double   | tol  | tolerance of residual magnitude squared
%   [1,1] double   | imax | maximum number of iterations
%
% OUTPUT:
%   [n,1] complex  | x    | approximate solution vector
%   [1,1] double   | err  | residual norm


function [x, err] = zgmres(A, b, x0, tol, imax)
    
    % count dimension
    n = numel(b);
    
    % default arguments
    if ~exist('x0', 'var') || isempty(x0)
        x0 = zeros(n, 1);
    end
    if ~exist('tol', 'var') || isempty(tol)
        tol = 1.0e-3;
    end
    if ~exist('imax', 'var') || isempty(imax)
        imax = 10;
    end
    
    % build matrix multiplication function
    if isa(A, 'function_handle')
        Af = A;
    else
        Af = @(x) A * x;
    end
    
    % set max number of iterations
    % theoretically gmres converges after n iterations
    imax = min(imax, n);
    
    % allocate memory
    D = zeros(n, imax);
    R = zeros(imax + 1, imax); % the R in QR. this doubles as a temporary H
    c = zeros(1, imax); % cos rotation coefficients
    s = zeros(1, imax); % sin rotation coefficients
    g = zeros(imax + 1, 1); % intermediate vector
    err = zeros(1, imax + 1); % norm of residual
    
    % initiate gmres
    x = x0;
    r = b - Af(x);
    g(1) = norm(r);
    D(:, 1) = r / g(1);
    err(1) = g(1);
    
    % iterate
    for i = 1 : imax
        
        % arnoldi process
        d = Af(D(:, i));
        for j = 1 : i
            R(j, i) = D(:, j)' * d;
            d = d - D(:, j) * R(j, i);
        end
        R(i + 1, i) = norm(d);
        if i ~= imax
            D(:, i + 1) = d / R(i + 1, i);
        end
        
        % transform H column into R column
        for j = 1 : i - 1
            R(j : j + 1, i) = [c(j)', s(j); -s(j)', c(j)] * R(j : j + 1, i);
        end
        tau = sqrt(R(i, i)' * R(i, i) + R(i + 1, i)' * R(i + 1, i));
        c(i) = R(i, i) / tau;
        s(i) = R(i + 1, i)' / tau;
        R(i : i + 1, i) = [tau; 0.0];
        
        % transform g
        g(i : i + 1) = [c(i)', s(i); -s(i)', c(i)] * g(i : i + 1);
        
        % tolerance test
        err(i + 1) = abs(g(i + 1));
        if err(i + 1) <= tol
            break
        end
        
    end
    
    % clip
    D = D(:, 1 : i);
    R = R(1 : i, 1 : i);
    g = g(1 : i);
    err = err(1 : i + 1);
    
    % compute solution vector
    alpha = R \ g;
    x = x0 + D * alpha;
    
end


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