%% Implementation of phase retrieval with P uniform/random phase illuminations
%% Images: 2D
%% Miao's oversampling technique with zero paddings
%% Wenjing Liao (wjliao@math.duke.edu)
clear all
close all
display('--------------------begin-------------------------')

%% parameters
i           = sqrt(-1);
P           = 1;  % the number illuminations
Sample2     = 4;  % total oversampling rate 
Sample1     = sqrt(Sample2); % oversampling rate in each dimension;
                  % oversampling rate at each dimension
constmask   = 0;  % Whether to use a uniform illumination. 
                  % constmask = 0: no uniform illumination
                  % constmask = 1: one uniform illumination


%% noise type
noisetype   = 0;   % 0 gaussian noise  1 poisson noise
noiselevel  = 0;   % noiselevel = 5 means 5% noise

testimage   = 4;   % 1 cameraman  % 2 virus   % 3 bull   % 4 phantom  
imagetype   = 1;   % 1 real positive   % 2  complex positive
                   % -2 complex

methodtype     = 1;   % 0: pure ER 
                      % 1: HIO followed by ER  
%% hio parameter
beta  = 0.9;

itermax = 200;      % number of iterations for HIO
ifer    = 100;      % number of iterations for ER

supptype    = 1;    % 0 tight support   1 loose support
if supptype == 0
    loose = 0;
elseif supptype == 1
    if testimage == 1
        loose = 13;
    elseif testimage == 2
        loose = 10;
    elseif testimage == 3
        loose = 10;
    elseif testimage == 4 
        loose = 0;
    end
end


%% load real positive images
if testimage == 1
    Xsingle  = imread('cameraman.tif');
elseif testimage == 2
    Xsingle  = imread('influenza.jpg');
    Xsingle  = rgb2gray(Xsingle);
elseif testimage == 3
    Xsingle  = imread('bull.jpg');
    Xsingle  = rgb2gray(Xsingle);  
elseif testimage == 4
    Xsingle  = phantom('Modified Shepp-Logan',200);
end
Xsingle  = im2double(Xsingle);

[N1 , N2]       = size(Xsingle);
if imagetype == 2
    Xsingle     = Xsingle.*exp(i*unifrnd(0,pi/2,N1,N2));   % complex positive images have random phases in [0,pi/2]
elseif imagetype == -2
    Xsingle     = Xsingle.*exp(i*unifrnd(0,2*pi,N1,N2));   % complex images have random phases in [0,2*pi]
end

%% add zero paddings to use the oversampling method by Miao. et. al.
N11       = round(Sample1 * (N1+2*loose));  % size of the image with zero paddings
N22       = round(Sample1 * (N2+2*loose));
Xdouble   = zeros(N11,N22);                 % image with zero paddings
N1start   = round(1 + (N11-N1)/2);          % position of the true image
N1end     = N1start + N1 - 1; 
N2start   = round(1 + (N22-N2)/2);   
N2end     = N2start + N2 - 1;
Xdouble(N1start:N1end,N2start:N2end) = Xsingle;


S1start   = N1start - loose;                % In the case where the image has a loose support, 
S1end     = N1end   + loose;                % only a loose position of the true image is given
S2start   = N2start - loose;
S2end     = N2end   + loose;


%% construct the illuminator
if constmask == 1
    indexstart = 2;            % the first illumination is uniform, not random
elseif constmask == 0
    indexstart = 1;
end
%% random phase
RandomPhase = zeros(N11,N22,P);
for k = indexstart : P
    RandomPhase(:,:,k) = 2*pi*rand(N11,N22);
end
%% construct the mask
Diagonal  = exp(i*RandomPhase);
  
%% true image
X  = Xdouble;


%% Fourier intensity measurement
Y           = zeros(N11,N22,P);    
Y_abs_fro   = 0;                 % Euclidean norm 
for k = 1 : P
    Y(:,:,k)            = fft2(Diagonal(:,:,k).*X);
    Y_abs_fro           = Y_abs_fro + norm(Y(:,:,k),'fro')^2;
end
Y_abs     = abs(Y);              % Fourier intensity data
Y_abs_fro = sqrt(Y_abs_fro);     % Euclidean norm 

%% Add external noise
Y_abs_noisy = zeros(N11,N22,P);
if noisetype == 0        % random gaussian noise
    sigma   = noiselevel * Y_abs_fro / (100 * sqrt(2*P*N11*N22));
    Enoise  = zeros(N11,N22,P);
    E_fro   = 0;
    for k = 1 : P
        Enoise(:,:,k) = normrnd(0,sigma,N11,N22) + i * normrnd(0,sigma,N11,N22);
        E_fro         = E_fro + norm(Enoise(:,:,k),'fro')^2;
    end
    E_fro       = sqrt(E_fro);
    Y_noisy     = Y + Enoise;     
    Y_abs_noisy = abs(Y_noisy);     % noisy Fourier intensity data
elseif noisetype == 1     % poisson noise
     if noiselevel > 0
        pscale = (100/noiselevel)^2*(sum(sum(sum(Y_abs))))/(Y_abs_fro^2);
        for k = 1 : P
            Enoise = poissrnd(pscale*Y_abs(:,:,k));
            Y_abs_noisy(:,:,k) = Enoise / pscale;
        end
        E_fro   = 0;
        for k = 1 : P
            E_fro = E_fro + norm(Y_abs(:,:,k)-Y_abs_noisy(:,:,k),'fro')^2;
        end
        E_fro       = sqrt(E_fro);
    elseif noiselevel == 0
        Y_abs_noisy = Y_abs;
        E_fro       = 0;
    end
end
%% the true noise to signal ratio
fprintf('noise to signal ratio = %6.4f percent\n',100*E_fro/Y_abs_fro)
    
%% Apply error reduction algorithm to recover X
Index  = [N1 , N2 ; N11 , N22 ; S1start , S1end; S2start , S2end]; % data to input 
%% starting point
X0  = generatestart(imagetype,Index);
%% run
if methodtype == 1     % HIO followed by ER
    [Xrec , iterres , iterchange] = HIOER(Diagonal,Y_abs_noisy,imagetype,Index,beta,X0,itermax,ifer);
elseif methodtype == 0 % pure ER
    [Xrec , iterres , iterchange] = ER(Diagonal,Y_abs_noisy,imagetype,Index,X0,itermax);
end


%% display error ||Xrec-X||/||X||
if imagetype == 1
    spaceerr  = norm(X-Xrec,'fro')/norm(X,'fro');
elseif imagetype == 2 || imagetype == -2
    theta     = globalphase(X,Xrec);  % adjust the global phase
    spaceerr  = norm(X-exp(i*theta)*Xrec,'fro')/norm(X,'fro');
end
    
fprintf('image error = %6.4f percent \n',100*spaceerr)


%% plot the magnitudes of the recovered image
figure
imshow(abs(Xrec(S1start:S1end,S2start:S2end)))
title(['image error = ' num2str(100*spaceerr) '%'])


%% plot residual curve
figure 
plot(1:length(iterres),iterres,'.-')
axis([1 length(iterres) 0 0.5])
title('relative residual at each iteration')
xlabel('iteration')
ylabel('relative residual')
    
