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;
mn = params.mn;
mx = params.mx;
A = params.A;
AA = params.AA;
PSF2 = params.PSF2;
PSF = params.PSF;
params.rhofact = 1;

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

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

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

%% Simulate noisy data
if ~params.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( N ) + sqrt( -1 ) * randn( N ) );
    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 = 'l1'; % 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

if(~isempty( strfind( Penalty, 'FP' ) ) )
    params.delta = 1e-4;
    regname = [Penalty num2str(params.delta, '%0.2E') Operator]; % Analysis prior
else
    params.delta = 1;
    regname = [Penalty Operator]; % Analysis prior
end
params.R = R;
params.Penalty = Penalty;
params.regname = regname;

%% Parameters
condnumC = 1000;
niter_j = 10;
niter_k = 10;
niter = 25;
nbeta = 200;
beta = logspace(log10(stdnois) - 6, log10(stdnois) + 2, nbeta); 
minmax(beta)

paramsname = [ ...
               '_lams' num2str(min(beta(:)), '%0.2E') 'to' num2str(max(beta(:)), '%0.2E') 'n' int2str(nbeta) ...
               '_Itrs' int2str( niter ) '.' int2str( niter_j ) '.' int2str( niter_k ) ...
             ];

params.niter = niter;
params.niter_j = niter_j;
params.niter_k = niter_k;
params.dispinfo = 0;
params.paramsname = paramsname;

%% AA, RR, CinvAA, CinvA, etc response
ei = zeros( N ); ei(1,1) = 1;
RRr = R' * ( R * ei);
RRr = abs( fft2( RRr ) );
params.RRr = RRr;

AAr = PSF2;
mxAA = max(AAr(:)); mnAA = min(AAr(:)); % find alpha to obtain given condnum value
    
alpha = (mxAA - mnAA * condnumC) / (condnumC - 1); % Choose alpha based on condition number of C = A'A + alpha * I

params.mnAA = mnAA;
params.mxAA = mxAA;
params.AAr = AAr;
params.RRr = RRr;
params.alpha = alpha;

%% 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);

    tS = tic;
    [WSUREdump WMSEdump MSEdump PSNRdump Cdump udump RGCVdump tEdump] = WSURE_IRLSMIL_CSMRI_Nobeta(y, betacurr, params);
    tE(ibeta) = toc(tS);

    WSURE(:, ibeta) = WSUREdump;
    WMSE(:, ibeta) = WMSEdump;
    MSE(:, ibeta) = MSEdump;
    PSNR(:, ibeta) = PSNRdump;
    C(:, ibeta) = Cdump;
    RGCV(:, ibeta) = RGCVdump;
    
    disp(['ibeta: ' int2str(ibeta) '; Beta = ' num2str(betacurr) ...
        '; PSNR = ' num2str(PSNRdump( niter + 1 ), '%0.4E') '; MSE = ' num2str(MSEdump( niter + 1 ), '%0.4E') ...
        '; WMSE = ' num2str(WMSEdump( niter + 1 ), '%0.4E') '; WSURE = ' num2str(WSUREdump( niter + 1 ), '%0.4E') ...
        '; RGCV = ' num2str(RGCVdump( niter + 1 ), '%0.4E') ...
        '; T/itr = ' num2str( tEdump( niter + 1 ) ) ... 
        '; Time = ' num2str( tE( ibeta ) ) ...
        ]);
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_IRLSMIL_' params.fnamesetup '_' bsnrname '_' regname '_' paramsname];
params.fname = fnamevarybetaresults;

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

betaoptwsure = beta(idoptwsure);
betaoptmse = beta(idoptmse);
betaoptrgcv = beta(idoptrgcv);

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

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

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

%% Plot results
plotmany(log10(beta), ...
          [10 * log10([MSEbeta; abs(WSUREbeta); WMSEbeta; RGCVbeta]); PSNRbeta], ... % Opens in a new figure
          {'min', 'min', 'min', 'min', 'min', 'min', 'max'}, ...
          {[0 0 1], [0 0 0], [1 0 1], [0 0.5 0], [0.5 0.5 0], [0 0.5 0.5], [1 0 0]}, ...
          {'*', 'o', 's', 'v', '^', '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 ...
          );

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