clear;
clc

%% Setup path to include Jeff's toolbox
if(isempty(strfind(path, '/alg/')))
    currdir = pwd;
    cd('~/mfiles/');
    mypath = pathdef;
    path(mypath);
    cd(currdir);
end

%% Test Image
try
    imgname = 'Cat';
    xtrue = imread('~/mfiles/testImages/cat.tif');
catch ME
    imgname = 'Cameraman';
    xtrue = imread('cameraman.tif');
end

imgname = 'Cameraman';
xtrue = imread('cameraman.tif');

xtrue = double(xtrue(:,:,1));
mx = max(xtrue(:));
mn = min(xtrue(:));
[nx1 nx2] = size(xtrue);
Npix = nx1 * nx2;

params.mn = mn;
params.mx = mx;
params.xtrue = xtrue;
params.Npix = Npix;

%% Blur, fatrices, etc
nblur = 15; 

% % Laplacian nonseparable
% nblurby2 = floor(nblur / 2);
% ii1 = -nblurby2 : nblurby2; ii2 = ii1;
% [ii1 ii2] = ndgrid(ii1, ii2);
% h = 1 ./ (1 + ii1 .^ 2 + ii2 .^ 2); blurname = ['Lap' int2str(nblur)]; 

% Gaussian
stddev = 2;
dx = 8 * stddev / (nblur - 1);
x = -4 * stddev : dx : 4 * stddev;
x2 = x .^ 2;
h = exp( -0.5 * x2 / stddev ^ 2); blurname = ['Gau' int2str(nblur)]; 

% % Uniform
% h = ones(1, nblur); blurname = ['Uni' int2str(nblur)]; 

% figure(1); plot(h); colorbar; axis equal off; title('psf');

if(min(size(h)) == 1) % check if it's a 1D blur
    h = h(:) * h; % 2D blur by outer product
end

psf = h / sum(h(:)); % Normalize blur to sum to 1

% figure(2); imagesc(log10((psf))); colorbar; axis equal off; title('psf');

tmp = ( fft2( psf, nx1, nx2 ) ); 
PSF2 = abs( tmp ) .^ 2;
% figure(3); imagesc(log10(fftshift(abs(PSF)))); colorbar; axis equal off; title('PSF');

A = Gblur_S( true( [nx1 nx2] ), 'psf', psf, 'type', 'fft,same');
AA = Gblur_S( true( [nx1 nx2] ), 'psffft', PSF2, 'type', 'fft,same'); % Sathish: Fatrix corresponding to (A' * A)

params.PSF2 = PSF2;
params.A = A;
params.AA = AA;

%% Simulate noisy data
BSNR = 40;
Axtrue = A * xtrue;
sigma_sq = norm( Axtrue(:) - mean( Axtrue(:) ) ) ^ 2 / numel( Axtrue(:) ) * 10 ^ ( -BSNR / 10 );
stdnois = sqrt(sigma_sq);
randn('state',1);
y = Axtrue + stdnois * randn( [nx1 nx2] );

params.sigma_sq = sigma_sq;

bsnrname = [num2str(BSNR, '%0.2f') 'dB'];

%% Sathish: Choose Regularization: this is decided apriori, so keep it outside of the algo
Operator = 'FD';
Penalty = 'l1'; % could be TV, or smooth edge-preserving like Huber or Fair
switch(Operator)
    case{'FD'}
        R = Cdiffs( [nx1 nx2], 'offsets', [1 nx1], 'type_diff', 'circshift'); % circulant
        
    case{'W'}
        R = Gwave2( 'mask', true( [nx1 nx2] ) ); % 2 levels of undecimated Haar exlucding approximation level
end

params.R = R;
params.Penalty = Penalty;

regname = [Penalty Operator]; % Analysis prior

%% Sathish: Initial guess for u: this is decided apriori, so keep it outside of the algo
Ay = A' * y; % Better to smooth the noise in y, also suits mathematically since generally A is rectangular

params.Ay = Ay;

%% Sathish: Generate random vector for WSURE calculation: this can be done offline and so keep it outside of the algo
r = ones( numel( xtrue ), 1 );
r( floor( numel( xtrue ) / 2 ) + 1 : numel( xtrue ) ) = -1;
rand('state',10);
for ii = 1:10
    permindex = randperm( numel( xtrue ) );
    r = r( permindex );
end
r = reshape( r, [nx1 nx2] );

params.r = r;

%% Constants
params.Constants.EPSL = 1.0E-10; %% Some constants needed for Brent iterations for finding optimal lambda
params.Constants.CGOLD = 0.3819660;
params.Constants.GOLD = 1.618034;
params.Constants.GLIMIT = 100.0;
params.Constants.Brenttol = 0.001;
params.Constants.maxBitr = 100;

%% RR response
ei = zeros( [nx1 nx2] );
ei(1,1) = 1;
RRr = R' * ( R * ei);
RRr = abs( fft2( RRr ) );
params.RRr = RRr;

%% AA response
AAr = AA * ei;
AAr = abs( fft2( AAr ) );
params.AAr = AAr;

%% Min condition number of A'A + R'R
params.todo = 2; %% Find condition number of the system S'F'FS + mu*RR where the first term has been approximated by a circulant matrix
params.BRbx = 1e-2; %% Center of first bracket
[minCondNum, mumin] = Brent_Linmin([], real(max(AAr(:))/min(AAr(:))), params, []); %% Input arguments unused except for "params"; Need to simplify code using varargin!

mu = mumin / 100;
Gresp = 1 ./ (AAr + mu * RRr);
G = Gblur_S(true( [nx1 nx2] ), 'psffft', Gresp, 'type', 'fft,same');

params.mu = mu;
params.G = G;

%% Parameters
niter = 100;
nbeta = 80;
beta = logspace(log10(stdnois) - 4, log10(stdnois) + 2, nbeta); 
minmax(beta)

params.niter = niter;
params.dispinfo = 0;
params.dispBrentinfo = 1;

paramsname = ['ALmu' num2str(mu, '%0.4f') '_nlams' int2str(nbeta) '_Itrs' int2str(niter)];

%% AL algo
WSURE = zeros(niter + 1, nbeta); WMSE = WSURE; MSE = WMSE; ISNR = MSE; C = ISNR;
nworkers = nbeta; if(nworkers > 8) nworkers = 8; end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
matlabpool close force local; if(nworkers > 1) matlabpool('open', nworkers); end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

parfor ibeta = 1 : nbeta
    betacurr = beta(ibeta);
    
    [WSUREdump WMSEdump MSEdump ISNRdump Cdump] = WSURE_SB(y, betacurr, params); % (y, Ay, A, G, R, sigma_sq, betacurr, mu, r, Penalty, niter, params);
    WSURE(:, ibeta) = WSUREdump;
    MSE(:, ibeta) = MSEdump;
    WMSE(:, ibeta) = WMSEdump;
    ISNR(:, ibeta) = ISNRdump;
    C(:, ibeta) = Cdump;
    
    disp(['ibeta: ' int2str(ibeta) '; Beta = ' num2str(betacurr) '; ISNR = ' num2str(ISNRdump(end), '%0.4f') ...
        '; MSE = ' num2str(MSEdump(end), '%0.4f') '; WMSE = ' num2str(WMSEdump(end), '%0.4f') '; WSURE = ' num2str(WSUREdump(end), '%0.4f') ...
        ]);
end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
matlabpool close force local;
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% Plot WSURE versus beta
WSUREbeta = WSURE(end, :); 
WMSEbeta = WMSE(end, :);
MSEbeta = MSE(end, :);
ISNRbeta = ISNR(end, :);

savefname = ['Z_varybeta_Deblur_SB_' imgname '_' blurname '_' bsnrname '_' regname '_' paramsname '.mat'];

%% Run the algo for optimal betas
[dump idoptwsure] = min(WSUREbeta(:)); betaoptwsure = beta(idoptwsure);
[dump idoptmse] = min(MSEbeta(:)); betaoptmse = beta(idoptmse);

[WSUREoptwsure MSEoptwsure WMSEoptwsure ISNRoptwsure Coptwsure xoptwsure] = WSURE_SB(y, betaoptwsure, params); % (y, Ay, A, G, R, sigma_sq, betaoptwsure, mu, r, Penalty, niter, params);
[WSUREoptpmse MSEoptmse WMSEoptmse ISNRoptmse Coptmse xoptmse] = WSURE_SB(y, betaoptmse, params); % (y, Ay, A, G, R, sigma_sq, betaoptmse, mu, r, Penalty, niter, params);

%% Plot results
plot2axes(log10(beta), [10 * log10( [WSUREbeta; WMSEbeta]); ISNRbeta], ... % Opens in a new figure
          {'WSURE', 'WMSE', 'ISNR'}, ...
          'Reg Param (log_{10} scale)', ...
          {'WMSE, WSURE in dB', 'ISNR in dB'}, ...
          'ISNR');
print( '-depsc2', ['~/mfiles/WSURE/ZEPS/' savefname( 1 : end-4 ) 'vsbeta.eps'] );

figure; clf; 
imagesc( [ xtrue y; xoptmse xoptwsure ], [ min( xtrue(:) ) max( xtrue(:) ) ] );
colormap gray; axis equal off;
print( '-depsc2', ['~/mfiles/WSURE/ZEPS/' savefname( 1 : end-4 ) 'recon.eps'] );
title( [ 'Beta (orac, wsure) = ' num2str( betaoptmse , '%0.2E' ) ', ' num2str( betaoptwsure, '%0.2E' ) ...
    '; ISNR (orac, wsure) = ' num2str( ISNRoptmse( end ), '%0.2f' ) ', ' num2str( ISNRoptwsure( end ), '%0.2f' )  ...
    ] ); 
print( '-depsc2', ['~/mfiles/WSURE/ZEPS/' savefname( 1 : end-4 ) 'reconISNR.eps'] );

% plot( log10(beta), 10 * log10( [ WSUREbeta; WMSEbeta ] ) ); legend('WSURE', 'WMSE'); xlabel('Reg Param (log_{10} scale)'); ylabel('WMSE, WSURE in dB'); hold on;
% plot( log10(beta(idbetaoptmse)), 10 * log10(WMSEbeta(idbetaoptmse)), 'r*'); 

%% Save results
save(['~/mfiles/WSURE/ZMAT/' savefname]);