function [xs,info] = L1SPIRiT_SB_NonCartesian(kData, x0, SPIRiTop, NUFFTop, lambda, wavWeight, tvWeight, mu, varargin)

% solve the L1/TV regularized non-Cartesian SPIRiT reconstruction using
% Split-Bregman (ADMM with just the regularizer split(s)):
%
% min lambda/2*||(G-I)*x||^2 + 1/2*||NUFFTop*x-kData||^2
%  x    + wavWeight*||DWTop*x||_1,2 + tvWeight*TV(x)
%
% becomes
%
%    min    lambda/2*||(G-I)*x||^2 + 1/2*||NUFFTop*x-kData||^2
% x,w,gx,gy   + wavWeight*||w||_1,2 + tvWeight*||[gx,gy]||_1,2
%
% with w = DWTop*x, gx = Dx*x, gy = Dy*x
%
% with Split-Bregman, this becomes
% 
% x <- min lambda*||(G-I)*x||^2 + ||NUFFTop*x-kData||^2
%       x                             
%                 ||[DWTop]*x-[ w -  u_w]||^2
%            + mu*||[  Dx ]   [gx - u_gx]||
%                 ||[  Dy ]   [gy - u_gy]||
%
% w <- min wavWeight*||w||_1,2 + mu/2*||w - (DWTop*x + u_w)||^2
%       w
%
% [gx,gy] <- min tvWeight*||[gx,gy]||_1,2 + mu/2*||[gx] - [Dx*x + u_gx]||^2
%           gx,gy                                ||[gy]   [Dy*x + u_gy]||
%
% u_w <- u_w + (DWTop*x - w)
% u_gx <- u_gx + (Dx*x - gx)
% u_gy <- u_gy + (Dy*x - gy)
%
% (we do the x-update last)
% 
% Copyright 2013, Daniel Weller, University of Michigan. All rights
% reserved. Email wellerds@umich.edu.

opts.maxiters = 10;
opts.miniters = 1;
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.abs_tol = 1e-4;
opts.rel_tol = 1e-4;
opts.xtol = 1e-4;
opts.update_x_cg_opts = {}; % see options in pre_conj_grad.m
opts.update_x_precon = ''; % also accepts 'diag' for diagonal preconditioning
opts.wavelet_args = {'Daubechies',4,4}; % wavelet type, type spec, and # of levels
opts.tv_isotropic = true; % toggles isotropic/anisotropic TV
opts.output = 0; % 0 = quiet, 1 = summary + warnings only, 2 = iterations, 3 = display image
opts.test_convergence = false; % set to true to compute r's and s's
opts.evalfun = []; % function out = evalfun(x) to call at each iteration

opts = vararg_pair(opts,varargin);

te = tic();
tinit = tic();

x = x0;
sz_images = size(x); sz_images(end+1:3) = 1;
sz_data = size(kData); sz_data(end+1:3) = 1;

if wavWeight > 0
    sz_impad = 2^nextpow2(max(sz_images(1:2)));
    sz_impad = [sz_impad,sz_impad,sz_images(3:end)];
    
    DWTop = Wavelet(opts.wavelet_args{:});
    DWT_x = DWTop*zpad(x,sz_impad);
    w = DWT_x;
    u_w = 0;
else
    w = []; u_w = []; DWTop = [];
end
if tvWeight > 0
    TVop = TVOP(sz_images(1:2));
    TV_x = TVop*x;
    gxy = TV_x;
    u_gxy = 0;
else
    gxy = []; u_gxy = []; TVop = [];
end

res_x = []; % cg residuals

info.mu = mu;

if isempty(mu)
    fail('mu needs to be specified');
end

% other initialization
switch lower(opts.update_x_precon)
    case 'diag'
        % generate preconditioner for x-update

        % SPIRiT gram matrix piece
        
        % used for computing diagonal preconditioner
        x_precon_helper = permute(getK(SPIRiTop),[4,3,1,2]); % [Nc Nc N1 N2]
        x_precon_helper = col(num2cell(x_precon_helper,[1,2]));
        x_precon_helper = cellfun(@(x) x-eye(sz_images(3)),x_precon_helper,'UniformOutput',false);
        x_precon_helper = cellfun(@(x) sum(abs(x).^2,1),x_precon_helper,'UniformOutput',false); % [N1*N2 Nc]
        x_precon_helper = lambda.*col(cat(1,x_precon_helper{:})); % [N1*N2*Nc]
        
        % NUFFT approx'ed with DTFT for gram matrix, which is Toeplitz
        % diagonal of DTFT gram matrix is just # of samples
        x_precon_helper = x_precon_helper + prod(sz_data(1:2))./prod(sz_images(1:2));

        % regularizer pieces
        if wavWeight > 0
            x_precon_helper = x_precon_helper + mu; % assume orthonormal DWT
%             im0 = zeros(sz_images(1:2));
%             Nd_center = floor(sz_images(1:2)./2)+1;
%             Nd_center_cells = num2cell(Nd_center);
%             im0(Nd_center_cells{:}) = 1;
%             x_precon_helper = x_precon_helper + mu.*sum(col(abs(DWTop*zpad(im0,sz_impad(1:2))).^2)); % circulant approx
        end
        if tvWeight > 0
            im0 = zeros(sz_images(1:2));
            Nd_center = floor(sz_images(1:2)./2)+1;
            Nd_center_cells = num2cell(Nd_center);
            im0(Nd_center_cells{:}) = 1;
            x_precon_helper = x_precon_helper + mu.*sum(col(abs(TVop*im0).^2)); % circulant
        end
        
        x_precon_helper = 1./x_precon_helper; % invert
end

tinit = toc(tinit);

% compute objective function
obj = ffun();

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

% output storage
if opts.test_convergence
    rs = zeros(opts.maxiters,1);
    ss = zeros(opts.maxiters,1);
    converged = zeros(opts.maxiters,1);
end
normxdiffs = zeros(opts.maxiters,1);
objs = [obj;zeros(opts.maxiters,1)];
x_cg_info = cell(opts.maxiters,1);
if ~isempty(opts.evalfun) && isa(opts.evalfun,'function_handle')
    evaloutputs = opts.evalfun(x);
end
times = zeros(opts.maxiters,1);

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

iter = 1;
while iter <= opts.maxiters
    if opts.output == 1 % just update completion
        ticker(mfilename, iter, opts.maxiters);
    end
    
    titer = tic();
    
    % update w
    if wavWeight > 0
        w_prev = w;
        w = DWT_x + u_w;
        wabs = sqrt(sum(abs(w).^2,3));
        wabs = 1 - (wavWeight/mu)./max(wabs,wavWeight/mu);
        wabs(~isfinite(wabs)) = 0; % safeguard for when wabs_n = wavWeight/mu = 0
        w = w.*repmat(wabs,[1,1,size(w,3)]);
    end
    
    % update gxy
    if tvWeight > 0
        gxy_prev = gxy;
        gxy = TV_x + u_gxy;
        if opts.tv_isotropic
            gxy_abs = sqrt(sum(sum(abs(gxy).^2,4),3));
        else
            gxy_abs = sqrt(sum(abs(gxy).^2,3));
        end
        gxy_abs = 1 - (tvWeight/mu)./max(gxy_abs,tvWeight/mu);
        gxy_abs(~isfinite(gxy_abs)) = 0; % safeguard for when wabs_n = tvWeight/mu = 0
        if opts.tv_isotropic
            gxy = gxy.*repmat(gxy_abs,[1,1,size(gxy,3),size(gxy,4)]);
        else
            gxy = gxy.*repmat(gxy_abs,[1,1,size(gxy,3),1]);
        end
    end
    
    % update u's, r's, s's, thresholds
    if opts.test_convergence
        if wavWeight > 0
            A_w = DWT_x; B_w = -w; r_w = A_w + B_w; u_w = u_w + r_w;
            A_w = norm(col(A_w)); B_w = norm(col(B_w)); r_w = norm(col(r_w));
            s_w = mu*norm(col(crop(DWTop'*(w-w_prev),sz_images)));
            n_r_w = numel(u_w);
            Au_w = norm(col(crop(DWTop'*u_w,sz_images)));
            A = A_w; B = B_w; r = r_w; s = s_w; n_r = n_r_w; Au = Au_w; 
        else
            A = []; B = []; r = []; s = []; n_r = 0; Au = [];
        end
        
        if tvWeight > 0
            A_gxy = TV_x; B_gxy = -gxy; r_gxy = A_gxy + B_gxy; u_gxy = u_gxy + r_gxy;
            A_gxy = norm(col(A_gxy)); B_gxy = norm(col(B_gxy)); r_gxy = norm(col(r_gxy));
            s_gxy = mu*norm(col(TVop'*(gxy-gxy_prev)));
            n_r_gxy = numel(u_gxy);
            Au_gxy = norm(col(TVop'*u_gxy));
            A = [A,A_gxy]; B = [B,B_gxy]; r = [r,r_gxy]; s = [s,s_gxy]; n_r = n_r + n_r_gxy; Au = [Au,Au_gxy]; %#ok<AGROW>
        end
        
        A = norm(A); B = norm(B); r = norm(r); s = norm(s); Au = norm(Au);
        rs(iter) = r; ss(iter) = s;
    else
        if wavWeight > 0, u_w = u_w + (DWT_x - w); end
        if tvWeight > 0, u_gxy = u_gxy + (TV_x - gxy); end
    end
    
    % update x
    x_prev = x;
    b_x = update_x_cg_AHbfun();
    if ~isempty(res_x), res_x = res_x - b_x; end
    if isempty(opts.update_x_precon)
        x_cg_opts = {'resvec',res_x};
    else
        x_cg_opts = {'M',@update_x_cg_Mfun,'resvec',res_x};
    end
    [x,x_cg_info{iter},res_x] = pre_conj_grad(col(x),@update_x_cg_AHAfun,b_x,opts.update_x_cg_opts{:},x_cg_opts{:});
    x = reshape(x,sz_images);
    res_x = res_x + b_x;
    normxdiffs(iter) = norm(col(x-x_prev));
    
    % update intermediate variables
    if wavWeight > 0
        DWT_x = DWTop*zpad(x,sz_impad);
    end
    if tvWeight > 0
        TV_x = TVop*x;
    end

    if opts.test_convergence
        % residual tolerances from Boyd
        tol_r = sqrt(n_r)*opts.abs_tol + max([A,B])*opts.rel_tol;
        tol_s = sqrt(prod(sz_images))*opts.abs_tol + mu*Au*opts.rel_tol;

        if ~(wavWeight > 0 || tvWeight > 0) || (r <= tol_r && s <= tol_s)
            converged(iter) = 1; % primal and dual residual convergence

            if normxdiffs(iter) <= opts.xtol*norm(col(x_prev))
                converged(iter) = 2; % x-convergence as well
            end
        end
    end
    
    times(iter) = toc(titer);

    % compute updated objective
    obj = ffun();
    objs(iter+1) = obj;
    
    if ~isempty(opts.evalfun) && isa(opts.evalfun,'function_handle')
        evaloutputs(iter+1) = opts.evalfun(x);
    end
    
    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
    
    % output
    if opts.output >= 2
        fprintf(1,'Iter #%d: ',iter);
        if opts.test_convergence
            if wavWeight > 0 || tvWeight > 0
                if r <= tol_r, str_r = '<='; else str_r = '>'; end
                if s <= tol_s, str_s = '<='; else str_s = '>'; end
                fprintf(1,'r = %g %s %g, s = %g %s %g, ',r,str_r,tol_r,s,str_s,tol_s);
            end
        end
        if normxdiffs(iter) <= opts.xtol*norm(col(x_prev)), str_x = '<='; else str_x = '>'; end
        fprintf(1,'dx = %g %s %g',normxdiffs(iter),str_x,opts.xtol*norm(col(x_prev)));

        if opts.test_convergence
            switch converged(iter)
                case 1
                    fprintf(1,' (residual convergence).\n');
                case 2
                    fprintf(1,' (x convergence).\n');
                otherwise
                    fprintf(1,'.\n');
            end
        else
            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 && opts.test_convergence && 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;
if opts.test_convergence
    info.rs = rs(1:iter);
    info.ss = ss(1:iter);
    info.converged = converged(1:iter);
end
info.objs = objs(1:iter+1);
info.normxdiffs = normxdiffs(1:iter);
info.x_cg_info = x_cg_info(1:iter);
if ~isempty(h) && ishandle(h), info.hfigure = h; end % handle of figure containing sos-image
if ~isempty(opts.evalfun) && isa(opts.evalfun,'function_handle')
    info.evaloutputs = evaloutputs(1:iter+1);
end
info.tinit = tinit;
info.times = times(1:iter);
xs(1).x = x; xs(1).t = sum(times(1:iter)); xs(1).iter = iter;

if opts.output == 1
    fprintf(1,'Iter #%d: ',iter);
    if opts.test_convergence
        if wavWeight > 0 || tvWeight > 0
            if r <= tol_r, str_r = '<='; else str_r = '>'; end
            if s <= tol_s, str_s = '<='; else str_s = '>'; end
            fprintf(1,'r = %g %s %g, s = %g %s %g, ',r,str_r,tol_r,s,str_s,tol_s);
        end
    end
    if normxdiffs(iter) <= opts.xtol*norm(col(x_prev)), str_x = '<='; else str_x = '>'; end
    fprintf(1,'dx = %g %s %g',normxdiffs(iter),str_x,opts.xtol*norm(col(x_prev)));

    if opts.test_convergence
        switch converged(iter)
            case 1
                fprintf(1,' (residual convergence).\n');
            case 2
                fprintf(1,' (x convergence).\n');
            otherwise
                fprintf(1,'.\n');
        end
    else
        fprintf(1,'.\n');
    end
    drawnow;
end

% nested helper functions (for CG and the like)
    function f = ffun()
        f = 0.5*lambda*norm(col(SPIRiTop*x))^2 + 0.5*norm(col(NUFFTop*x - kData))^2;
        if wavWeight > 0
            f = f + wavWeight*sum(sum(sqrt(sum(abs(DWTop*zpad(x,sz_impad)).^2,3)),2),1);
        end
        if tvWeight > 0
            if opts.tv_isotropic
                f = f + tvWeight*sum(sum(sqrt(sum(sum(abs(TVop*x).^2,4),3)),2),1);
            else
                f = f + tvWeight*sum(sum(sum(sqrt(sum(abs(TVop*x).^2,3)),4),2),1);
            end
        end
    end

    function out = update_x_cg_AHAfun(in)
        % out <- (lambda*(G-I)'*(G-I) + NUFFTop'*NUFFTop + mu*(DWT'*DWT + TV'*TV))*in
        
        in = reshape(in,sz_images);
        out = NUFFTop'*(NUFFTop*in) + lambda.*(SPIRiTop'*(SPIRiTop*in));
        if wavWeight > 0
            out = out + mu.*in; % assume DWT is orthonormal
%             out = out + mu.*crop(DWTop'*(DWTop*zpad(x,sz_impad)),sz_images);
        end
        if tvWeight > 0
            out = out + mu.*(TVop'*(TVop*in));
        end
        out = out(:);
        drawnow;
        
    end

    function out = update_x_cg_AHbfun()
        % out <- NUFFTop'*kData + mu*(DWT'*(w-u_w) + TV'*(gxy-u_gxy))
        
        out = NUFFTop'*kData;
        if wavWeight > 0
            out = out + mu.*crop(DWTop'*(w - u_w),sz_images);
        end
        if tvWeight > 0
            out = out + mu.*(TVop'*(gxy-u_gxy));
        end
        out = out(:);
        
    end

    function out = update_x_cg_Mfun(in)
        % out <- M.\in
        
        if ~isempty(x_precon_helper), out = in.*x_precon_helper; else out = in; end
    end



end
