function [WSURE WMSE MSE PSNR C u RGCV] = WSURE_SB(y, beta, params)
% Deblur and CSMRI recon using split-Bregman algorithm
% Splitting only the regularization term
% Both problems are shift-invariant and requires inversion of a circulant matrix that is implemented using FFTs
% Jacobian matrix estimation requires both Ju(y) and Ju(y*)
%
% Computes RGCV (a Monte-Carlo version of GCV for nonlinear algos) for comparison; RGCV = (1 / Nsamp) * || y - A * f(y) ||^2 / ( 1 - (1 / Nsamp) trace{ J_f(y) } )^2
% Computes Predicted MSE for deblur and Projected MSE for undersampled MR image recon (same as Predicted in this case)

%% Load params
xtrue = params.xtrue;
dispinfo = params.dispinfo;
A = params.A;
Ay = params.Ay;
G = params.G;
R = params.R;
sigma_sq = params.sigma_sq;
if(isfield(params, 'sigma_sq_const'))
    sigma_sq_const = params.sigma_sq_const;
else
    sigma_sq_const = sigma_sq;
end
mu = params.mu;
penalt = params.Penalty;
niter = params.niter;
r = params.r;
Npix = params.Npix;
Ndata = params.Ndata;

Ndinv = 1 / Ndata;
Ninv = 1 / Npix;
mn = params.mn;
mx = params.mx;
isnrnum = Ninv * sum( abs( y(:) - xtrue(:) ) .^ 2 );

%% Initial estimates for AL algo and WSURE
u = Ay;
Ru = R * u;
v = Ru; 
eta = zeros( size( Ru ) );
dimsRu = size(Ru);

% Jacobian matrices times random vr
Jur = A' * r; cJur = zeros(size(Jur));
GAr = G * Jur; 
RJur = R * Jur; cRJur = R * cJur;
Jvr = RJur; cJvr = cRJur;
Jetar = eta; cJetar = eta;

% Is complex ?
if( any( imag( u(:) ) ) )  iscomplex = 1; else iscomplex = 0; end

% WSURE, etc storage-vectors
WSURE = zeros(niter + 1, 1); WMSE = WSURE; MSE = WMSE; PSNR = MSE; C = PSNR; div = C; RGCV = div;

%% Compute cost and display
Au = A * u; % Redundant, can be avoided if not monitoring cost
switch(penalt)
    case{'lange3'}
        regterm = sum( abs( Ru(:) ) ) / delta - sum( log( 1 + abs( Ru(:) ) / delta ) ); % Redundant, can be avoided if not monitoring cost
        
    case{'l1'}
        regterm = sum( abs( Ru(:) ) ) ; % Redundant, can be avoided if not monitoring cost
        
    case{'TV'}
        regterm = sqrt( sum( abs( Ru ) .^ 2, 3 ) ); regterm = sum( regterm(:) );
end
cnew = 0.5 * sum( abs( y(:) - Au(:) ) .^ 2 ) + beta * regterm;
C(1) = cnew;

difs = xtrue - u;
Adifs = A * difs;
MSE(1) = Ninv * sum( abs( difs(:) ) .^ 2 );
WMSE(1) = Ndinv * sum( abs( Adifs(:) ) .^ 2 );
WSURE(1) = WMSE(1);

RGCVnume = Ndinv * sum( abs( y(:) - Au(:) ) .^ 2 ); 
RGCVdeno = r .* real( A * Jur ); RGCVdeno = ( 1 - Ndinv * sum( RGCVdeno(:) ) ) ^ 2; if(RGCVdeno == 0) RGCVdeno = RGCVnume ^ 2; end
RGCV(1) = RGCVnume / RGCVdeno;

if( iscomplex ) PSNR(1) = 10 * log10( mx / MSE(1) ); else PSNR(1) = 10 * log10( isnrnum / MSE(1) ); end

if(dispinfo)
    cold = 2 * cnew; % Redundant, can be avoided if not monitoring cost
    dcost = cnew - cold; % Redundant, can be avoided if not monitoring cost
    disp(['SB0; Lam=' num2str(beta, '%0.4E') '; C=' num2str(cnew, '%0.5E') '; DC= ' num2str(dcost, '%0.5E') '; PSNR=' num2str(PSNR(1), '%0.4E') ...
        '; MSE=' num2str(MSE(1), '%0.4E') '; WMSE=' num2str(WMSE(1), '%0.4E') '; WSURE=' num2str(WSURE(1), '%0.4E') ...
        '; RGCV=' num2str(RGCV(1), '%0.4E') ...
        ]);
    figure(3); clf; 
    subplot(2, 2, 1); cla;
    uini = Ay;
    if( iscomplex ) imagesc( abs( uini ), [mn mx]); else imagesc( u, [mn mx]); end
    colormap gray; colorbar; axis equal off; title(['PSNR: ' num2str(PSNR(1))]);
end

%% ========================= AL algo =========================
for ii = 1 : niter
    % Update u vector
    u = G * ( Ay + mu * ( R' * ( v - eta ) ) );
    Ru = R * u;
    
    Jur = GAr + mu * G * ( R' * ( Jvr - Jetar ) );
    cJur = mu * G * ( R' * ( cJvr - cJetar ) );
    
    RJur = R * Jur;
    cRJur = R * cJur;
    
    % Update v vector
    rhsv = Ru + eta;
    thresh = beta / mu;
    switch(penalt)
        case{'lange3'}
           error('lange3: To be done');

        case{'l1'}
            v = wthresh(rhsv , 's', thresh);
            signv = abs( sign( v ) );
            if( iscomplex )
                absrhsv = abs( rhsv );
                absrhsv( absrhsv == 0 ) = eps; % to Avoid NAN's; this does not affect the solution since when ||rhsv|| == 0,  v = 0 too!
                Jdrho = signv .* ( 1 - thresh / 2 ./ absrhsv ); % see TIP manuscript on WSURE for MRI using split-Bregman
                cJdrho = 0.5 * thresh * signv .* rhsv .^ 2 ./ absrhsv .^ 3;
                Jvr = Jdrho .* ( RJur + Jetar ) + cJdrho .* conj( cRJur + cJetar );
                cJvr = Jdrho .* ( cRJur + cJetar ) + cJdrho .* conj( RJur + Jetar );
            else
                Jvr = abs( sign( v ) ) .* ( RJur + Jetar );
            end
            
        case{'TV'}
            rhsvnorm = repmat( sqrt( sum( abs( rhsv .^ 2 ), 3) ), [1 1 dimsRu(3)] );
            rhsvnorm( rhsvnorm == 0 ) = eps; % to Avoid NAN's; this does not affect the solution since when rhsv == 0,  v = 0 too!
            maxop = rhsvnorm - thresh;
            maxop = ( maxop + abs( maxop ) ) / 2;
            v = rhsv ./ rhsvnorm .* maxop;
            signv = abs( sign( v ) );
            if( iscomplex )
                Jdrho_Jrhoy = signv .* ( ...
                        ( 1 - thresh ./ rhsvnorm ) .* ( RJur + Jetar ) + ...
                        0.5 * thresh * rhsv ./ rhsvnorm .^ 2 .* repmat( sum( conj( rhsv ) .* ( RJur + Jetar ) ./ rhsvnorm, 3 ), [1 1 dimsRu(3)] ) ...
                        );
                    
                cJdrho_cJrhoy = 0.5 * thresh * signv .* rhsv ./ rhsvnorm .^ 2 .* repmat( sum( rhsv .* conj( cRJur + cJetar ) ./ rhsvnorm, 3 ), [1 1 dimsRu(3)] );
                
                Jdrho_cJrhoy = signv .* ( ...
                        ( 1 - thresh ./ rhsvnorm ) .* ( cRJur + cJetar ) + ...
                        0.5 * thresh * rhsv ./ rhsvnorm .^ 2 .* repmat( sum( conj( rhsv ) .* ( cRJur + cJetar ) ./ rhsvnorm, 3 ), [1 1 dimsRu(3)] ) ...
                        );
                    
                cJdrho_Jrhoy = 0.5 * thresh * signv .* rhsv ./ rhsvnorm .^ 2 .* repmat( sum( rhsv .* conj( RJur + Jetar ) ./ rhsvnorm, 3 ), [1 1 dimsRu(3)] );

                Jvr = Jdrho_Jrhoy + cJdrho_cJrhoy;
                cJvr = Jdrho_cJrhoy + cJdrho_Jrhoy;
            else
                Jvr = zeros( size( v ) );
                % Jvr = abs( sign( v ) ) .* ( ...
                %       ( 1 - thresh / rhsvnorm ) .* ( RJur + Jetar ) + ...
                %       thresh * rhsv ./ rhsvnorm .^ 3 .* repmat( sum( rhs .* ( RJur + Jetar ), 3 ), [1 1 dimsRu(3)] ) ...
                %       ); % This is probably WRONG! No experiments have been done for deblur with TV and SB, so have not verified if this works correctly!
            end
    end

    % Update eta
    eta = eta - ( v - Ru );
    Jetar = Jetar - ( Jvr - RJur );
    cJetar = cJetar - ( cJvr - cRJur );
    
    % Compute the divergence term
    if( iscomplex )
        divvr = r .* real( A * Jur ); % Use Projected SURE
        divc = r .* ( A * Jur );
    else    
        divvr = r .* ( A * Jur ); % Use Projected SURE
        divc = divvr;
    end
    div(ii + 1) = sum( divvr(:) );
    divc = sum( divc(:) );
    
    RGCVnume = Ndinv * sum( abs( y(:) - Au(:) ) .^ 2 );
    RGCVdeno = ( 1 - Ndinv * div(ii + 1) ) ^ 2; if(RGCVdeno == 0) RGCVdeno = RGCVnume ^ 2; end
    
    % Compute WSURE, etc
    Au = A * u;
    difs = xtrue - u;
    Adifs = A * difs;
    MSE(ii + 1) = Ninv * sum( abs( difs(:) ) .^ 2 );
    WMSE(ii + 1) = Ndinv * sum( abs( Adifs(:) ) .^ 2 );
    WSURE(ii + 1) = Ndinv * sum( abs( y(:) - Au(:) ) .^ 2 ) - sigma_sq_const + 2 * Ndinv * sigma_sq * div(ii + 1);
    RGCV(ii + 1) = RGCVnume / RGCVdeno;
    
    if( iscomplex )
        PSNR(ii + 1) = 10 * log10(  mx / MSE(ii + 1) );
    else
        PSNR(ii + 1) = 10 * log10( isnrnum / MSE(ii + 1) );
    end
    
    % Compute cost and display
    switch(penalt)
        case{'lange3'}
            regterm = sum(abs(Ru(:))) / delta - sum( log(1 + abs(Ru(:)) / delta) ); % Redundant, can be avoided if not monitoring cost
            
        case{'l1'}
            regterm = sum(abs(Ru(:))) ; % Redundant, can be avoided if not monitoring cost
            
        case{'TV'}
            regterm = sqrt( sum( abs( Ru ) .^ 2, 3 ) ); regterm = sum( regterm( : ) );
    end
    cold = cnew; % Redundant, can be avoided if not monitoring cost
    cnew = 0.5 * sum( abs( y(:) - Au(:) ) .^ 2 ) + beta * regterm;
    C(ii + 1) = cnew;
    if(dispinfo)
        dcost = cnew - cold; % Redundant, can be avoided if not monitoring cost
        disp(['SB' int2str(ii) '; Lam=' num2str(beta, '%0.4E') '; C=' num2str(cnew, '%0.5E') '; DC= ' num2str(dcost, '%0.5E') '; PSNR=' num2str(PSNR(ii + 1), '%0.4E') ...
            '; MSE=' num2str(MSE(ii + 1), '%0.4E') '; WMSE=' num2str(WMSE(ii + 1), '%0.4E') '; WSURE=' num2str(WSURE(ii + 1), '%0.4E') ...
            '; RGCV=' num2str(RGCV(ii + 1), '%0.4E') ...
            '; Div=' num2str(divc, '%0.4E') ...
            ]);

        figure(3); clf;
        subplot(2, 2, 1); cla;
        if( iscomplex ) imagesc( abs( uini ), [mn mx]); else imagesc( u, [mn mx]); end
        colormap gray; colorbar; axis equal off; title(['PSNR: ' num2str(PSNR(1))]);
        
        subplot(2, 2, 2); cla; 
        if( iscomplex ) imagesc( abs( u ), [mn mx]); else imagesc( u, [mn mx]); end
        colormap gray; colorbar; axis equal off; title(['PSNR: ' num2str(PSNR(ii))]);
        
        subplot(2, 2, 3); cla; plot( log10( MSE( 1 : ii ) ) );
        subplot(2, 2, 4); cla; plot( log10( WMSE( 1 : ii ) ), 'r'); hold on; plot( log10( abs(WSURE( 1 : ii ) ) ), 'k');
    end
end

if(dispinfo)
    figure(3); clf;
    subplot(2, 2, 1); cla;
    if( iscomplex ) imagesc( abs( uini ), [mn mx]); else imagesc( u, [mn mx]); end
    colormap gray; colorbar; axis equal off; title(['PSNR: ' num2str(PSNR(1))]);

    subplot(2, 2, 2); cla;
    if( iscomplex ) imagesc( abs( u ), [mn mx]); else imagesc( u, [mn mx]); end
    colormap gray; colorbar; axis equal off; title(['PSNR: ' num2str(PSNR(end))]);
    
    subplot(2, 2, 3); cla; plot( log10( MSE ) );
    subplot(2, 2, 4); cla; plot( log10( WMSE ), 'r'); hold on; plot( log10( abs( WSURE ) ), 'k'); legend('WMSE in red', 'WSURE in black');
end