function varargout = pre_conj_grad_old(x0, A, b, varargin)
%
% function [xs, info, resvec] = pre_conj_grad_old(x0, A, b, varargin)
%
% Compute (optionally) preconditioned conjugate gradient solution of
% ||Ax-b|| starting at initial point x0. Code based on pseudocode from
% Nocedal and Wright (2ed.).
%
% Inputs:
%   x0            initial estimate
%   A             *atrix or function handle to evaluate A*x
%   b             vector or function handle to evaluate b
%
% Options (name / value pairs):
%   maxiters      maximum # of iterations (default: min(dof, 20))
%   miniters      minimum # of iterations (default: 1)
%   nrestart      # of iterations between restarts (default: Inf (do not restart))
%   resvec        initial residual vector A*x0 - b (default: []) (useful when pcg is inside another algorithm)
%   xtol          relative convergence tolerance in change in x (default: 10^-4)
%   rtol          relative convergence tolerance in change in r (default: 10^-4)
%   dxtol         relative convergence tolerance in gradient step (default: 10^-4)
%   output        0 = quiet, 1 = warnings only, 2 = each iteration (default: 0)
%   outputfun     function to call when generating output instead of default
%   forcemaxiters if true, will always run for maxiters iterations (default: false)
%   S             *atrix S, for solving preconditioned system S^T*A*S*y = S^T*b, where x = S*y (default: [])
%   M             preconditioning *atrix M = S*S^T or function handle for computing M*y (default: [])
%
% Outputs:
%   xs            solution to minimization problem
%   info          structure containing minimization information with fields
%     .niters       # of iterations performed (beyond zeroth iteration)
%     .converged    0 = did not converge; 1 = convergence in x, 2 = convergence in r, 3 = stagnation in conjugate direction
%     .normrs       norm of residual r = A*x-b at each iteration
%     .normps       norm of descent direction p at each iteration
%     .alphas       step size alpha at each iteration
%     .betas        conjugate parameter beta at each iteration
%     .times        computation time for each step
%     .ttotal       total computation time for nlcg method
%   resvec        residual vector A*xs - b
% 
%  Copyright 2012-10-10, Daniel Weller, University of Michigan

if isequal(x0,'test'), pre_conj_grad_test(); return; end

t0 = tic();

use_real = isreal(x0); % supports complex, too
ndof = length(x0);
if ~use_real, ndof = 2*ndof; end % complex has twice as many dof

% set up options
opts.maxiters = min(ndof,20);
opts.nrestart = Inf;
opts.resvec = [];
opts.xtol = 1e-4;
opts.rtol = 1e-4;
opts.dxtol = 1e-4;
opts.output = 0;
opts.outputfun = @outputfun_default;
opts.miniters = 1;
opts.forcemaxiters = false;
opts.S = [];
opts.M = [];

opts = vararg_pair(opts,varargin);

if opts.forcemaxiters, opts.miniters = opts.maxiters; end

if isempty(opts.M)
    if ~isempty(opts.S) % preconditioner specified as S
        opts.M = @(x) opts.S*opts.S'*x;
    end
end

% initialization
if isa(b,'function_handle'), b = b(); end
x = x0;
if ~isempty(opts.resvec)
    r = opts.resvec;
    opts = rmfield(opts,'resvec'); % save memory
else
    r = evalA(A,x) - b;
end
if nargout < 2 && ~isempty(opts.outputfun), r0 = r; end
normr0 = norm(r);

if isempty(opts.M)
    p = -r;
    ryprod = normr0^2;
else
    y = evalM(opts.M,r);
    ryprod = real(r'*y);
    if ryprod < 0
        if opts.output >= 1, warn('Init: PCG preconditioner does not appear to be psd, ignoring'); end
        ryprod = normr0^2;
        opts.M = [];
        p = -r;
    else
        if nargout < 2 && ~isempty(opts.outputfun), y0 = y; end
        p = -y;
    end
end
normp0 = norm(p);

xs = x; niters_rmin = 0;
normrs = [normr0;zeros(opts.maxiters,1)];
normps = [normp0;zeros(opts.maxiters,1)];
normxdiffs = zeros(opts.maxiters,1);
alphas = zeros(opts.maxiters,1);
betas = zeros(opts.maxiters,1);
times = zeros(opts.maxiters,1);

niters = 1;
converged = 0;
while niters <= opts.maxiters
    te = tic();
    Ap = evalA(A,p);
    pApprod = real(p'*Ap);
    if pApprod <= 0
        if opts.output >= 1 && pApprod < 0, warn('Iter #%d: PCG system matrix does not appear to be psd',niters); end
        alphas(niters) = 0;
    else
        alphas(niters) = ryprod/pApprod;
    end
    norm_xprev = norm(x);
    norm_xdiff = abs(alphas(niters))*normps(niters);
    normxdiffs(niters) = norm_xdiff;
    if opts.output >= 2, xprev = x; end
    x = x + alphas(niters)*p;

    if opts.output >= 2, rprev = r; end
    if isfinite(opts.nrestart) && mod(niters,opts.nrestart) == 0 % restart; recompute residual
        if opts.output >= 1, fprintf(1,'Iter #%d: restarting.\n',niters); end
        r = evalA(A,x) - b;
    else
        r = r + alphas(niters)*Ap;
    end
    normrs(niters+1) = norm(r);
    if niters <= opts.miniters || normrs(niters+1) <= normrs(niters_rmin+1) % check about updating
        xs = x; niters_rmin = niters;
    end

    ryprodprev = ryprod;
    if isempty(opts.M)
        ryprod = normrs(niters+1)^2;
    else
        if opts.output >= 2, yprev = y; end
        y = evalM(opts.M,r);
        ryprod = real(r'*y);
        if ryprod < 0
            if opts.output >= 1, warn('Iter #%d: PCG preconditioner does not appear to be psd, ignoring',niters); end
            ryprodprev = normrs(niters)^2;
            ryprod = normrs(niters+1)^2;
            opts.M = [];
        end
    end
    
    if ryprod >= ryprodprev
        if ryprod > ryprodprev
            if opts.output >= 1, warn('Iter #%d: residual increased, resetting descent direction...',niters); end
        end
        betas(niters) = 0; % reset
    else
        betas(niters) = ryprod/ryprodprev;
    end
    
    if opts.output >= 2, pprev = p; end
    if isfinite(opts.nrestart) && mod(niters,opts.nrestart) == 0 % restart; use gradient instead
        if isempty(opts.M)
            p = -r;
        else
            p = -y;
        end
    else
        if isempty(opts.M)
            p = -r + betas(niters)*p;
        else
            p = -y + betas(niters)*p;
        end
        if real(r'*p) > 0
            if opts.output >= 1, warn('Iter #%d: direction not descent, resetting descent direction...',niters); end
            betas(niters) = 0; % reset
            if isempty(opts.M)
                p = -r;
            else
                p = -y;
            end
        end
    end
    normps(niters+1) = norm(p);
    
    % check for convergence
    if norm_xdiff <= opts.xtol*norm_xprev % in x
        converged = 1;
    elseif abs(normrs(niters)-normrs(niters+1)) <= opts.rtol*normrs(niters) % in residual r
        converged = 2;
    elseif normps(niters+1) <= opts.dxtol*normp0; % in conjugate direction p
        converged = 3;
    end
    times(niters) = toc(te);
    
    if converged && niters >= opts.miniters, break; end
    
    if opts.output >= 2 && ~isempty(opts.outputfun)
        if isempty(opts.M)
            opts.outputfun(niters,x,r,r,p,xprev,rprev,rprev,pprev,converged,times(niters));
        else
            opts.outputfun(niters,x,r,y,p,xprev,rprev,yprev,pprev,converged,times(niters));
        end
    end
    
    niters = niters + 1;
end

ttotal = toc(t0);
if nargout >= 1, varargout{1} = xs; end
if niters > opts.maxiters, niters = opts.maxiters; end
info.niters = niters_rmin;
info.converged = converged;
info.normrs = normrs(1:niters_rmin+1);
info.normps = normps(1:niters_rmin+1);
info.normxdiffs = normxdiffs(1:niters_rmin);
info.alphas = alphas(1:niters_rmin);
info.betas = betas(1:niters_rmin);
info.times = times(1:niters_rmin);
info.ttotal = ttotal;
if nargout >= 2
    varargout{2} = info;
else
    if ~isempty(opts.outputfun)
        if isempty(opts.M)
            opts.outputfun(niters,x,r,r,p,x0,r0,r0,-r0,converged,ttotal);
        else
            opts.outputfun(niters,x,r,y,p,x0,r0,y0,-y0,converged,ttotal);
        end
    end
end
if nargout >= 3
    varargout{3} = r;
end

end


%
% function out = evalA(A,in)
%
% Evaluate *atrix out = A*in, or function handle out = A(in).
%
function out = evalA(A,in)

if isa(A,'function_handle')
    out = A(in);
else
    out = A*in;
end

end


%
% function out = evalM(M,in)
%
% Evaluate *atrix out = M*in, or function handle out = M(in).
%
function out = evalM(M,in)

if isempty(M)
    out = in; % no preconditioner
else
    if isa(M,'function_handle')
        out = M(in);
    else
        out = M*in;
    end
end

end


%
% function outputfun_default(iter,xval,rval,yval,pval,xprev,rprev,yprev,pprev,converged,te)
%
% Default output function that prints info about the current (or final)
% iteration.
%
function outputfun_default(iter,xval,rval,yval,pval,xprev,rprev,yprev,pprev,converged,te) %#ok<INUSL>

switch converged
    case 1
        status = ' Done: x converged.';
    case 2
        status = ' Done: r converged.';
    case 3
        status = ' Done: p stagnated.';
    otherwise
        status = '';
end
fprintf(1,'Iter #%d: ||x-x_prev|| = %g, ||r|| = %g -> %g, ||p|| = %g -> %g, time = %g s.%s\n',iter,norm(xval-xprev),norm(rprev),norm(rval),norm(pprev),norm(pval),te,status);
drawnow;

end

%
% function pre_conj_grad_test()
%
% Perform test on an arbitrary pos. def. matrix with eigenvalues uniformly
% distributed.
%
function pre_conj_grad_test()

% generate random matrix A
rng('default');
N = 1000;
dmin = 1/N;
d = rand(N,1).*(1-dmin)+dmin;
D = diag(d);
% V is IDFT matrix
W_N = exp(-1j*2*pi/N);
ns = 0:N-1;
ks = 0:N-1;
V = W_N.^(-(ns.')*ks)./sqrt(N);
A = V*D*V';
x = complex(randn(N,1),randn(N,1));
b = A*x;
x0 = complex(zeros(N,1));

[xs,info] = pre_conj_grad(x0,A,b,'output',2,'forcemaxiters',true);

fprintf(1,'Solution: ||xs-x||/||x|| = %g, ||r|| = %g, in %g s.\n',norm(xs-x)/norm(x),info.normrs(end),info.ttotal);
drawnow;

[xs_pre,info_pre] = pre_conj_grad(x0,A,b,'output',2,'forcemaxiters',true,'M',V*diag(1./d)*V');

fprintf(1,'Solution: ||xs-x||/||x|| = %g, ||r|| = %g, in %g s.\n',norm(xs_pre-x)/norm(x),info_pre.normrs(end),info_pre.ttotal);
drawnow;

figure;
semilogy(0:info.niters,info.normrs,0:info_pre.niters,info_pre.normrs);
xlabel('Iteration'); ylabel('||r||'); title('Residual norm decay for uniformly distributed eigenvalues');
legend({'No pre.';'Exact pre.'});

end
