% example_xct_beam_hardening.m
%
% X-ray CT analytic and iterative reconstruction of an object
% consisting of water, bone and fat for a polyenergetic spectrum.
%
% Jeff Fessler, University of Michigan
% Cristóbal Martínez and Mónica Abella, Universidad Carlos III de Madrid
% Updated: April 2019

clear all
close all
prompt run
warn 'work in progress'
warning('off','all')

% Acquisition variables
na_val=60; % Number of projections
I0=1e6; % Number of photons. % 1e5 low snr  1e6; high snr
dx_val=0.08; % Pixel size

% Phantom variables
% Density
mu_bone=1.9;
mu_bone2=2.1;
mu_fat=0.95;
mu_ST=1.06;

%&Create phantom
nx_val=512;
ny_val=512;

ig = image_geom('nx', nx_val, 'ny', ny_val, 'dx', dx_val);
ig.mask = ig.circ > 0;     


% Soft tissue ellipse
ellw = [0 0 15 12 0 mu_ST
        -10 3 0.8 0.8 0 -mu_ST
        -10 -3 0.8 0.8 0 -mu_ST
        -5 4 1.2 1.2 0 -mu_ST
        -5 -4 1.2 1.2 0 -mu_ST
        -0 5 1.2 1.8 0 -mu_ST
        -0 -5 1.2 1.8 0 -mu_ST
        5 6 1.4 1.9 0 -mu_ST
        5 -6 1.4 1.9 0 -mu_ST
        10 3 1.4 1.4 0 -mu_ST
        10 -3 1.4 1.4 0 -mu_ST];
wtrue = ellipse_im(ig, ellw, 'oversample', 4);

% Bone ellipses
ellb = [-10 3 0.8 0.8 0 mu_bone
        -10 -3 0.8 0.8 0 mu_bone
        -5 4 1.2 1.2 0 mu_bone2
        -5 -4 1.2 1.2 0 mu_bone2
        -0 5 1.2 1.8 0 mu_bone
        -0 -5 1.2 1.8 0 mu_bone
        5 6 1.4 1.9 0 mu_bone2
        5 -6 1.4 1.9 0 mu_bone2];
btrue = ellipse_im(ig, ellb, 'oversample', 4);     

% Fat ellipses
ella=[10 3 1.4 1.4 0 mu_fat
 10 -3 1.4 1.4 0 mu_fat];

atrue=ellipse_im(ig, ella, 'oversample', 4);               

% Density true phantom
dens_true = wtrue + btrue+atrue;
    
% Image and sinogram variables
ig = image_geom('nx', nx_val, 'ny', ny_val,'dx',dx_val); 
sg = sino_geom('par', 'nb', nx_val, 'na', na_val, 'dr', ig.dx); 
ig.mask = ig.circ > 0;

% Create projection matrix
if ~isvar('G'), printm 'G'
	if has_mex_jf
		f.tab_type = {'square/strip', 'chat', 0, 'Ltab', ...
			1000, 'strip_width', ig.dx};
		G = Gtomo2_table(sg, ig, f.tab_type, 'nthread', 1);
	else
		G = Gtomo2_strip(sg, ig);
	end
end

% Complete sinogram
sgComplete = sino_geom('par', 'nb', nx_val, 'na', 180, 'dr', ig.dx); 
% Create Fatrix for complete reconstruction
GComplete = Gtomo2_strip(sgComplete, ig);

wsino = GComplete*wtrue;
bsino = GComplete*btrue;
asino=GComplete*atrue;
ssinoGold = cat(3, wsino, bsino,asino);

% Gold standard reconstruction
f.fbpComplete = fbp2(sgComplete, ig);
dens_true_rec=fbp2(sum(ssinoGold,3), f.fbpComplete, 'window', 'hanning,0.8'); % FBP of the sinogram without energy effects  
    
% Sinogram for sparse data
wsino = G*wtrue;
bsino = G*btrue;
asino=G*atrue;
ssino = cat(3, wsino, bsino,asino);

% Polychromatic variables
sls1 = de_ftab_sls('max', 40, 'n', 81);
xrs = xray_read_spectra('poly1,100', 'filters', {{'aluminum', 0.25}});
sls2= de_ftab_sls('max', [6 3 3], 'n', [49 25 25],'min', [0 0 0]);

mas1 = xray_read_mac('water');

ftab1 = de_ftab(xrs, mas1, 'sls', sls1, 'ctype', 'newt', ...
    'ftype','exp', 'fit_args', {'kev', 10:5:160, 'mac', []});    

mas2 = xray_read_mac({'water', 'bone','adipose'});

ftab2 = de_ftab(xrs, mas2, 'sls', sls2, 'ctype', 'newt', ...
    'ftype','exp', 'fit_args', {'kev', 10:5:160, 'mac', []});

% Polychromatic sinogram
fsino = ftab2.fit.fmfun(ssino);
ri=0; % Scatter
f.I0 = I0; 
yb=f.I0 * exp(-fsino);
wi = poisson(yb); % Adding poisson noise
yi = -log(max((wi-ri)./f.I0,1/f.I0));
prompt

% Analytical reconstruction
% Uncorrected FBP
f.fbp = fbp2(sg, ig);
tmp = yi ; % Scale correct for effective water mac
fbpu = fbp2(tmp, f.fbp, 'window', 'hanning,0.8');
clim = [0 mu_bone]; 



%% Statistical reconstruction

% Regularization

ci=wi;
kappa = sqrt(div0(G' * ci, G' * sg.ones));

    
% Different regularization terms depending on the acquisition parameters

f.beta = 15;
R = Reg1(kappa, 'type_denom', 'matlab', 'beta', f.beta, ...
'pot_arg', {'huber', 0.005});  % 'pot_arg', {'hyper3', 0.1});
f.niter = 10;
    
% Variables for statistical method
xinit = fbpu;
mu_w=1;
bi_vol=f.I0;  
ri_vol=ri;
bi=bi_vol;

xinit_row=reshape(xinit,ig.nx*ig.ny,1);
mask_row=reshape(ig.mask,ig.nx*ig.ny,1);

display = @(i,x) im(i,reshape(embed(x,mask_row),ig.nx,ig.ny));

% Parameters for BH estimation
AOpt=1.475;
BOpt=0.0100;

% Number of subsets
nsubset =15;
Gb = Gblock(G,nsubset);
% Thresholds for segmentation functions
Th=[0 0.4 1.1 1.5];

% Statistical method 
temp= sir_bh_os(AOpt,BOpt, ftab2, xinit/ ftab2.mac.bar(1), Gb, wi,bi, ri,ftab1, R, f.niter,mask_row,Th);

% New estimate of the image
xcStatistical = reshape(temp,ig.nx,ig.ny,f.niter+1);
% figure(6)
% ult=size(xc3);
% clf, im(xc3(:,:,ult(3)))
xcStatistical_rec=xcStatistical(:,:,end);

%% Monochromatic method

% Parameters for monochromatic method
f.beta = 20;
R = Reg1(kappa, 'type_denom', 'matlab', 'beta', f.beta, ...
'pot_arg', {'huber', 0.01});  
f.niter = 10;
nsubset =8;
Gb = Gblock(G,nsubset);
      
[xs, info] = pwls_sps_os(xinit(ig.mask), yi, wi, Gb, R, f.niter);

xcMonochromatic=embed(xs,ig.mask);
xcMonochromatic_rec=xcMonochromatic(:,:,end);

%% JS method

% Parameters for JS method
A=0.335;
B=0.0405;
[ xsJosephSpital   ] = f_correctionJS( A,B,yi,ftab1,G,f );
%% Showing all the images

close all
f1=figure('Position', [200 200 900 600])

subplot(2,3,[1 4])
imshow(dens_true_rec,[0 2.1])
title('Gold standard')
subplot(2,3,2)
imshow(fbpu/ ftab2.mac.bar(1),[0 2.1])
title('FBPU')
subplot(2,3,3)
imshow(xsJosephSpital,[0 2.1])
title('Joseph Spital Correction')
subplot(2,3,5)
imshow(xcMonochromatic_rec/ ftab1.mac.bar(1),[0 2.1])
title('Statistical with no BH')
subplot(2,3,6)
imshow(xcStatistical_rec,[0 2.1])
title('Statistical with BH Correction')

