clear all;

% load raw data
[filename, path_data] = uigetfile('*.tif','Choose any one of raw data');
View = ReadTifStack(strcat(path_data, filename));
N = size(View);

% load PSF data 
[filename, path_psf] = uigetfile('*.tif','Choose any one of PSF image',path_data)
FullPSF = double(ReadTifStack(strcat(path_psf, 'PSFA.tif')));

% set parameters
dlg_title = 'Set Parameters';
prompt = {'Enter pixel size (microns):','Enter light sheet thickness FWHM (microns):', 'Enter angle of the light sheet relative to coverslip', 'Enter the increment delta along the mirror in the x direction (pixels)','Enter iteration number:', 'Enter time points to be processed'};
num_lines = 2;
defaultans = {'0.1625','1.6','45','1','20','0-10'};
answer = inputdlg(prompt,dlg_title,num_lines,defaultans);
pixel = str2num(answer{1});
thickness = str2num(answer{2});
angle = str2num(answer{3});
interval = str2num(answer{4})
iter_num = str2num(answer{5});
timepoints = strsplit(answer{6},'-');
t1 = str2num(timepoints{1});
t2 = str2num(timepoints{2});

% interval number 
ny = N(2);
nx_img = N(1);
nz_img = N(3)/interval;

N_Image = [ny, nx_img, nz_img];
nx_obj = max(nx_img, nz_img * interval);
nz_obj = nx_obj;
N_Obj = [ny, nx_obj, nz_obj];

[ny_PSF, nx_PSF, nz_PSF] = size(FullPSF);
px_range = (nx_PSF - max(N_Obj(2),256))/2+1 : (nx_PSF + max(N_Obj(2),256))/2;
py_range = (ny_PSF - max(N_Obj(1),256))/2+1:(max(N_Obj(1),256)+ny_PSF)/2;
pz_range = (nz_PSF - max(N_Obj(3),256))/2+1 : (nz_PSF + max(N_Obj(3),256))/2;
PSF = FullPSF(py_range, px_range, pz_range)/max(FullPSF(:))*65535 + 0.001;
PSF = PSF/sum(PSF(:));
N_PSF = size(PSF);

g = gpuDevice(1); wait(g); g.FreeMemory
gpu_OTF_A = fftn(circshift(gpuArray(PSF),-floor(N_PSF/2)));
disp(['OTF Done!  GPU Memory: ',num2str(g.FreeMemory)]);

N_Exc = [ny, 3*nx_obj, 3*nz_obj];
mode = 0;
Excitation = IlluminationPattern(thickness,N_Exc,mode, pixel, angle);

WriteTifStack(Excitation, strcat(path_psf, 'Excitation.tif'), 32);

gpu_Excitation = gpuArray(single(ReadTifStack(strcat(path_psf,'Excitation.tif'))));

disp(['Excitation Done!  GPU Memory: ',num2str(g.FreeMemory)]);

%calculationg the sensitivity of A 
SensA = zeros([ny, nx_obj, nz_obj],'single', 'gpuArray');  
for z = 1 : nz_img
        delta = (z - nz_img/2) * interval;
        InterMediate = zeros([ny, nx_obj, nz_obj],'double', 'gpuArray');  
        InterMediate(:,:,delta + nz_obj/2) =  ones([ny,nx_obj],'double','gpuArray');
         px_range = N_Exc(2)/2 - N_Obj(2)/2 + 1 - delta : N_Exc(2)/2 + N_Obj(2)/2 - delta;
        pz_range = N_Exc(3)/2 - N_Obj(3)/2 + 1 - delta : N_Exc(3)/2 + N_Obj(3)/2 - delta;  
        Sample_Illuminate = ConvFFT3_S(InterMediate,conj(gpu_OTF_A)).*gpu_Excitation(:,px_range, pz_range); 
        SensA = SensA + single(Sample_Illuminate);
end

WriteTifStack(gather(SensA), strcat(path_psf, 'SensA.tif'), 32);

Sens = gpuArray(single(ReadTifStack(strcat(path_psf,'SensA.tif'))));
disp(['Sensitivity Done!  GPU Memory: ',num2str(g.FreeMemory)]);

for time= t1:t2
time

ViewA = double(ReadTifStack(strcat(path_data,'SPIMA_',num2str(time),'.tif'))); 
ViewB = double(ReadTifStack(strcat(path_data,'SPIMB_',num2str(time),'.tif'))); 

gpu_ViewA = gpuArray(max(ViewA, 0.001));
gpu_ViewB = gpuArray(max(ViewB, 0.001));

Estimate = zeros([ny, nx_obj, nz_obj],'single', 'gpuArray') + 0.001;

disp(['GPU Memory: ',num2str(g.FreeMemory)]);

for iteration = 1:iter_num
   disp(['Processing iteration # ' num2str(iteration)]);
   t_start = clock;
    
   % View A Forward and Adjoint Model
   Correction = zeros([ny, nx_obj, nz_obj],'single', 'gpuArray');    % correction should be the same size with obj
   for z = 1 : nz_img  
        delta = (z - nz_img/2) * interval;         
        px_range = N_Exc(2)/2 - N_Obj(2)/2 + 1 - delta : N_Exc(2)/2 + N_Obj(2)/2 - delta;
        pz_range = N_Exc(3)/2 - N_Obj(3)/2 + 1 - delta : N_Exc(3)/2 + N_Obj(3)/2 - delta;    
        Sample_Blur = ConvFFT3_S(Estimate.*gpu_Excitation(:,px_range, pz_range), gpu_OTF_A);
        InterMediate = zeros([ny, nx_obj, nz_obj],'single', 'gpuArray'); 
        InterMediate(:,:,delta + nz_obj/2) = gpu_ViewA(:,:,z)./Sample_Blur(:,:,delta + N_Obj(3)/2);       
        Sample_Illuminate = ConvFFT3_S(InterMediate,conj(gpu_OTF_A)).*gpu_Excitation(:,px_range,pz_range); 
        Correction = Correction + Sample_Illuminate;
    end
        
    Estimate = Estimate.* Correction./Sens; 
     
    disp(['View A Model Done!  GPU Memory: ',num2str(g.FreeMemory)]);
    
    % View B Forward and Adjoint Model
    Correction =  zeros([ny, nx_obj, nz_obj],'single', 'gpuArray');
    for z = 1 : nz_img
        delta = (z - nz_img/2) * interval;                 
        px_range = N_Exc(2)/2 - N_Obj(2)/2 + 1 - delta : N_Exc(2)/2 + N_Obj(2)/2 - delta;
        pz_range = N_Exc(3)/2 - N_Obj(3)/2 + 1 - delta : N_Exc(3)/2 + N_Obj(3)/2 - delta;                    
        Sample_Blur = ConvFFT3_S(Estimate.*gpu_Excitation(:,px_range, pz_range), permute(gpu_OTF_A,[1,3,2]));  
        InterMediate = zeros([ny, nx_obj, nz_obj],'single', 'gpuArray');  
        InterMediate(:,delta + nx_obj/2,:) = squeeze(gpu_ViewB(:,z,:))./squeeze(Sample_Blur(:,delta + nx_obj/2,:));       
        Sample_Illuminate = ConvFFT3_S(InterMediate,conj(permute(gpu_OTF_A,[1,3,2]))).*gpu_Excitation(:,px_range, pz_range); 
        Correction = Correction + Sample_Illuminate;
    end
   
    Estimate = Estimate.* Correction./Sens; 
     
    disp(['View B Model Done! GPU Memory: ',num2str(g.FreeMemory)]);
    disp(['Iteration # ' num2str(iteration), ' takes ', num2str(etime(clock,t_start)), ' s']); 
        
end
 
Final = gather(Estimate);
FinalTop = zeros(ny,nx_img,nx_img);

for i=1:nx_img
    for k=1:nx_img
        if i<=k
            FinalTop(:,i,k) = Final(:,i,k);
        end
     end
end

WriteTifStack(FinalTop, strcat(path_data, 'Decon_',num2str(time),'.tif'), 32);

end

  




    
    
    
    
