% %============================================================================%
% % Duke University                                                            %
% % K. P. Trofatter                                                            %
% % kpt2@duke.edu                                                              %
% %============================================================================%
% zcg() - complex conjugate gradient
% 
% USAGE:
%   [x, rr] = zcg(A, b, x0=zeros(size(b)), Pif=2(x)x, 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] function | Pif  | preconditioner inverse function
%   [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   | rr   | residual magnitude squared


function [x, rr] = zcg(A, b, x0, Pif, tol, imax)
    
    % count dimension
    n = numel(b);
    
    % default arguments
    if ~exist('x0', 'var') || isempty(x0)
        x0 = zeros(n);
    end
    if ~exist('Pif', 'var') || isempty(Pif)
        Pif = @(x) x;
    end
    if ~exist('tol', 'var') || isempty(tol)
        tol = 1.0e-3;
    end
    if ~exist('imax', 'var') || isempty(imax)
        imax = n;
    end
    
    % compute max number of iterations
    imax = min(imax, n);
    
    % build matrix multiplication function
    if isa(A, 'function_handle')
        Af = A;
    else
        Af = @(x) A * x;
    end
    
    % compute initial residual
    x = x0;
    r = b - Af(x);
    
    % iterate
    rz = zeros(1, imax);
    rr = zeros(1, imax + 1);
    for i = 1 : imax
        
        % tolerance test
        rr(i) = r' * r; % this adds an extra inner product, but that's ok
        if abs(rr(i)) <= tol ^ 2
            break
        end
        
        % compute search direction vector
        z = Pif(r);
        rz(i) = r' * z;
        if i == 1
            d = z;
        else
            beta = rz(i) / rz(i - 1);
            d = z + beta * d;
        end
        
        % update
        Ad = Af(d);
        alpha = rz(i) / (d' * Ad);
        x = x + alpha * d;
        r = r - alpha * Ad;
        
    end
    
    % handle last err
    if abs(rr(i)) <= tol ^ 2
        rr = rr(1 : i);
    else
        rr(i + 1) = r' * r;
    end
    rr = sqrt(rr);
    
end


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