%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%  Edge-Aware Volume Smoothing Using L0 Gradient Minimization
%  Qichao Wang, Yubo Tao, Hai Lin
%  For EuroVis 2015
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%  Original code from
%       "Image Smoothing via L0 Gradient Minimization", Li Xu, Cewu Lu, Yi Xu, Jiaya Jia, 
%       ACM Transactions on Graphics, (SIGGRAPH Asia 2011), 2011. 
%       http://www.cse.cuhk.edu.hk/leojia/projects/L0smoothing/index.html
%  Adapted to our paper code.
%  Version 1.0 - March 2015.
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Main Steps:
% 1: Read 3D Volume Dataset
% 2: 3D Volume L0 Smoothing
%   Computation using Alternative Minimization
%     2.1: a shrinkage problem (Equation (7))
%     2.2: the screened Poisson equation (Equation (8))
% 3: Write 3D L0 Smoothing Volume
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
tic;

%% Read 3D Volume Dataset
% Select correct .raw file and x y z size
xsize = 128;% volume size
ysize = 128;
zsize = 128;
volumesize = xsize * ysize * zsize;
fname = 'Sphere_128128128';% volume file name
fileid = fopen([fname,'.raw'],'rb');
data = fread(fileid, volumesize, 'uint8');
fclose(fileid);
r = size(data);
if(volumesize ~= r)
    disp('readsize != volumesize');
return;
end
res = reshape(data, xsize, ysize, zsize);
res = res/255;
V = res;

%% 3D Volume L0 Smoothing
fx = [1, -1];
fy = [1; -1];
[N,M,D] = size(V);
sizeI3D = [N,M,D];
otfFx = psf2otf(fx,sizeI3D);
otfFy = psf2otf(fy,sizeI3D);

tempotf = psf2otf(fx,[D,D]);
temp = tempotf(1,1:end);
otfFz = zeros(sizeI3D);
for i=1:D
    otfFz(:,:,i) = otfFz(:,:,i) + temp(i);
end

Normin1 = fftn(V);
Denormin2 = abs(otfFx).^2 + abs(otfFy ).^2 + abs(otfFz ).^2;

kappa = 1.414;
%The critical parameter lambda is allowed to be adjusted to effectively
%control the level of structure coarseness.
lambda = 0.002;
betamax = 1e5;

beta = 2*lambda;
% Computation using Alternative Minimization
while beta < betamax
    Denormin   = 1 + beta*Denormin2;
    
    % Equation (7)
    h = [diff(V,1,2), V(:,1,:) - V(:,end,:)];%x direction
    v = [diff(V,1,1); V(1,:,:) - V(end,:,:)];%y direction
    u = zeros(xsize,ysize,zsize);%z direction
    u(:,:,1:end-1) = diff(V,1,3); u(:,:,end) = V(:,:,1)-V(:,:,end);

    t = (h.^2 + u.^2 + v.^2)<lambda/beta;

    h(t) = 0; v(t) = 0; u(t) = 0;
    
    % Equation (8)
    Normin2 = [h(:,end,:) - h(:,1,:), -diff(h,1,2)];%x direction
    Normin2 = Normin2 + [v(end,:,:) - v(1,:,:); -diff(v,1,1)];%y direction

    temp2 = -diff(u,1,3);%z direction
    temp1 = u(:,:,end) - u(:,:,1);
    tempu = zeros(size(u));
    tempu(:,:,1) = temp1;tempu(:,:,2:end) = temp2;
    Normin2 = Normin2 + tempu;

    FV = (Normin1 + beta*fftn(Normin2))./Denormin;
    V = real(ifftn(FV));

    beta = beta*kappa;
    fprintf('.');
end
fprintf('\n');

%% Write 3D L0 Smoothing Volume
V = floor(V*255);
writefield = fopen([fname,'_lambda_',num2str(lambda),'.raw'],'wb');
fwrite(writefield, V, 'uint8');
fclose(writefield);

toc;