% Setup Experiment to run IRLSMIL with varying regulariation parameter values and compute SURE-Metrics
%
% Sathish Ramani, University of Michigan, 2011

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

%% Load experimental setup workspace
savedfname = 'ZMAT/DEBLUR/Z_Deblur_Cameraman_Uni9_Experiment.mat';
load(savedfname, 'params');

xtrue = params.xtrue;
Npix = params.Npix;
N = params.N; nx1 = N(1); nx2 = N(2);
mn = params.mn;
mx = params.mx;
A = params.A;
AA = params.AA;
ei = params.ei;
PSF2 = params.PSF2;
PSF = params.PSF;
Axtrue = params.Axtrue;
sigpow = params.sigpow;

%% Fatrices for Projected MSE and Projected SURE
% Numerically stable Pseudo-inverse of PSF2 for Projected MSE 
svtol = max( [ eps * max([nx1 nx2]) * max(PSF2(:)) 1e-10 ] );
iPSF2 = PSF2(:); 
id1 = iPSF2 <= svtol; iPSF2(id1) = 0;
id2 = iPSF2 > svtol; iPSF2(id2) = 1 ./ PSF2(id2);
iPSF2 = reshape(iPSF2, [nx1 nx2]);
iPSF2_sqrt = sqrt(iPSF2);
% ndata_PJMSE = length(find(id2));
condnumPSF2 = max(PSF2(:)) / min(PSF2(:));

% Numerically stable Pseudo-inverse of PSF2 for Projected SURE
if(~isempty(strfind(savedfname, 'Gau'))) || (~isempty(strfind(savedfname, 'Uni')))
    svtol_PJ = 1e-3; % uniform 9 x 9 blur and % Gaussian 17 x 17 blur with std = 2
else
    svtol_PJ = svtol; % Lorentzian 15 x 15 blur
end

iPSF2_PJ = PSF2(:); 
id1 = iPSF2_PJ <= svtol_PJ; iPSF2_PJ(id1) = 0;
id2 = iPSF2_PJ > svtol_PJ; iPSF2_PJ(id2) = 1 ./ PSF2(id2);
iPSF2_PJ = reshape(iPSF2_PJ, [nx1 nx2]);
iPSF2_PJ_sqrt = sqrt(iPSF2_PJ);
trW = sum(iPSF2_PJ(:));
ndata_PJSURE = length(find(id2));

VA = Gblur_S( true( [nx1 nx2] ), 'psffft', iPSF2_sqrt .* PSF, 'type', 'fft,same'); % Fatrix corresponding to V * A
V = Gblur_S( true( [nx1 nx2] ), 'psffft', iPSF2_PJ_sqrt, 'type', 'fft,same'); % Fatrix corresponding to V; used in ProJ SURE only
WA = Gblur_S( true( [nx1 nx2] ), 'psffft', iPSF2_PJ .* PSF, 'type', 'fft,same'); % Fatrix corresponding to W * A; used in ProJ SURE only

params.svtol = svtol;
params.svtol_PJ = svtol_PJ;
params.iPSF2_sqrt = iPSF2_sqrt;
params.iPSF2_PJ_sqrt = iPSF2_PJ_sqrt;
params.iPSF2 = iPSF2;
params.iPSF2_PJ = iPSF2_PJ;
params.trW = trW;
params.V = V;
params.VA = VA;
params.WA = WA;
% params.ndata_PJMSE = ndata_PJMSE;
params.ndata_PJSURE = ndata_PJSURE;

%% Choose Regularization: this is decided apriori, so keep it outside of the algo
Operator = 'W';
Penalty = 'l1'; % could be TV, or smooth edge-preserving like Huber or Fair
switch(Operator)
    case{'FD'}
        R = GFD( 'mask', true( N ) );
        % 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
if(~isempty( strfind( Penalty, 'FP' ) ) )
    params.delta = 1e-4;
    regname = [Penalty num2str(params.delta, '%0.2E') Operator]; % Analysis prior
else
    regname = [Penalty Operator]; % Analysis prior
end
params.R = R;
params.Penalty = Penalty;
params.regname = regname;

%% Simulate noisy data
BSNR = 40;
sigma_sq = sigpow * 10 ^ ( -BSNR / 10 );
stdnois = sqrt(sigma_sq);
randn('state', 0);
y = Axtrue + stdnois * randn( [nx1 nx2] );
params.sigma_sq = sigma_sq;

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

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

%% Parameters
condnumC = 20;
niter = 100;
niter_in = 1;
nbeta = 200;
% beta = logspace(log10(stdnois) - 8, log10(stdnois) - 3, nbeta); % Gau 17 x 17 blur
% beta = logspace(log10(stdnois) - 6, log10(stdnois) + 1, nbeta); % Uni 9 x 9 blur
beta = logspace(log10(stdnois) - 4, log10(stdnois) + 2, nbeta); % Lap 15 x 15 blur
minmax(beta)
paramsname = ['CondC' num2str(condnumC, '%0.0f') '_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.dispBrentinfo = 1;
params.niter_in = niter_in;
params.paramsname = paramsname;

%% AA, RR, CinvAA, CinvA, etc response
RRr = R' * ( R * ei);
RRr = abs( fft2( RRr ) );

AAr = AA * ei;
AAr = abs( fft2( AAr ) );
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; ISNR = MSE; C = ISNR; RGCV = ISNR; PJMSE = ISNR; PJWSURE = ISNR; AGCV = 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 udump RGCVdump PJSUREdump PJMSEdump] = WSURE_IRLSMIL_Nobeta(y, betacurr, params);
    AGCVdump = AGCV_IRLSMIL(betacurr, params);
    
    WSURE(:, ibeta) = WSUREdump;
    WMSE(:, ibeta) = WMSEdump;
    MSE(:, ibeta) = MSEdump;
    ISNR(:, ibeta) = ISNRdump;
    C(:, ibeta) = Cdump;
    RGCV(:, ibeta) = RGCVdump;
    PJSURE(:, ibeta) = PJSUREdump;
    PJMSE(:, ibeta) = PJMSEdump;
    AGCV(:, ibeta) = AGCVdump;
    
    disp(['ibeta: ' int2str(ibeta) '; Beta = ' num2str(betacurr) ...
        '; ISNR = ' num2str(ISNRdump(end), '%0.4f') '; MSE = ' num2str(MSEdump(end), '%0.4f') ...
        '; PJMSE = ' num2str(PJMSEdump(end), '%0.4f') '; PJSURE = ' num2str(PJSUREdump(end), '%0.4f') ...
        '; WMSE = ' num2str(WMSEdump(end), '%0.4f') '; WSURE = ' num2str(WSUREdump(end), '%0.4f') ...
        '; RGCV = ' num2str(RGCVdump(end), '%0.4f') '; AGCV = ' num2str(AGCVdump(end), '%0.4f') ...
        ]);
end
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
matlabpool close force local;
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% Plot WSURE versus beta
WSUREbeta = WSURE(end, :); 
WMSEbeta = WMSE(end, :);
MSEbeta = MSE(end, :);
ISNRbeta = ISNR(end, :);
RGCVbeta = RGCV(end, :);
PJSUREbeta = PJSURE(end, :);
PJMSEbeta = PJMSE(end, :);
AGCVbeta = AGCV(end, :);

fnamevarybetaresults = ['Z_varybeta_Deblur_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(:)); 
[dump idoptpjsure] = min(PJSUREbeta(:)); 
[dump idoptagcv] = min(AGCVbeta(:)); 

betaoptwsure = beta(idoptwsure);
betaoptmse = beta(idoptmse);
betaoptrgcv = beta(idoptrgcv);
betaoptpjsure = beta(idoptpjsure);
betaoptagcv = beta(idoptagcv);

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

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

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

disp('Running IRLSMIL for opt reg param by PJSURE');
[WSUREoptpjsure MSEoptpjsure WMSEoptpjsure ISNRoptpjsure Coptpjsure xoptpjsure RGCVoptpjsure PJSUREpjsure PJMSEpjsure] = WSURE_IRLSMIL_Nobeta(y, betaoptpjsure, params);

disp('Running IRLSMIL for opt reg param by AGCV');
[WSUREoptagcv MSEoptagcv WMSEoptagcv ISNRoptagcv Coptagcv xoptagcv RGCVoptagcv PJSUREagcv PJMSEagcv] = WSURE_IRLSMIL_Nobeta(y, betaoptagcv, params);

%% Plot results
plotmany(log10(beta), ...
          [10 * log10([MSEbeta; abs(WSUREbeta); WMSEbeta; abs(PJSUREbeta); PJMSEbeta; RGCVbeta]); ISNRbeta], ... % 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', 'PJSURE', 'PJMSE', 'RGCV', 'ISNR'}, ...
          'Reg Param (log_{10} scale)', ...
          {'MSE in dB', 'ISNR in dB'}, ...
          max([1 5]), ... % No. of markers to plot
          'bold', 12 ...
          );
% plot(log10(beta), ISNRbeta, 'k', ... % Opens in a new figure
%      log10(betaoptmse), ISNRbeta(idoptmse), 'r*', ...
%      log10(betaoptpjsure), ISNRbeta(idoptpjsure), 'mo', ...
%      log10(betaoptwsure), ISNRbeta(idoptwsure), 'ms', ...
%      log10(betaoptrgcv), ISNRbeta(idoptrgcv), 'b+', ...
%      log10(betaoptagcv), ISNRbeta(idoptagcv), 'bx' ...
%     );
% legend('ISNR', 'MSE-optimal', 'ProjSURE-based', 'PredSURE-based', 'NGCV-based', 'AGCV-based');
% 
% print( '-depsc2', [params.dirsetupeps fnamevarybetaresults( 1 : end-4 ) 'vsbeta.eps'] );
% 
% figure; clf; 
% imagesc( [ xtrue y xoptmse xoptwsure xoptpjsure xoptrgcv xoptrgcv], [ mn mx ] );
% colormap gray; axis equal off;
% print( '-depsc2', [params.dirsetupeps fnamevarybetaresults( 1 : end-4 ) 'recon.eps'] );
% 
% title( [ '(O,W,P,G,A): ' num2str( betaoptmse, '%0.2E' ) ', ' num2str( betaoptwsure, '%0.2E' ) ...
%     ', ' num2str( betaoptpjsure, '%0.2E' ) ', ' num2str( betaoptrgcv, '%0.2E' ) ', ' num2str( betaoptagcv, '%0.2E' )...
%     '; ISNR: ' num2str( ISNRoptmse( end ), '%0.2f' ) ', ' num2str( ISNRoptwsure( end ), '%0.2f' )  ...
%     ', ' num2str( ISNRoptpjsure( end ), '%0.2f' ) ', ' num2str( ISNRoptrgcv( end ), '%0.2f' ) ', ' num2str( ISNRoptagcv( end ), '%0.2f' ) ...
%     ] ); 
% print( '-depsc2', [params.dirsetupeps fnamevarybetaresults( 1 : end-4 ) 'reconISNR.eps'] );

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