clear;
clc;

%% Setup path to include Jeff's toolbox
if(isempty(strfind(path, '/alg/')))
    try
        addpath('~fessler/l/src/matlab/alg/');
        setup;
    catch ME
        currdir = pwd;
        cd('~/mfiles/');
        mypath = pathdef;
        path(mypath);
        cd(currdir);
    end
end

%% Load experimental setup workspace
savedworkspace = 'ZMAT/CSMRI/Z_CSMRI_SL_Rad30_psamp11.36_Experiment.mat';
load(savedworkspace, 'params');

xtrue = params.xtrue;
Npix = params.Npix;
Ndata = params.Ndata;
N = params.N; nx1 = N(1); nx2 = N(2);
mn = params.mn;
mx = params.mx;
A = params.A;
AA = params.AA;
AAr = params.AAr;
PSF2 = params.PSF2;
PSF = params.PSF;

if( ~isempty( strfind( savedworkspace, 'GE' ) ) ) 
    realdata = 1; 
else
    realdata = 0; 
end
params.realdata = realdata;

if( ~realdata )
    sigpow = params.sigpow; 
    Axtrue = params.Axtrue;
end

if( realdata )
    sigma_sq = params.sigma_sq;
    y = params.y;
    stdnois = sqrt(sigma_sq) / sqrt( 2 );
end

%% Simulate noisy data
if( ~realdata )
    BSNR = 40;
    sigma_sq = sigpow * 10 ^ ( -BSNR / 10 );
    stdnois = sqrt(sigma_sq) / sqrt( 2 ); % Circular Gaussian random vector with variance sigma_sq
    randn('state', 0);
    noisrealz = ( randn( [nx1 nx2] ) + sqrt( -1 ) * randn( [nx1 nx2] ) );
    noissamp = PSF .* noisrealz;
    nois = stdnois * noissamp;
    y = Axtrue + nois;
    
    Ay = A' * y; % Initial guess for u: suits mathematically since generally A is rectangular

    bsnrname = [num2str(BSNR, '%0.2f') 'dB'];
    params.Ay = Ay;
    params.sigma_sq = sigma_sq;
    params.bsnrname = bsnrname;
else
    bsnrname = 'RealData';
    params.bsnrname = bsnrname;
end

%% Choose Regularization: this is decided apriori, so keep it outside of the algo
Operator = 'FD';
Penalty = 'TV'; % could be TV, or smooth edge-preserving like Huber or Fair
switch Operator
    case{'FD'}
        R = GFD( 'mask', true( N ) ); % Finite differences in horizontal and vertical directions
        
    case{'W'}
        R = Gwave2( 'mask', true( N ) ); % 2 levels of undecimated Haar exlucding approximation level
end

regname = [Penalty Operator]; % Analysis prior
params.R = R;
params.Penalty = Penalty;
params.regname = regname;

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

%% Min condition number of A'A + R'R
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 = 1e-6;
params.Constants.maxBitr = 100;

params.dispBrentinfo = 0;
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 * 1e-4;
Gresp = 1 ./ (AAr + mu * RRr);
G = Gblur_S(true( [nx1 nx2] ), 'psffft', Gresp, 'type', 'fft,same');

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

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

paramsname = ['ALmu' num2str(mu, '%0.4E') '_lams' num2str(min(beta(:)), '%0.2E') 'to' num2str(max(beta(:)), '%0.2E') 'n' int2str(nbeta) '_Itrs' int2str(niter)];
params.niter = niter;
params.dispinfo = 0;
params.paramsname = paramsname;

%% AL algo
WSURE = zeros(niter + 1, nbeta); WMSE = WSURE; MSE = WMSE; PSNR = MSE; C = PSNR; RGCV = C;
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 PSNRdump Cdump udump RGCVdump] = WSURE_SB(y, betacurr, params);
    WSURE(:, ibeta) = WSUREdump;
    MSE(:, ibeta) = MSEdump;
    WMSE(:, ibeta) = WMSEdump;
    PSNR(:, ibeta) = PSNRdump;
    C(:, ibeta) = Cdump;
    RGCV(:, ibeta) = RGCVdump;
    
    disp(['ibeta: ' int2str(ibeta) '; Beta = ' num2str(betacurr, '%0.4E') '; PSNR = ' num2str(PSNRdump(end), '%0.4E') ...
        '; MSE = ' num2str(MSEdump(end), '%0.4E') '; WMSE = ' num2str(WMSEdump(end), '%0.4E') '; WSURE = ' num2str(WSUREdump(end), '%0.4E') ...
        '; RGCG = ' num2str(RGCVdump(end), '%0.4E') ...
        ]);
end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
matlabpool close force local;
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% Plot WSURE versus beta
WSUREbeta = WSURE(end, :); 
WMSEbeta = WMSE(end, :);
MSEbeta = MSE(end, :);
PSNRbeta = PSNR(end, :);
RGCVbeta = RGCV(end, :);

fnamevarybetaresults = ['Z_varybeta_CSMRI_SB_' params.fnamesetup '_' bsnrname '_' regname '_' paramsname];
params.fname = fnamevarybetaresults;

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

disp('Running SB for opt reg param by WSURE');
[WSUREoptwsure WMSEoptwsure MSEoptwsure PSNRoptwsure Coptwsure xoptwsure RGCVwsure] = WSURE_SB(y, betaoptwsure, params);

disp('Running SB for opt reg param by RGCV');
[WSUREoptrgcv WMSEoptrgcv MSEoptrgcv PSNRoptrgcv Coptrgcv xoptrgcv RGCVrgcv] = WSURE_SB(y, betaoptrgcv, params);

disp('Running SB for opt reg param by MSE');
[WSUREoptmse WMSEoptmse MSEoptmse PSNRoptmse Coptmse xoptmse RGCVoptmse] = WSURE_SB(y, betaoptmse, params);

%% Plot results
plotmany(log10(beta), ...
          [10 * log10([MSEbeta; abs(WSUREbeta); WMSEbeta; RGCVbeta]); PSNRbeta], ... % Opens in a new figure
          {'min', 'min', 'min', 'min', 'max'}, ...
          {[0 0 1], [0 0 0], [1 0 1], [0 0.5 0.5], [1 0 0]}, ...
          {'*', 'o', 's', 'd', 'x'}, ...
          {'MSE', 'WSURE', 'WMSE', 'RGCV', 'PSNR'}, ...
          'Reg Param (log_{10} scale)', ...
          {'MSE in dB', 'PSNR in dB'}, ...
          max([1 5]), ... % No. of markers to plot
          'bold', 12 ...
          );
% print( '-depsc2', ['~/mfiles/WSURE/ZEPS/' savefname( 1 : end-4 ) 'vsbeta.eps'] );
% 
% figure; clf; 
% if( any( imag( Ay(:) ) ) )
%     imagesc( abs( [ xtrue Ay; xoptmse xoptwsure ] ), [ min( xtrue(:) ) max( xtrue(:) ) ] );
% else
%     imagesc( [ xtrue y; xoptmse xoptwsure ], [ min( xtrue(:) ) max( xtrue(:) ) ] );
% end
% 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' ) ...
%     '; PSNR (orac, wsure) = ' num2str( PSNRoptmse( end ), '%0.2f' ) ', ' num2str( PSNRoptwsure( end ), '%0.2f' )  ...
%     ] ); 
% print( '-depsc2', ['~/mfiles/WSURE/ZEPS/' savefname( 1 : end-4 ) 'reconPSNR.eps'] );

%% Save results
% save([params.dirsetupmat params.fname '.mat']);