function [xs,info] = L1SPIRiT_NLCG_NonCartesian(kData, x0, SPIRiTop, NUFFTop, lambda, wavWeight, tvWeight, varargin)
% L1SPIRiT using nonlinear CG + backtracking line search for comparison

% Solve least squares + l1 + TV problem:
%
% min 0.5*lambda*||(G-I)*x||^2 + 0.5*||NUFFTop*x - kData||^2
%  x    + wavWeight*||Psi*x||_{1,2,epsilon} + tvWeight*||TVop*x||_{1,2,epsilon}
%
% Gradient is:
%
% grad(x) = lambda*(G-I)'*(G-I)*x + NUFFTop'*(NUFFTop*x-kData)
%           + wavWeight*Psi'*D(Psi*x)*Psi*x + tvWeight*TVop'*D(TVop*x)*TVop*x
%
% where diagonal [D(W)]_n,: = (||W_{n,:}||^2 + epsilon)^(-1/2);
%
% Based on code found in sparseMRI (fnlCg.m) and SPIRiT (cgNUSPIRiT.m) toolboxes
%
% Copyright 2013-01-02, Daniel Weller, University of Michigan

% options
opts.maxiters = 10; % maximum # of iterations
opts.miniters = 1; % set to maxiters to force maximum # of iterations
opts.itersave = []; % if not empty, iterations to save current output in "xs"
opts.tsave = []; % if not empty, times to save current output in "xs"
opts.xtol = 1e-4; % relative tolerance in change in x
opts.ftol = 1e-4; % relative tolerance in objective
opts.dxtol = 1e-4; % scale of descent vector
opts.wavelet_args = {'Daubechies',4,4}; % wavelet type, type spec, and # of levels
opts.wav_epsilon = 1e-15; % corner rounding for l1 -- value from sparseMRI's init.m
opts.tv_isotropic = true; % toggles isotropic/anisotropic TV
opts.tv_epsilon = 1e-15; % corner rounding for TV -- value from sparseMRI's init.m
opts.lsmaxiters = 10; % maximum # of line search iterations
opts.lsalpha0 = 1; % initial alpha for line search
opts.lsbeta = 0.5; % geometric scale factor for line search
opts.lssecant = 0.3; % secant line scale for line search
opts.lsadaptive = true; % adapt alpha0 to # of line search iterations
opts.output = 0; % 0 = quiet, 1 = warnings + final only, 2 = iters, 3 = images too
opts.evalfun = []; % function handle to per-iteration evaluation

opts = vararg_pair(opts,varargin);

x = x0;

sz_images = size(x); sz_images(end+1:3) = 1;
sz_impad = 2^nextpow2(max(sz_images(1:2)));
sz_impad = [sz_impad,sz_impad,sz_images(3:end)];

if wavWeight > 0, DWTop = Wavelet(opts.wavelet_args{:}); else DWTop = []; end
if tvWeight > 0, TVop = TVOP(); else TVop = []; end

% initialize NLCG
SPIRiTx = []; SPIRiTdx = [];
NUFFTx = []; NUFFTdx = [];
DWTx = []; DWTdx = [];
TVx = []; TVdx = [];

alpha0 = opts.lsalpha0;

te = tic();
% pre-compute matrices for making line search fast
[NUFFTx,SPIRiTx,DWTx,TVx] = pre_compute(x);
gradx = gradfun();
dx = -gradx; % initial direction is negative gradient
[NUFFTdx,SPIRiTdx,DWTdx,TVdx] = pre_compute(dx);
drawnow;

f = ffun(0); % initial objective value
normdx = norm(col(dx));
fscale = -normdx^2; % better be negative

fs = [f;zeros(opts.maxiters,1)];
converged = zeros(opts.maxiters,1);
lsiters = zeros(opts.maxiters,1);
normdxs = [normdx;zeros(opts.maxiters,1)];
normxdiffs = zeros(opts.maxiters,1);
alphas = zeros(opts.maxiters,1);
bs = zeros(opts.maxiters,1);
if ~isempty(opts.evalfun) && isa(opts.evalfun,'function_handle')
    evaloutputs = opts.evalfun(x);
end
times = zeros(opts.maxiters,1);

if opts.output == 3 && usejava('awt')
    h = figure();
    imshow(sos(x),[]);
    title('Initial guess'); drawnow;
else
    h = [];
end

xs = struct('x',x,'t',0,'iter',0);

iter = 1;
while iter <= opts.maxiters
    titer = tic();
    
    % initialize line search
    fstart = f;
    alpha = alpha0;
    
    % perform line search
    lsiter = 1;
    while lsiter <= opts.lsmaxiters
        f = ffun(alpha);
        if f <= fstart + opts.lssecant*alpha*fscale
            break; % found a suitable alpha between 0 and alpha0
        end
        
        alpha = alpha*opts.lsbeta;
        lsiter = lsiter + 1;
    end
    if lsiter > opts.lsmaxiters % didn't find anything
        if opts.output >= 1, warn('Iter #%d: backtracking line search did not converge',iter); end
%         alpha = 0;
%         f = fstart;
    end
    alphas(iter) = alpha;
    normxdiff = abs(alpha)*normdx;
    normxdiffs(iter) = normxdiff;
    normx = norm(col(x));
	x = x + alpha.*dx; % new x
    lsiters(iter) = lsiter;
    fs(iter+1) = f;
    
    if opts.lsadaptive
        % adaptive code from sparseMRI's fnlCg.m
        if lsiter > 3
            alpha0 = alpha0*opts.lsbeta;
        elseif lsiter < 2
            alpha0 = alpha0/opts.lsbeta;
        end
    end
    
    % update gradient
    gradx_prev = gradx;
    [NUFFTx,SPIRiTx,DWTx,TVx] = pre_compute(x);
    gradx = gradfun();
    b = polak_ribiere(); % could also use Fletcher-Reeves, Hestenes-Stiefel, or other method here
    bs(iter) = b;
    dx = b.*dx - gradx;
    
    % verify that direction is towards descent
    fscale = real(sum(sum(sum(conj(gradx).*dx,3),2),1)); % better be negative
    if fscale >= 0 % moving away from minimum
        if opts.output >= 1, warn('Iter #%d: dx points in direction of slope >= 0',iter); end
        dx = -gradx; % try negative gradient direction instead
        fscale = -norm(col(gradx))^2;
    end
    
    [NUFFTdx,SPIRiTdx,DWTdx,TVdx] = pre_compute(dx);
    normdx = norm(col(dx));
    normdxs(iter+1) = normdx;
    
    if normxdiff <= opts.xtol*normx, converged(iter) = 1; end % x-convergence
    if fstart - f <= opts.ftol*fstart, converged(iter) = 2; end % objective convergence
    if normdx <= opts.dxtol*normdxs(1), converged(iter) = 3; end % gradient convergence
    
    times(iter) = toc(titer);
    
    if any(opts.itersave == iter) || any(opts.tsave > sum(times(1:iter-1)) & opts.tsave <= sum(times(1:iter)))
        xs(end+1) = struct('x',x,'t',sum(times(1:iter)),'iter',iter); %#ok<AGROW>
    end
    
    if ~isempty(opts.evalfun)
        evaloutputs(iter+1) = opts.evalfun(x);
    end
    
    if opts.output >= 2
        if normxdiff <= opts.xtol*normx, str_x = '<='; else str_x = '>'; end
        if fstart - f <= opts.ftol*fstart, str_f = '>='; else str_f = '<'; end
        if normdx <= opts.dxtol*normdxs(1), str_dx = '<='; else str_dx = '>'; end
        fprintf(1,'Iter #%d: ||xdiff|| = %g %s %g, f = %g %s %g, ||dx|| = %g %s %g',iter,normxdiff,str_x,opts.xtol*normx,f,str_f,(1-opts.ftol)*fstart,normdx,str_dx,opts.dxtol*normdxs(1));
        switch converged(iter)
            case 1
                fprintf(1,' (x convergence).\n');
            case 2
                fprintf(1,' (objective convergence).\n');
            case 3
                fprintf(1,' (gradient convergence).\n');
            otherwise
                fprintf(1,'.\n');
        end
        
        if ~isempty(h) && ishandle(h)
            set(0,'CurrentFigure',h);
            imshow(sos(x),[]);
            title(sprintf('Iter #%d',iter));
        end
        drawnow;
    end
    
    if iter >= opts.miniters && converged(iter) > 0
        break;
    end
    
    iter = iter + 1;
end
te = toc(te);

if iter > opts.maxiters, iter = opts.maxiters; end
info.te = te;
info.niters = iter; % number of iterations performed
info.converged = converged(1:iter);
info.fs = fs(1:iter+1);
info.normdxs = normdxs(1:iter+1);
info.lsiters = lsiters(1:iter);
info.normxdiffs = normxdiffs(1:iter);
info.alphas = alphas(1:iter);
info.bs = bs(1:iter);
info.times = times(1:iter);
if ~isempty(h) && ishandle(h), info.hfig = h; end
if ~isempty(opts.evalfun), info.evaloutputs = evaloutputs(1:iter+1); end
xs(1).x = x; xs(1).t = sum(times(1:iter)); xs(1).iter = iter;

if opts.output == 1
    if normxdiff <= opts.xtol*normx, str_x = '<='; else str_x = '>'; end
    if fstart - f <= opts.ftol*fstart, str_f = '>='; else str_f = '<'; end
    if normdx <= opts.dxtol*normdxs(1), str_dx = '<='; else str_dx = '>'; end
    fprintf(1,'Iter #%d: ||xdiff|| = %g %s %g, f = %g %s %g, ||dx|| = %g %s %g',iter,normxdiff,str_x,opts.xtol*normx,f,str_f,(1-opts.ftol)*fstart,normdx,str_dx,opts.dxtol*normdxs(1));
    switch converged(iter)
        case 1
            fprintf(1,' (x convergence).\n');
        case 2
            fprintf(1,' (objective convergence).\n');
        case 3
            fprintf(1,' (gradient convergence).\n');
        otherwise
            fprintf(1,'.\n');
    end
    drawnow;
end

    function [NUFFTval,SPIRiTval,DWTval,TVval] = pre_compute(x)
        NUFFTval = NUFFTop*x;
        SPIRiTval = SPIRiTop*x;
        if wavWeight > 0, DWTval = DWTop*zpad(x,sz_impad); else DWTval = []; end
        if tvWeight > 0, TVval = TVop*x; else TVval = []; end
    end

    function f = ffun(alpha)
        f = 0.5*lambda*norm(reshape(SPIRiTx + alpha.*SPIRiTdx,[],1))^2 ...
            + 0.5*norm(reshape(NUFFTx + alpha.*NUFFTdx - kData,[],1))^2;
        if wavWeight > 0
            f = f + wavWeight*sum(sum(sqrt(sum(abs(DWTx + alpha.*DWTdx).^2,3)),2),1);
        end
        if tvWeight > 0
            if opts.tv_isotropic
                f = f + tvWeight*sum(sum(sqrt(sum(sum(abs(TVx + alpha.*TVdx).^2,4),3)),2),1);
            else
                f = f + tvWeight*sum(sum(sum(sqrt(sum(abs(TVx + alpha.*TVdx).^2,3)),4),2),1);
            end
        end
    end

    function gradx = gradfun()
        gradx = lambda.*(SPIRiTop'*SPIRiTx) + NUFFTop'*(NUFFTx-kData);
        if wavWeight > 0
            DinvPsi = sqrt(sum(abs(DWTx).^2,3) + opts.wav_epsilon);
            gradx = gradx + wavWeight.*crop(DWTop'*(DWTx./repmat(DinvPsi,[1,1,size(DWTx,3)])),sz_images);
        end
        if tvWeight > 0
            if opts.tv_isotropic
                DinvTV = sqrt(sum(sum(abs(TVx).^2,4),3) + opts.tv_epsilon);
                gradx = gradx + tvWeight.*(TVop'*(TVx./repmat(DinvTV,[1,1,size(TVx,3),size(TVx,4)])));
            else
                DinvTV = sqrt(sum(abs(TVx).^2,3) + opts.tv_epsilon);
                gradx = gradx + tvWeight.*(TVop'*(TVx./repmat(DinvTV,[1,1,size(TVx,3),1])));
            end
        end
    end

    function b = polak_ribiere()
        gradx_diff = gradx-gradx_prev;
        gradx_prev_normsqd = sum(sum(sum(abs(gradx_prev).^2,3),2),1);
        gradx_gradx_diff_prod = real(sum(sum(sum(conj(gradx).*gradx_diff,3),2),1));
        b = max(0,gradx_gradx_diff_prod/(gradx_prev_normsqd + eps));
    end

end


