addpath src;
addpath data;

diffpat = importdata("data\yeast_diffraction.tif");
diffpat = im2double(diffpat);
diffpat = croppedOut(diffpat,[784,784]);

[N1,N2] = size(diffpat);



stopper = ones(20);
stopper = My_paddzero(stopper,[N1,N2]);

stopper = stopper & diffpat==0;


diffpat(stopper) = -1;

figure(32); img(log(abs(diffpat)),'abs','off')

%% GPS reconstruction
width  = 136;%130;
height = 136;%114;
support = ones(height, width);
support = My_paddzero(support, [N1,N2]);
support = support>0;

clear GPS_input
GPS_input.diffpat    = diffpat;
GPS_input.support    = support;
GPS_input.isreal     = true;        %whether you reconstruct a real or complex object
GPS_input.iterations = 2000;
GPS_input.sigma      = 0.01;        %parameter to tune, sigma in (0.0, 0.1]



% reconstruction
[rec_gps,y] = GPS_algorithm(GPS_input);
rec_gps_crop = croppedOut( rec_gps, [height, width]);


figure(41); img(rec_gps_crop,'GPS round 1')

%% shrinkwrap reconstruction
norm_dp1 = norm ( diffpat(~stopper), 1 );
norm_dp2 = norm ( diffpat(~stopper), 2 );
ker = fspecial('gaussian', 10,4);


u = rec_gps;
u_best  = u;
u_best0 = u;
err_best2 = 1;
err_best  = 1;
u_sum=0;
non_stopper = ~stopper;
no_supp = ~support;

thresh =0.02;
iterations=1000;
%u = u_initial;
%u = u_best;%circshift(u,[0,0]);
for t=1:iterations
    Fu = fftshift(fft2(ifftshift(u)));
    diff = abs(Fu(non_stopper)) - diffpat(non_stopper);
    z_stopper = Fu(stopper);
    Fu = diffpat .* exp(1i*angle(Fu));
    Fu(stopper) = z_stopper;

    u_new = fftshift(ifft2(ifftshift(Fu)));

    %shrinkwrap
    if mod(t,2)==0
        u_abs = abs(u_new);
        u_abs = conv2(u_abs, ker, 'same');      
        support_refined = u_abs.*support;

        support_refined = support_refined/max(support_refined(:)); % make in range [0,1]
        support_refined = support_refined > thresh;

        %u_new = u_new.*support_refined;% - 0*u_hat;
        %u_new = u_new.*support_refined;
        if mod(t,20)==0
            figure(201); img(support_refined,''); drawnow;
        end
        support = support_refined;
        no_supp = ~support;
    end

    %u_new(no_supp) = 0;
    u_new(no_supp) = u(no_supp) - 0.9*u_new(no_supp);
    %u_new = real(u_new);
    u = u_new;


    %error
    %
    Fu2  = fftshift(fft2(ifftshift(u.*support)));
    diff_z  = abs(Fu2(non_stopper))  - diffpat(non_stopper);
    err2 = norm(diff_z, 'fro') / norm_dp2;
    if err2<err_best2
        err_best2 =err2;
        u_best=u;
    end
    %}

    u_sum=u_sum+u;


    if mod(t,10)==0
        figure(4);
        rec = u_new.*support;
        img( rec, ['mag ',num2str(t)], 'colormap','gray')
        drawnow
        err1 = norm(diff_z(:),1) / norm_dp1;
        fprintf('%d. err1=%.3f, err2=%.3f\n',t,err1, err2);
    end
end


%% GPS reconstruction 2nd round
clear GPS_input
GPS_input.diffpat    = diffpat;
GPS_input.initial    = u_best;
GPS_input.support    = support;
GPS_input.isreal     = true;        %whether you reconstruct a real or complex object
GPS_input.iterations = 2000;
GPS_input.sigma      = 0.01;        %parameter to tune, sigma in (0.0, 0.1]



% reconstruction
[rec_gps,y] = GPS_algorithm(GPS_input);
rec_gps_crop = croppedOut( rec_gps, [height, width]);


figure(42); img(rec_gps_crop,'GPS 2nd round')


