%% Revised by Q. Peter Su on Oct 24 2012
%% ST4_genXcorrDrift_3D.m
%        Routine to 3D bin STORM data per 'N' frames 
%        this will be used finally to calculate the Sample drift
%        Modified from Rahul's genXcorrDrift5.m
%% Input: 
%       0. Croped *.tif file of a Single Bead
%          *.tif file is generated by AvImage2.m 
%          and single bead is cropped in ImageJ 
%       1. Total Numbers of STORM movie Frames
%       2. Number of Frames Pooled Together
%          (same with the number in AvImage2.m)
%       3. Polynomial fitting factor for drift path curve
%% Output: 
%       1. Polynomial fitting curve for drift path
%       2. 3D Drift in X,Y and Z axis all in um ( Micro-Meter) 
%% Revised by Q.Peter Su on Nov 12-14 2012
%  different color drift correction -- Red, Green, Rahul Roy, or without Z
%% Revised by Q.Peter Su on Nov 16 2012
%  add Drift Correction Dimension judge input
%  2D only in X Y, or 3D in X Y Z
%% Revised by Q. Peter Su on Nov 19 2012
%  Use strcmp(colorjudge,'n') to judge the string
%% Revised by Q. Peter Su on Nov 20 2012
%  scatter3 by colormap
%% Revised by Q. Peter Su on Dec 28 2012
%  Input the X and Y pixel size in nm


%% 
function Driftave = ST04_genXcorrDrift_3D(Xpixel,Ypixel)
clc
clear all;
close all;
%% Data import
[FileName,PathName] = uigetfile('*.tif',...
       'Select Croped *.tif file of a Single Bead','MultiSelect', 'off');
cd(PathName);
     frames = input('Enter Total Numbers of Frames of Movie.   ');
        frm = input('Enter Number of Frames Pooled Together.   ');
offset = 400*(input('Enter the First Movie no. of the Series.  ')-1);
Afname=char(FileName);
disp('Import Data ...ing...');
Images=load_images4(Afname);

%% Change the cells to matrices
[row noofframes]  = size(Images);
[M N]= size(Images{1,1});
for hp = 1:noofframes
    Images{1,hp} = Images{1,hp}/max(Images{1,hp}(:));
end
II = reshape(cell2mat(Images), M, N, noofframes);

%% Cross-Correlation
% calculate the xcorr between the two adjacent 2D histograms 
    window = 6;
    CC = cell(1, noofframes);
    CCd = cell(1, noofframes);
    h = fspecial('gaussian', 3, 1);
%% delete former XCC.tif file    
    %if exist('XCC.tif', 'file')
    %  delete('XCC.tif');
    %end
%% Creat the XCC.tif file
  for m = 1:noofframes
      ccn = normxcorr2(II(:,:,1),II(:,:,m));
      CC{1,m} = ccn(M-window:M+window,N-window:N+window);
      CCd{1,m} = imfilter(CC{1,m},h);
      xccfilename=strcat(Afname(1:end-4), '-XCC.tif');
      imwrite(uint16(CCd{1,m}*10000), xccfilename,'tif',...
              'Compression', 'none', 'WriteMode', 'append');
      % revised by Q.Peter Su on Oct 24 2012
  end
  
%% Fitting of X-correlation images to gaussians
CC1 = cellfun(@(x) uint16(x*20000), CC, 'UniformOutput', false);
Gdata = cell(size(CC1));
denoisedCC = CC1;

for j = 1:length(Gdata)    
% Denoising by background substraction and weiner filter
    background = imopen(CC1{j},strel('disk',3));
    CC1{j} = CC1{j} -background;
    denoisedCC{j} = CC1{j} - imopen(CC1{j},strel('disk',3));
    denoisedCC{j} = imfilter(denoisedCC{j},h);
    [meQ, seQ] = image_background(denoisedCC{j});
    SNR = 30;
    InsT = meQ+ SNR*seQ;
    fitParams = struct('curvatureThreshold',   10000, ...
                       'intensityThreshold',   InsT, ...
                       'windowSize',           7, ...
                       'minWidth',             0, ...
                       'maxWidth',             7, ...
                       'angle',                -0 * pi/180);
            %negative because imagesc displays the y-axis inverted
            
     Gdata{j} = [window+1, window+1];
     %Gdata{j} = fit_spots1S(double(denoisedCC{j}), fitParams);
     Gdata{j} = fit_spots2S(double(CC1{j}), 'gaussian', fitParams, Gdata{j});
     strg=['Frame No. : ', num2str(j), ' / ', num2str(frames/frm)];
     disp(strg);   % revised by Q. Peter Su on oct 18 2012
                   % Revised by Q. Peter Su on Nov 16 2012
end

%% Data check and reshape 
% to make sure only one gaussian is selected per frame
% Data check to ensure only one peak is selected 
corr_offset = zeros(noofframes,5);
[Cnrows,Cncols] = cellfun(@size, Gdata,  'UniformOutput', false);
irregular = find(cell2mat(Cnrows)~=1);

if ~isempty(irregular)
    sss = strcat('Correlation Frames with irregular number of peaks', num2str(irregular));
    display(sss);
end

for m =1:length(irregular)
       if Cnrows{1,irregular(m)}~=0
           peakdis = sqrt((Gdata{1,irregular(m)}(:,1)-(window+2)).^2 + (Gdata{1,irregular(m)}(:,2)-(window+2)).^2);
           [sp, ix] = min(peakdis);
           Gdata{1,irregular(m)} = Gdata{1,irregular(m)}(ix,:);
       else
           Gdata{1,irregular(m)} = 12*ones(1,8);
       end
       
end
%% Data reshape
 G = cell2mat(Gdata);
 fbins = length(Gdata);
 RG = reshape(G, 8, fbins)';
 RGt = linspace(1, fbins, fbins);
%% X and Y Pixel Size
if nargin<2 % Revised by Q.Peter Su on Jan 08 2013
    Xpixel    =input('Please input X pixel size (in nm)   - ');
    Ypixel    =input('Please input Y pixel size (in nm)   - ');
end 
%% Z Calculation and Information 
 % Revised by Q. Peter Su on Nov 12 2012
disp('Z pixel size : Press r for Red Color, g for Green Color.  ');
colorjudge=input('rr for Rahul Roy, n without Z value.- ','s');
if strcmp(colorjudge,'r')
    Zcalfile = 'F:\Q. Peter Su\0-Matlab codes for STORM\Calibration Values Red.dat';
    ZCalibration = importdata(Zcalfile);
    p = polyfit (ZCalibration(:,2),ZCalibration(:,1),3);
    z_offcal = p(1)*(RG(:,4)).^3+...
               p(2)*(RG(:,4)).^2+...
               p(3)*(RG(:,4)).^1+...
               p(4);                  % in nm
end
if strcmp(colorjudge,'g')
    Zcalfile = 'F:\Q. Peter Su\0-Matlab codes for STORM\Calibration Values Green.dat';
    ZCalibration = importdata(Zcalfile);
    p = polyfit (ZCalibration(:,2),ZCalibration(:,1),3);
    z_offcal = p(1)*(RG(:,4)).^3+...
               p(2)*(RG(:,4)).^2+...
               p(3)*(RG(:,4)).^1+...
               p(4);                  % in nm
end
if strcmp(colorjudge,'rr')
    Zcalfile = 'F:\Q. Peter Su\0-Matlab codes for STORM\Calibration Values.dat';
    ZCalibration = importdata(Zcalfile);
    p = polyfit (ZCalibration(:,2),ZCalibration(:,1),3);
    z_offcal = p(1)*(RG(:,4)).^3+...
               p(2)*(RG(:,4)).^2+...
               p(3)*(RG(:,4)).^1+...
               p(4);                  % in nm
end

if strcmp(colorjudge,'n')
    z_offcal = 0*(RG(:,4));
end

%% RGf matrix       
RGf = [ RG z_offcal RGt'];
%% X Y and Z Drift calculation  
corr_offset = [RGf(:,1)-RGf(1,1) RGf(:,2)-RGf(1,2) RGf(:,9)-RGf(1,9) RGf(:,10)];
%% median filter data
% They are delta X Y and Z
x_off = Xpixel/1000*medfilt1(corr_offset(:,1),15);        % in um
y_off = Ypixel/1000*medfilt1(corr_offset(:,2),15);        % in um
z_off = 0.001*medfilt1(corr_offset(:,3),15);              % in um
t = [fix(frm/2)+1:frm:frames+1];
%% Polynomial Fitting revised by Q. Peter Su on Oct 19 2012
% Add polynomial factor input on Oct 19 2012

judge='n';
while judge == 'n'
  polynomial=input('Enter the factor for Polynomial fitting.  ');
  [p,sx] = polyfit(t',x_off,polynomial);
  [fx,errorx] = polyval(p,t',sx);
  errorX=sqrt(sum(errorx.^2)/length(errorx));
  
  [p,sy] = polyfit(t',y_off,polynomial);
  [fy,errory] = polyval(p,t',sy); 
  errorY=sqrt(sum(errory.^2)/length(errory));
  
  [p,sz] = polyfit(t',z_off,polynomial);
  [fz,errorz] = polyval(p,t',sz); 
  errorZ=sqrt(sum(errorz.^2)/length(errorz));
  
  ti = [1:frames];
  xi = interp1(t,fx,ti, 'cubic');
  yi = interp1(t,fy,ti, 'cubic');
  zi = interp1(t,fz,ti, 'cubic');
  
  
figure1 = figure('Color',...
    [0.313725501298904 0.313725501298904 0.313725501298904]); 
  axes1 = axes('Parent',figure1,'ZColor',[1 1 1],'YColor',[1 1 1],...
         'XColor',[1 1 1],...
         'Color',[0 0 0]);
  hold(axes1,'all');
  scatter3(x_off,y_off,z_off,5,t','filled'),view(-60,60)
  hold on
  
  plot3(xi,yi,zi,...
      'Parent',axes1,'MarkerFaceColor',[1 0 0],...
      'MarkerEdgeColor',[1 0 0],...
      'Marker','.',...
      'LineWidth',3,...
      'Color',[1 0 0]);
  title({[FileName(1:end-4)];...
         ['No. of Frames Pooled Together = ' num2str(frm)];...
         ['Polynomial Factor = '  num2str(polynomial)];...
         ['ErrorX = ' num2str(errorX)];...
         ['ErrorY = ' num2str(errorY)];...
         ['ErrorZ = ' num2str(errorZ)];},...
         'FontWeight','bold','FontSize',14,'FontName','Arial',...
         'color',[1,1,1]);
  xlabel('X in um');
  ylabel('Y in um');
  zlabel('Z in um');
  hold off
  judge=input('Are you happy with the fitting?  ','s');
end

%% generate driftave
dimensionjudge=input('Generate the Driftave in 2D or 3D, press 2 or 3?  ');
if dimensionjudge==2
    Driftave = zeros(frames, 5);
    Driftave(:,1) = ti+offset;
    Driftave(:,2) = xi;   %in um
    Driftave(:,3) = yi;   %in um
    %Driftave(:,4) = zi;  %in um
    Driftave(:,4) = 0;    %in um
elseif dimensionjudge==3
    Driftave = zeros(frames, 5);
    Driftave(:,1) = ti+offset;
    Driftave(:,2) = xi;   %in um
    Driftave(:,3) = yi;   %in um
    Driftave(:,4) = zi;   %in um
end
%% auto save file
saveas(gcf, [FileName(1:end-4) '-Poly=' num2str(polynomial) '-drift.fig'], 'fig');
save([FileName(1:end-4) '-Poly=' num2str(polynomial) '-drift.mat'],'Driftave');
end
