
%%  Read Image

% Below X is n1 x n2 x 3; i.e. we have three n1 x n2 images, one for each of the 3 color channels
%load('10%noise_Diffraction_Pattern.mat'); h = 52; w = 28; load('model.mat');
%load('diffpat.mat'); h = 52; w = 28;
%load('dp1e8.mat'); diffpat = dp1e8; h=52; w=28;
%load('cameraman_128c.mat'); h=128; w=128;
%load('lena_128_1e8.mat'); h=128; w=128;
%load('F2Db.mat'); diffpat = F2Dc; h = 82; w=90; load('yeast_model.mat');
%load('F2Dc_500.mat'); diffpat = F2Dc_500; h=75; w=85; load('model_500.mat');
%load('nanorice1.mat'); diffpat= nanorice1.F2Dbgs; supp = nanorice1.supp; h=16; w=16;
%load('nanorice2.mat'); diffpat= nanorice2.F2Dbgs; supp = nanorice2.supp; h=20; w=20;
%h=94;w=94;

function [u_best] = my_OSS(input)
diffpat  = input.diffpat;
[n1, n2] = size(diffpat);
support  = input.support;
ntimes   = input.iterations; 

not_support = ~support;
[ind1,ind2] = find(support~=0);
x_min = min(ind1);
x_max = max(ind1);
y_min = min(ind2);
y_max = max(ind2);

x_center = ceil(n1/2); y_center = ceil(n2/2);

% filter parameter
X = 1:ntimes; filtercount = 10;
sub_iter = ntimes/filtercount;
FX=(filtercount + 1-ceil(X*filtercount/ntimes))*ceil(ntimes/(1*filtercount));
FX=((FX-ceil(ntimes/filtercount))*(2*n1)/max(FX))+(2*n1/10);
[KK,JJ] = meshgrid(1:n1,1:n2);
R2 = (KK - x_center).^2 + (JJ - y_center).^2 ;

%% OSS
Z_initial = diffpat;
stopper = find(diffpat==-1);
known = diffpat ~=-1;
Z_initial (stopper) = 0;

if isfield(input, 'beta')
    beta = input.beta;
else
    beta = 0.8;
end

if isfield(input, 'initial')
    Z_new = fftshift(fft2((input.initial)));
    Z_new = Z_new /norm(Z_new,'fro')*norm(diffpat,'fro');
else
    b = (randn(n1,n2)+0*1j*randn(n1,n2)).*support;
    b = fftshift(fft2(b));
    Z_new  = exp(1i*angle(b)) .* Z_initial;
end
u = ifft2(ifftshift(Z_new)).*support;

u_best = u;

fprintf('Done with initialization, starting loop\n')

dt=1;
errorF_best=1;
sum_dp = sum(sum(diffpat(known)));
%% Loop
for t = 1:ntimes
    u_old = u;

    Z = fftshift(fft2((u)));
    %Z_stopper = Z(stopper);
    phase = angle(Z);

    Z_new = exp(1i*phase) .* diffpat;
    %Z_new(stopper) = Z_stopper;
    %u = ifft2(ifftshift(Z_new) ) ;

    diff = Z - Z_new;
    diff(stopper) = 0;
    grad = ifft2(ifftshift(diff) ) ;
    u = u - dt* grad;

    u = real(u);
    fail_set = not_support | u<0;
    u(fail_set) = u_old(fail_set) - beta*u(fail_set);
    rec = u.*support;    
    
    % smooth u
    if mod(t,sub_iter)==1
        kfilter = exp( -R2 /(2* FX(t)^2 )); kfilter = kfilter/max(max(kfilter));
    end

    ktemp   = fftshift(fftn(u));
    ktemp   = ktemp.*kfilter;
    u = real(ifftn(ifftshift(ktemp)));

    u(support) = rec(support);
    
    if mod(t,100)==0
        u = u_best;
    end
    
    if mod(t,2)==0
        z_temp = abs( fftshift(fft2(rec)) );
        errorF = sum(abs(z_temp(known)-diffpat(known))) / sum_dp;
        
        if errorF <= errorF_best
            errorF_best = errorF;
            u_best = u;
        end
                
    end
    
    if mod(t,10)==0
        fprintf('%d. ErrorF = %f\n', t, errorF) ;
        rec = rec(x_min:x_max, y_min: y_max);
        figure(101); img( real(rec), 'real', imag(rec),'imag'); drawnow();
    end
    
end

fprintf('All done!\n')














