%========================================================================
%Fractional gray value Cahn-Hilliard inpainting in 2D, Version 1.0
%Copyright(c) 2015 Jessica Bosch
%All Rights Reserved.
%
%This is an implementation of a fractional inpainting model based on the 
%vector-valued Cahn--Hilliard equation in two spatial dimensions.
%Please refer to the following paper:
%
%J. Bosch, and M. Stoll, "A fractional inpainting model based on the 
%vector-valued Cahn--Hilliard equation" 
%Preprint MPIMD/15-04, Max Planck Institute Magdeburg, March 2015.
%
%This is an example code for inpainting of a simple stripe image with a 
% relative large inpainting gap, see Section 6.1 in the paper above.
%
%Kindly report any suggestions or corrections to bosch@mpi-magdeburg.mpg.de

close all;
clear all;

format long g;

%%
start = tic;        % computational costs for the whole computation

%% read initial data:
% N1 = number of mesh points in x-direction
% N2 = number of mesh points in y-direction
% nphase = number of phases (gray values)
% gray_values = vector (size: nphase x 1) that contains the gray values
% orig_im = given original image (size: N1 x N2)
% u_0 = initial phase variable (size: nphase x N1 x N2)
% D = damaged parts (size: N1 x N2)
%     D(i,j)=0 -> image information available
%     D(i,j)=1 -> damaged parts
%%%
load Stripes.mat

%% model parameters
      
tau=1;              % time step size
omega=10^(6);       % fidelity parameter
C2=3*omega;         % constant for the convexity splitting
zeta=1.5;           % fractional power

mesh=1/N1;       % mesh size
a=0;                % left boundary node
bx=N1*mesh;      % right boundary node in x-direction
by=N2*mesh;      % right boundary node in y-direction

Lx=bx-a;
Ly=by-a;

%   epsilon = interface parameter for the 1st simulation - smoothing
%             (relative large value)
%   epsilon_small = interface parameter for the 2nd simulation - sharpening
%                    (smalle value, proportional to the mesh size)
epsilon=10;
epsilon_small=mesh;
C1=3/epsilon;       % C1 = constant for the convexity splitting


%% matrices

% derivative of the smooth potential
W=zeros(nphase,N1,N2);

% current solution
u_neu=zeros(nphase,N1,N2);

% solution from the previous time step
u_alt=u_0;    
    
% 1D wavenumbers
[kx,ky]=meshgrid((0:N2-1)*pi/Ly,(0:N1-1)*pi/Lx);

% 2D eigenvalues of the fractional Neumann Laplacian
M=(kx.^2+ky.^2).^(zeta/2);

%% start the iteration

it = 0;                 % time step
eps_switch=0;           % marker for the switch from epsilon to 
                        % epsilon_small

energy_tol_1=2e-4;      % tolerance for 1st stopping criterion
energy_tol_2=2e-5;      % tolerance for 2nd stopping criterion

norm_sol=1;             % 1st term for the norm in the stopping criterion
aux_var=0;              % 2nd term for the norm in the stopping criterion

while (it<2)
%while  (eps_switch<0.5 || (eps_switch>0.5 && sqrt(norm_sol)>energy_tol_2*sqrt(aux_var)))
    
   % switch to epsilon_small
    if ((sqrt(norm_sol)<energy_tol_1*sqrt(aux_var)) && (eps_switch<0.5))
        epsilon=epsilon_small;
        eps_switch=1;  
        C1=3/epsilon;
    end
   

   % calculate the derivative of the smooth potential
    for j=1:nphase 
        W(j,:,:)=u_alt(j,:,:).^(3)-(3/2)*u_alt(j,:,:).^(2)+0.5*u_alt(j,:,:);
    end
    
    W_sum=zeros(1,N1,N2);
    for j=1:nphase 
        W_sum = W_sum + W(j,:,:);
    end
    
    for j=1:nphase 
        W(j,:,:)=W(j,:,:)-(1/nphase)*W_sum;
    end
    
 
   % calculate the fidelity term     
    tmp=u_0;
    tmp2=u_alt;    

    for k=1:N1
        for l=1:N2
            if (D(k,l)>0.5)
                tmp(:,k,l)=0;              
                tmp2(:,k,l)=0;                   
            end
        end
    end

    P=omega*tmp+C2*u_alt-omega*tmp2;
    
    
   % compute the current solution using the DCT      
    for j=1:nphase 
        u_neu(j,:,:)=idct2(squeeze((M.*((-1/epsilon)*dct2(squeeze(W(j,:,:)))+C1*dct2(squeeze(u_alt(j,:,:))))+dct2(squeeze(P(j,:,:)))+(1/tau)*dct2(squeeze(u_alt(j,:,:))))./((1/tau)+epsilon*M.^(2)+C1*M+C2)));
    end


   % compute the norm for the stopping criterion             
    norm_sol=0;
    aux_var=0;
    for j=1:nphase  
        norm_sol=norm_sol+norm(squeeze(u_alt(j,:,:)-u_neu(j,:,:)),2)^(2);
        aux_var=aux_var+norm(squeeze(u_alt(j,:,:)),2)^(2);
    end
      
   % update the old solution 
    u_alt=u_neu;
       
   % update the time step 
    it = it+1;   
       
end

%%
% create the computed solution image data
final_im=imadd(gray_values(1)*squeeze(u_alt(1,:,:)),gray_values(2)*squeeze(u_alt(2,:,:)));
for k=3:nphase
    final_im=imadd(final_im,gray_values(k)*squeeze(u_alt(k,:,:)));
end
save(sprintf('Stripes-solution-%d', it),'final_im','it');

%% save figures of the reconstructed image

h1=figure(1);    
set(h1,'visible','off');
imshow(final_im);       
saveas(h1,sprintf('CH_stripes-%d.fig', it));  
saveas(h1,sprintf('CH_stripes-%d.eps', it));
close(h1);
    

%% print results:
% used model parameters
% CPU = CPU time
% PSNR and SSIM value regarding slice 1

fileID = fopen('Stipes_results.txt','a');
fprintf(fileID,'iter %d CPU %f norm_stop %e norm_U %e dofs %d omega %f eps %f eps_small %f tau %f C1 %f C2 %f phases %d frac %e\n', it, toc(start), sqrt(norm_sol)/sqrt(aux_var), norm_sol, nphase*N1*N2, omega, epsilon, epsilon_small, tau, C1, C2,nphase,zeta);
fprintf(fileID,'PSNR %f\n', 20*log10(max(max(abs(orig_im)))/(sqrt(mean(mean((orig_im-final_im).^2))))));
fprintf(fileID,'SSIM %f\n', ssim(orig_im,final_im, [0.01 0.03], fspecial('gaussian', 11, 1.5), 1));
fclose(fileID);