%% ST5_DriftCorrect_3D()
% Revised by Q.Peter Su on Oct 24 2012
% routine to correct drift from .dat files 

%% Input
%      1. *.dat file (outputs from Merge_and_plotfiles.m)
%      2. *.mat file Driftave XYZ all in um (Micron)
%      3. frame no.

%% Output 
%      1. -Corr2D.dat file, fig file, png file
%         -Corr3D.dat file, fig file, png file
  
%% 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 Jan 25 2013
%  frames = size(Driftave,1);
%% Revised by Q.Peter Su on May 20 2013
%  ST05_DriftCorrect_3D(Afname)
%  dimensionjudge=2;%input('Correct the Drift 2D or 3D, press 2 or 3 - ');

%% Revised by Q.Peter Su on Nov 21 2013
%     plot3k([x,y,z],z,[-1000 1000],{'o',1});
%     %scatter(x,y,3,z,'filled');
%     view([0 -90])
%     title({[FileName];[FileNameD];...
%            ['DriftCorrected-' num2str(dimensionjudge) 'D Image'];...
%            ['Z range [-1000 1000], color means position']},...
%          'FontWeight','bold','FontSize',10,'FontName','Arial',...
%           'Color',[1 1 1]);
%     xlabel('X in nm');
%     ylabel('Y in nm');

%% Revised by Q.Peter Su on Jan 24 2014
% if DeltaZ==0
%     disp('Delta Z = 0, Correct the Drift in 2D ...')
%     dimensionjudge=2;
% else
%     dimensionjudge=input('Correct the Drift 2D or 3D, press 2 or 3 - ');
% end

%% Revised by Q.Peter Su on June 10 2015
%  to avoid the coordinates of dots which is less than ZERO
%  index_zero=find(XYZt(:,1)<0); XYZt(index_zero,:)=[];
%  index_zero=find(XYZt(:,2)<0); XYZt(index_zero,:)=[];

function ST05_DriftCorrect_3D(Afname,FileNameD)
% clc
close all;
%% import Driftave information
if nargin<2
    [FileNameD,PathNameD] = uigetfile({'*drift*.mat','Driftave *.mat file';...
        '*.mat','All mat file'},...
        'Select *.mat file for the Drift Correction Information',...
        'MultiSelect', 'off');
    cd(PathNameD);
    disp(PathNameD);disp(FileNameD);
end
                                    %%%%%%%%%%%%%%%%%%%%%%%
AfnameD=char(FileNameD);            %%% X Y Z all in um %%%
Driftave = importdata(AfnameD);     %%%%%%%%%%%%%%%%%%%%%%%
%% select DAT file to correct
if nargin<1
    [FileName,PathName] = uigetfile('*.dat',...
        'Select .dat file to correct STORM localizations','MultiSelect', 'off');
    cd(PathName);
    
    % import x y z
    Afname=char(FileName);
end


frames = size(Driftave,1);
disp(Afname);
tic
disp('--------------------------------------------------------------');
disp('Import Data ... ing ...');
XYZt = importdata(Afname);      %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%x = XYZt(:,1);  % x in um      %% X Y are in Micro meter, not in pixels %%
%y = XYZt(:,2);  % y in um      %%% Z is in Nano meter, not in pixels %%%%%
%z = XYZt(:,3);  % z in nm      %%%%(need to change for actual data)%%%%%%%
%t = XYZt(:,12); % frame no.    %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
% else
%     XYZt = InputXYZt;
%end
toc

    
disp('--------------------------------------------------------------');
disp('Driftave Value in um');
disp('      No.       X         Y        Z       (in um)');
disp(Driftave(1,:));
disp('              Max X     Max Y    Max Z     (in um)');
strg=[0 max(Driftave(:,2)) max(Driftave(:,3)) max(Driftave(:,4)) 0];
disp(strg);
disp('              Min X     Min Y    Min Z     (in um)');
strgg=[0 min(Driftave(:,2)) min(Driftave(:,3)) min(Driftave(:,4)) 0];
disp(strgg);
strggg=[0 max(Driftave(:,2))-min(Driftave(:,2))...
          max(Driftave(:,3))-min(Driftave(:,3))...
          max(Driftave(:,4))-min(Driftave(:,4)) 0];
disp('             Delta X   Delta Y   Delta Z   (in um)');
disp(strggg);
%% run through all the frames and apply drift correction based on frame no.
DeltaX=max(Driftave(:,2))-min(Driftave(:,2));
DeltaY=max(Driftave(:,3))-min(Driftave(:,3));
DeltaZ=max(Driftave(:,4))-min(Driftave(:,4));

disp('--------------------------------------------------------------');
if DeltaZ==0
    disp('Delta Z = 0, Correct the Drift in 2D ...')
    dimensionjudge=2;
else
    dimensionjudge=input('Correct the Drift 2D or 3D, press 2 or 3 - ');
end


disp('...Drift Correcting ...ing...')
tic
if DeltaX==0 && DeltaY==0 && DeltaZ==0
    disp('Delta X = 0, Delta Y = 0, Correct Two Color Drift ...')
    XYZt(:,1)=XYZt(:,1)-Driftave(1,2);
    XYZt(:,2)=XYZt(:,2)-Driftave(1,3);
    XYZt(:,3)=XYZt(:,3)-Driftave(1,4);
else
    if dimensionjudge==2
        for i = 1:frames
            ind = find(XYZt(:,12) == Driftave(i,1));
            XYZt(ind,1) = XYZt(ind,1) - Driftave(i, 2);       % Correct X in um
            XYZt(ind,2) = XYZt(ind,2) - Driftave(i, 3);       % Correct Y in um
            %XYZt(ind,3) = XYZt(ind,3) - Driftave(i, 4)*1000;  % Correct Z in nm
            if mod(i,5000)==0
            dispFrame=strcat('Frame number = ',num2str(frames-i));
            disp(dispFrame);    
            end
        end
    elseif dimensionjudge==3
        for i = 1:frames
            ind = find(XYZt(:,12) == Driftave(i,1));
            XYZt(ind,1) = XYZt(ind,1) - Driftave(i, 2);       % Correct X in um
            XYZt(ind,2) = XYZt(ind,2) - Driftave(i, 3);       % Correct Y in um
            XYZt(ind,3) = XYZt(ind,3) - Driftave(i, 4)*1000;  % Correct Z in nm
            if mod(i,5000)==0
            dispFrame=strcat('Frame number = ',num2str(frames-i));
            disp(dispFrame);    
            end
        end
    end
end
toc
%% Revised by Q.Peter Su on June 10 2015
%  to avoid the coordinates of dots which is less than ZERO
index_zero=find(XYZt(:,1)<=0); XYZt(index_zero,:)=[];
index_zero=find(XYZt(:,2)<=0); XYZt(index_zero,:)=[];

%% Creat figure
close all
figure1 = figure('Color',...
    [0.313725501298904 0.313725501298904 0.313725501298904]);
% Create axes
axes1 = axes('Parent',figure1,'ZColor',[1 1 1],'YColor',[1 1 1],...
    'XColor',[1 1 1],...
    'Color',[0 0 0]);
hold(axes1,'all');
x=XYZt(:,1)*1000;
y=XYZt(:,2)*1000;
z=XYZt(:,3);
disp('...Plotting Localizations ...ing...')
tic
plot3k([x,y,z],z,[-750 750],{'o',1});
%scatter(x,y,3,z,'filled');
view([0 -90])
title({[Afname];[FileNameD];...
       ['DriftCorrected-' num2str(dimensionjudge) 'D Image'];...
       ['Z range [-750 750], color means position']},...
     'FontWeight','bold','FontSize',10,'FontName','Arial',...
      'Color',[1 1 1]);
xlabel('X in nm');
ylabel('Y in nm');
% colorbar
axis equal
toc
%% output file
if dimensionjudge==2
    outfile = [Afname(1:end-4)];
    outfilename = [outfile '-C2D.dat'];
    save(outfilename, 'XYZt', '-ascii', '-tabs');
    figurefilename1 = strcat(outfile, '-C2D.fig'); 
    saveas(gcf,figurefilename1,'fig');
    figurefilename2 = strcat(outfile, '-C2D.bmp'); 
    saveas(gcf,figurefilename2,'bmp');
elseif dimensionjudge==3
    outfile = [Afname(1:end-4)];
    outfilename = [outfile '-C3D.dat'];
    save(outfilename, 'XYZt', '-ascii', '-tabs');
    figurefilename1 = strcat(outfile, '-C3D.fig'); 
    saveas(gcf,figurefilename1,'fig'); 
    figurefilename2 = strcat(outfile, '-C3D.bmp'); 
    saveas(gcf,figurefilename2,'bmp');
end
end