%% Get directory to folder with images

path = [];
[path] = uigetdir('C:\Users\oderm\Desktop\Revision_QPI\MatlabCode_GitHub\Data\wt_30C'); % select folder with images
if isequal(path,0)
   disp('User selected Cancel');
else
   disp(['User selected ', fullfile(path)]);
end

cd (path)   % go to directory
files = [];
files = dir('*tif');  % get information of image files stored as *tif

z_stack =[];
z_stack = 25; % indicate how many images per stack were recorded

number_timepoints = [];
number_timepoints = floor(size(files,1)/z_stack); % calculate number of timepoints in folder

%
std_all = [];
images = [];
for n = 0:number_timepoints-1
    im = [];
    q=1;
    for i = (n*z_stack)+1:(n*z_stack)+z_stack
        im = imread(files(i).name);
        std_all(q,n+1) = std2(im); % get standard deviation for each image
        images{q,n+1} =  files(i).name; % get list of corresponding filenames
        q=q+1;
        disp([n i]);
    end
    clear im;
end

% get images with minimal standard deviation (in-focus plane) in each z-stack
min_std = [];
infoc_filename = [];
for j = 1:size(std_all,2)
    [M,I] = min(std_all(:,j));
    min_std(j,1) = I+(z_stack*(j-1));
    min_std(j,2) = I;
    infoc_filename{j,1} = images{I,j};
end

% TIE QPI

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% DEMO: Phase Imaging via Transport-of-Intensity Approach
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
z_step = [];
z_step = 2; % for analysis take every second image, corresponding to 500 nm between images in z-stack

for n=1:size(infoc_filename,1)
home;
clearvars -except path infoc_filename n min_std files z_step analysed
close all;

fprintf('-------------------------------------------------------\n');
fprintf('Phase Imaging via Transport-of-Intensity Approach      \n');
fprintf('-------------------------------------------------------\n');
fprintf('Variations of intracellular density during the cell    \n');
fprintf('cycle arise from tip-growth regulation in fission yeast\n')
fprintf('-------------------------------------------------------\n\n');

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Compute measurements
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
cd (path)
mkdir results % create folder to save results in
cd results
resdir = pwd;
cd 'C:\Users\oderm\Desktop\Revision_QPI\MatlabCode_GitHub' % change this to the directory containing 'solvers' and 'utils'

data.id = 'yeast_phase';
    
%% load images for analysis. Here use 3 images above and 3 below the focal image, with a z-step of 500 nm

% get relevant filenames
%     f_5=i(min_std(n,1)-(z_step*5),1).name;
%     f_4=files(min_std(n,1)-(z_step*4),1).name;
    f_3=files(min_std(n,1)-(z_step*3),1).name;
    f_2=files(min_std(n,1)-(z_step*2),1).name;
    f_1=files(min_std(n,1)-(z_step*1),1).name;
    f_0=files(min_std(n,1),1).name;
    f1=files(min_std(n,1)+(z_step*1),1).name;
    f2=files(min_std(n,1)+(z_step*2),1).name;
    f3=files(min_std(n,1)+(z_step*3),1).name;
%     f4=files(min_std(n,1)+(z_step*4),1).name;
%     f5=files(min_std(n,1)+(z_step*5),1).name;

% load brightfield images
% Im1 = im2double(imread([path '\' (f_5)]));
% Im2 = im2double(imread([path '\' (f_4)]));
Im3 = im2double(imread([path '\' (f_3)]));
Im4 = im2double(imread([path '\' (f_2)]));
Im5 = im2double(imread([path '\' (f_1)]));
I0  = im2double(imread([path '\' (f_0)]));
Im7 = im2double(imread([path '\' (f1)]));
Im8 = im2double(imread([path '\' (f2)]));
Im9 = im2double(imread([path '\' (f3)]));
% Im10 = im2double(imread([path '\' (f4)]));
% Im11 = im2double(imread([path '\' (f5)]));


fprintf('::: Data is loaded!\n');

% all axial derivatives
% y5 = (Im11 - Im1)./I0;
% y4 = (Im10 - Im2)./I0;
y3 = (Im9 - Im3)./I0;
y2 = (Im8 - Im4)./I0;
y1 = (Im7 - Im5)./I0;

% choose small and large defocus measurements

% defocus distances
data.dzNear = 0.5e-6; % distance between images 500 nm
data.dzFar  = 1.5e-6; % distance to most out of focus image 1500 nm

% observations indicat which axial derivative is closest and most distant
data.bNear = y1;
data.bFar  = y3;

% data dimensions
data.dims = size(data.bNear);

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Physical parameters
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% wavelength in meters (?)
data.nu = 630e-9;

% magnification
data.mag = 60;

% pixel size in meters
data.px = 108.3e-9;

% numerical aperture
data.NA = 1.4;

% wave number (by definiton)
data.k = 2*pi/data.nu;

% maximal frequency
data.fmax = 1/(2*max(data.px/data.mag,data.nu/(2*data.NA)));

% Data output calculation (reconstruction)
% maximal radian 
maxRadian = 30;

% minimal radian
minRadian = -10;

% image bit size 16
bitsize = 16;
maxIntensity = (2^(bitsize))-1;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Experiment settings
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

%% Output options

% log experiment results?
VERBOSE = true;

% save experiment results?
SAVEMAT = false;

% save display outputs?
SAVEDISP = false;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Reconstruction settings
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% regularization parameter (can be an array)
LAMBDA = [2.5e-3];

% correction value for the singularity at the origin
data.sval = 1;

% number of iterations
data.T = 250;

% relative-error tolerance
data.tol = 1e-12;

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Initializations
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% initialize CPU clock
tinit = tic;

% generate necessary paths
if ~exist('NOSYNC','dir')
    mkdir NOSYNC
end

% include necessary paths
addpath(genpath(pwd));

% time stamp
stamp = datestr(now, 'yyyy_mmm_dd_HH_MM');

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Reconstruction
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

% construct filters
[data.W_low,data.W_high] = genFourierWeights(data.dims(2),...
    data.dims(1),...
    data.dzNear,...
    data.mag,...
    data.nu,...
    data.fmax,...
    5e-1,...
    pi/10,...
    true);

% forward operators (as Fourier multipliers!)
data.tieOtfNear = tieOtf(data.dims(2),...
    data.dims(1),...
    data.fmax,...
    data.nu,...
    data.dzNear);

data.tieOtfFar = tieOtf(data.dims(2),...
    data.dims(1),...
    data.fmax,...
    data.nu,...
    data.dzFar);

fprintf('::: WTIE-TV Reconstruction \n');

PHI_TIE = reconPhaseTieTvWeightedBatch(data,LAMBDA,VERBOSE);

for ii=1:length(LAMBDA)
    
    % get the regularization parameter
    lambda = LAMBDA(ii);
    
    % get the full FOV
    phi = PHI_TIE{:,ii};
   
    cd (resdir);
    
    % adjust dynamic range for saved images
    FinalImage=uint16((phi+abs(minRadian))*maxIntensity/(maxRadian-minRadian));

    imwrite(FinalImage,['Tie_' sprintf('%03d',n-1) '_'  sprintf('%03d',min_std(n,2))  '_' num2str(minRadian) '_' num2str(maxRadian) '.tif']);
    analysed{n,1}=['Tie_' sprintf('%03d',n-1) '_'  sprintf('%03d',min_std(n,2))  '_' num2str(minRadian) '_' num2str(maxRadian) '.tif'];
    cd 'C:\Users\oderm\Desktop\Revision_QPI\MatlabCode_GitHub\'
    if SAVEDISP
        fname = [pwd '/NOSYNC/results_tieVsDhm/tieVsDhm_results_' stamp num2str(ii)];
        export_fig(fname);
        close all;
    end
end

%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% Save results
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if SAVEMAT
    data.fid = sprintf('NOSYNC%MAT%sstieRes_%s_%s',filesep,filesep,data.id,stamp);
    % save results
    save(sprintf('%s.mat',data.fid));
    close all;
end
stime = toc(tinit);
fprintf('\n');
fprintf('-------------------------------------------------------\n');
fprintf('-------------------------------------------------------\n');
fprintf('Total reconstruction time is %.3f seconds.\n',stime);
fprintf('Run next code section for background correction.\n');
fprintf('-------------------------------------------------------\n');
fprintf('-------------------------------------------------------\n');
end
close all;
cd (resdir)
clearvars -except path infoc_filename n min_std files z_step analysed maxIntensity minRadian maxRadian
% EOF

%% background correction for analysed images to shift background intensity to 0 radians
% here 0 radians corresponds to intensity of 16384; -10 rad = 0; 30 rad = 2^16; 0 rad = (2^16)*(10/40);

cd ([path '\results'])
analysed = dir('*tif');

% create new folder to save background corrected images
mkdir bg_corr

% calculate pixel intensity corresponding to 0 phase shift
dispRange = maxRadian - minRadian;
zeroRad = floor(maxIntensity*(-minRadian/dispRange));
    
% load all results images, make histogram and correct for background
binedges=linspace(0,65536,4097);    
for fls=1:size(analysed,1)
    RawImage=imread(analysed(fls).name);
    histo=histogram(RawImage(:,:),binedges);
    histocounts(1,1:4096) = histo.BinCounts(1,1:4096);
    smoothedHisto=smooth(smooth(histocounts(1,:),10),10);

    % find maximum
    [M,I]=max(smoothedHisto(1:4000,1));

    %% fit Gaussian on background peak
    x(1,:)=I-300:I+299;
    y(1,:)=smoothedHisto(x(1,:),1);
    [xData, yData] = prepareCurveData( x, y );

    % Set up fittype and options.
    ft = fittype( 'gauss1' );
    opts = fitoptions( 'Method', 'NonlinearLeastSquares' );
    opts.Display = 'Off';
    opts.Lower = [-Inf -Inf 0];
    opts.StartPoint = [54.8559670781893 I 76.2382828044741];

    % Fit model to data, plot fit
    [fitresult, gof] = fit( xData, yData, ft, opts );
    h = plot( fitresult, xData, yData );
    pause(0.25);

    %% correct for background;
    FinalImage=RawImage+(zeroRad-(fitresult.b1*16)); 
    cd bg_corr
    imwrite(FinalImage,[analysed(fls).name(1:end-4) '_bg_corr.tif']);
    cd ../
end
close(gcf)
cd bg_corr