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);

% 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});
angleAB = str2num(answer{3});
angleCD = 90 - angleAB;
interval = str2num(answer{4});
iter_num = str2num(answer{5});
timepoints = strsplit(answer{6},'-');
t1 = str2num(timepoints{1});
t2 = str2num(timepoints{2});

ny = N(2);
nx_img = N(1);
nz_img = N(3);


N_Image = [ny, nx_img, nz_img];  %size of image
nx_obj = nx_img;          
nz_obj = nz_img;
N_Obj = [ny, nx_obj, nz_obj];  % size of object
N_PSF = N_Image;               % size of PSF
N_Exc = [ny, 2*nx_obj, nz_obj];  % size of Excitation pattern

g = gpuDevice(1); reset(g); wait(g); g.FreeMemory;

UnityLine = zeros(nz_obj,2*nx_obj,'single');
UnityLine(nz_obj/2,:) = 1.0;
UnityLineA = imrotate(UnityLine,angleAB,'bilinear','crop');
LineA = permute(repmat(UnityLineA,[1, 1, ny]),[3,2,1]);
WriteTifStack(LineA, strcat(path_psf, 'LineA.tif'), 32);
LineB = flip(LineA,3);
WriteTifStack(LineB, strcat(path_psf, 'LineB.tif'), 32);
gpu_LineA = gpuArray(LineA);
gpu_LineB = gpuArray(LineB);

UnityLine = zeros(nz_obj,2*nx_obj,'single');
UnityLine(nz_obj/2,:) = 1.0;
UnityLineC = imrotate(UnityLine,angleCD,'bilinear','crop');
LineC = permute(repmat(UnityLineC,[1, 1, ny]),[3,2,1]);
WriteTifStack(LineC, strcat(path_psf, 'LineC.tif'), 32);
LineD = flip(LineC,3);
WriteTifStack(LineD, strcat(path_psf, 'LineD.tif'), 32);
gpu_LineC = gpuArray(LineC);
gpu_LineD = gpuArray(LineD);

mode = 3;
ExcitationAB = IlluminationPattern(thickness,N_Exc,mode, pixel, angleAB);
WriteTifStack(ExcitationAB, strcat(path_psf, 'ExcitationAB.tif'), 32);
gpu_ExcitationAB = gpuArray(single(ReadTifStack(strcat(path_psf,'ExcitationAB.tif'))));

ExcitationCD = IlluminationPattern(thickness,N_Exc,mode, pixel, angleCD);
WriteTifStack(ExcitationCD, strcat(path_psf, 'ExcitationCD.tif'), 32);
gpu_ExcitationCD = gpuArray(single(ReadTifStack(strcat(path_psf,'ExcitationCD.tif'))));

FullPSFA = single(ReadTifStack(strcat(path_psf,'PSFA.tif')));
[ny_PSF, nx_PSF, nz_PSF] = size(FullPSFA);
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 : (ny_PSF + max(N_Obj(1),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(:));

FullPSFB = single(ReadTifStack(strcat(path_psf,'PSFB.tif')));
PSFB = FullPSFB(py_range, px_range, pz_range)/max(FullPSFB(:))*65535 + 0.001;
PSFB = PSFB/sum(PSFB(:));

FullPSFC = single(ReadTifStack(strcat(path_psf,'PSFC.tif')));
PSFC = FullPSFC(py_range, px_range, pz_range)/max(FullPSFC(:))*65535 + 0.001;
PSFC = PSFC/sum(PSFC(:));

FullPSFD = single(ReadTifStack(strcat(path_psf,'PSFD.tif')));
PSFD = FullPSFD(py_range, px_range, pz_range)/max(FullPSFD(:))*65535 + 0.001;
PSFD = PSFD/sum(PSFD(:));

gpu_OTF_A = fftn(circshift(single(gpuArray(PSFA)),-floor(N_PSF/2)));
gpu_OTF_B = fftn(circshift(single(gpuArray(PSFB)),-floor(N_PSF/2)));
gpu_OTF_C = fftn(circshift(single(gpuArray(PSFC)),-floor(N_PSF/2)));
gpu_OTF_D = fftn(circshift(single(gpuArray(PSFD)),-floor(N_PSF/2)));

SensA = zeros([ny, nx_obj, nz_obj],'single','gpuArray');  
for x = 1 : nx_img
        px_range = x: x + nx_img -1;           
        InterMediate = gpu_LineA(:,px_range,:); 
        Sample_Illuminate = ConvFFT3_S(InterMediate,conj(gpu_OTF_A)).* gpu_ExcitationAB(:, px_range,:); 
        SensA = SensA + single(Sample_Illuminate);
end
WriteTifStack(gather(SensA), strcat(path_psf, 'SensA.tif'), 32);

SensB = zeros([ny, nx_obj, nz_obj],'single','gpuArray');  
for x = 1 : nx_img
        px_range = x: x + nx_img -1;           
        InterMediate = gpu_LineB(:,px_range,:); 
        Sample_Illuminate = ConvFFT3_S(InterMediate,conj(gpu_OTF_B)).* gpu_ExcitationAB(:, px_range,:); 
        SensB = SensB + single(Sample_Illuminate);
end
WriteTifStack(gather(SensB), strcat(path_psf, 'SensB.tif'), 32);

SensC = zeros([ny, nx_obj, nz_obj],'single','gpuArray');  
for x = 1 : nx_img
        px_range = x: x + nx_img -1;           
        InterMediate = gpu_LineC(:,px_range,:); 
        Sample_Illuminate = ConvFFT3_S(InterMediate,conj(gpu_OTF_C)).* gpu_ExcitationCD(:, px_range,:); 
        SensC = SensC + single(Sample_Illuminate);
end
WriteTifStack(gather(SensC), strcat(path_psf, 'SensC.tif'), 32);

SensD = zeros([ny, nx_obj, nz_obj],'single','gpuArray');  
for x = 1 : nx_img
        px_range = x: x + nx_img -1;           
        InterMediate = gpu_LineD(:,px_range,:); 
        Sample_Illuminate = ConvFFT3_S(InterMediate,conj(gpu_OTF_D)).* gpu_ExcitationCD(:, px_range,:); 
        SensD = SensD + single(Sample_Illuminate);
end
WriteTifStack(gather(SensD), strcat(path_psf, 'SensD.tif'), 32);


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'))); 
ViewC = double(ReadTifStack(strcat(path_data,'SPIMC_',num2str(time),'.tif'))); 
ViewD = double(ReadTifStack(strcat(path_data,'SPIMD_',num2str(time),'.tif'))); 

gpu_ViewA = gpuArray(max(ViewA, 0.001));
gpu_ViewB = gpuArray(max(ViewB, 0.001));
gpu_ViewC = gpuArray(max(ViewC, 0.001));
gpu_ViewD = gpuArray(max(ViewD, 0.001));

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

for iteration = 1:iter_num
   disp(['Processing iteration # ' num2str(iteration)]);
   
   % View A Forward and Adjoint Model
   t_start = clock;
   
   CorrectionA = zeros([ny, nx_obj, nz_obj],'single','gpuArray');    % correction should be the same size with obj 
   for x = 1 : nx_img      
        px_range = x: x + nx_img -1;
        SampleBlur = ConvFFT3_S(gpu_Estimate.*gpu_ExcitationAB(:,px_range, :), gpu_OTF_A); 
        InterMediate = gpu_ViewA./SampleBlur.*gpu_LineA(:,px_range,:);
        CorrectionA = CorrectionA + ConvFFT3_S(InterMediate,conj(gpu_OTF_A)).*gpu_ExcitationAB(:,px_range,:);    
    end
        
    gpu_Estimate = gpu_Estimate.*CorrectionA./SensA; 
    disp(['View A Model Done!  GPU Memory: ',num2str(g.FreeMemory) '; it takes ', num2str(etime(clock,t_start)), ' s']);
     
    % View B Forward and Adjoint Model   
    t_start = clock;
     CorrectionB = zeros([ny, nx_obj, nz_obj],'single','gpuArray');    % correction should be the same size with obj 
    for x = 1 : nx_img      
        px_range = x: x + nx_img -1;
        SampleBlur = ConvFFT3_S(gpu_Estimate.*gpu_ExcitationAB(:,px_range, :), gpu_OTF_B);  
        InterMediate = gpu_ViewB./SampleBlur.*gpu_LineB(:,px_range,:);
        CorrectionB = CorrectionB + ConvFFT3_S(InterMediate,conj(gpu_OTF_B)).*gpu_ExcitationAB(:,px_range,:); 
    end
    gpu_Estimate = gpu_Estimate.*CorrectionB./SensB; 
    disp(['View B Model Done!  GPU Memory: ',num2str(g.FreeMemory) '; it takes ', num2str(etime(clock,t_start)), ' s']);
    
        % View C Forward and Adjoint Model   
    t_start = clock;
    CorrectionC = zeros([ny, nx_obj, nz_obj],'single','gpuArray');    % correction should be the same size with obj 
    for x = 1 : nx_img      
        px_range = x: x + nx_img -1;
        SampleBlur = ConvFFT3_S(gpu_Estimate.*gpu_ExcitationCD(:,px_range, :), gpu_OTF_C);  
        InterMediate = gpu_ViewC./SampleBlur.*gpu_LineC(:,px_range,:);
        CorrectionC = CorrectionC + ConvFFT3_S(InterMediate,conj(gpu_OTF_C)).*gpu_ExcitationCD(:,px_range,:); 
    end
    gpu_Estimate = gpu_Estimate.*CorrectionC./SensC; 
    disp(['View C Model Done!  GPU Memory: ',num2str(g.FreeMemory) '; it takes ', num2str(etime(clock,t_start)), ' s']);
    
    % View D Forward and Adjoint Model   
    t_start = clock;
    CorrectionD = zeros([ny, nx_obj, nz_obj],'single','gpuArray');    % correction should be the same size with obj 
    for x = 1 : nx_img      
        px_range = x: x + nx_img -1;
        SampleBlur = ConvFFT3_S(gpu_Estimate.*gpu_ExcitationCD(:,px_range, :), gpu_OTF_D);  
        InterMediate = gpu_ViewD./SampleBlur.*gpu_LineD(:,px_range,:);
        CorrectionD = CorrectionD + ConvFFT3_S(InterMediate,conj(gpu_OTF_D)).*gpu_ExcitationCD(:,px_range,:); 
    end
    gpu_Estimate = gpu_Estimate.*CorrectionD./SensD; 
    disp(['View D Model Done!  GPU Memory: ',num2str(g.FreeMemory) '; it takes ', num2str(etime(clock,t_start)), ' s']);
    
    if mod(iteration,iter_num)==0
        WriteTifStack(gather(gpu_Estimate), strcat(path_data, 'Decon_',num2str(time),'.tif'), 32);
    end
           
end

end

  




    
    
    
    
