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)
FullPSFA = double(ReadTifStack(strcat(path_psf, 'PSFA.tif')));
FullPSFB = double(ReadTifStack(strcat(path_psf, 'PSFB.tif')));

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

mode = 2;
%Excitation = IlluminationPattern(thickness, N, mode, pixel, angle);


for time = t1:t2

% loading data     
ViewA = double(ReadTifStack(strcat(path_data,'SPIMA_',num2str(time),'.tif'))); 
ViewB = double(ReadTifStack(strcat(path_data,'SPIMB_',num2str(time),'.tif'))); 
ViewA = max(ViewA,0.001);
ViewB = max(ViewB,0.001);
[ny, nx_img, nz_img] = size(ViewA);

N_Image = [ny, nx_img, nz_img];
N_Obj = N_Image;

[ny_PSF, nx_PSF, nz_PSF] = size(FullPSFA);
 
%FullPSFA = FullPSFA.*Excitation;
%FullPSFB = permute(permute(FullPSFB,[1,3,2]).*Excitation,[1,3,2]);

py_range = (ny_PSF - max(N_Obj(1),256))/2+1 : (ny_PSF + max(N_Obj(1),256))/2;
px_range = (nx_PSF - max(N_Obj(2),256))/2+1 : (nx_PSF + max(N_Obj(2),256))/2;
pz_range = (nz_PSF - max(N_Obj(3),256))/2+1 : (nz_PSF + max(N_Obj(3),256))/2;

PSFA = FullPSFA(py_range, px_range, pz_range)/max(FullPSFA(:))*65535 + 0.001;
PSFA = PSFA/sum(PSFA(:));

PSFB = FullPSFB(py_range, px_range, pz_range)/max(FullPSFB(:))*65535 + 0.001;
PSFB = PSFB/sum(PSFB(:));

N_PSF = size(PSFA);

g = gpuDevice; reset(g); wait(g); g.FreeMemory
Estimate = zeros(N_Obj,'single', 'gpuArray') + 0.001;

OTF_A = fftn(circshift(single(gpuArray(PSFA)),-floor(N_PSF/2)));
OTF_B = fftn(circshift(single(gpuArray(PSFB)),-floor(N_PSF/2)));

for iteration = 1:iter_num
   disp(['Processing iteration # ' num2str(iteration)]);
   t_start = clock;
   
   % View A Decon
   BlurA = ConvFFT3_S(Estimate, OTF_A);
   Correction = ConvFFT3_S(ViewA./BlurA,conj(OTF_A));              
   Estimate = Estimate.* Correction;         
   disp(['View A Model Done!  GPU Memory: ',num2str(g.FreeMemory)]);    
   
   % View B Decon
   BlurB = ConvFFT3_S(Estimate,OTF_B);
   Correction = ConvFFT3_S(ViewB./BlurB,conj(OTF_B)); 
   Estimate = Estimate.* Correction; 
   disp(['View B Model Done!  GPU Memory: ',num2str(g.FreeMemory)]);
    
   disp(['Iteration # ' num2str(iteration), ' takes ', num2str(etime(clock,t_start)), ' s']); 

    if mod(iteration, iter_num)==0
        WriteTifStack(gather(Estimate), strcat(path_data, 'Decon_',num2str(time),'.tif'), 32);
    end
    
end
end

