%% Function to load two movies and use a pair of frames from the two
%  to create a offset factor for 3D DriftCorrection between two color dots
%  Output is a drift vector for X Y Z all in um 
       %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
       %%%%%% Drift vector is from Red bead to Green bead %%%%%%%
       %%% So Drift correction is : Green dots - Drift vector %%%
       %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Revised by Q.Peter Su on Oct 25 2012
% based on Rahul's code
%% Modified by Q.Peter Su on Nov 12-14 2012
% different color drift correction -- Red, Green, Rahul Roy or without Z
%% Modified by Q. Peter Su on Nov 19 2012
%  Use strcmp(colorjudge,'n') to judge the string
%% Revised by Q. Peter Su on Dec 28 2012
%  Input the X and Y pixel size in nm




%%
function Driftave =ST06_2colorXcorr_3D(Xpixel,Ypixel)
clc
clear all;
close all;
% red tiff image
[fileR pathR] = uigetfile('*.tif', 'Choose the Red Single Bead *.tif File',...
                        'Multiselect', 'off');
cd(pathR);disp('Red tiff file : ');disp(pathR);disp(fileR);
% green tiff image
[fileG pathG] = uigetfile('*.tif', 'Choose the Green Single Bead *.tif File',...
                        'Multiselect', 'off');
cd(pathG);disp('Green tiff file : ');disp(pathG);disp(fileG);
% Frame number
frames = input('Enter Total No. of Frames in Green Movies?  ');
% import data
disp('Import Data ...ing... ');
red_images = load_images4([pathR fileR]);
green_images = load_images4([pathG fileG]);
%% display first

green = green_images{1,1};
red = red_images{1,1}; 
scrsz = get(0,'ScreenSize');
figure('OuterPosition',[1 0.4*scrsz(4) 1200 600]);

% colorbar set to enhance contrast
norm_G = double(max(green(:)));
nNormG = double(green)./norm_G;
glevel = graythresh(nNormG); 
cmaxG = glevel*norm_G;
clims1 = [0 norm_G];

colormap gray;
subplot(1,2,1) 
imagesc(green, clims1); %colorbar;

norm_R = double(max(red(:)));
nNormG = double(red)./norm_R;
rlevel = graythresh(nNormG); 
cmaxR = glevel*norm_R;
clims2 = [0 norm_R];
%colormap red;
subplot(1,2,2) 
imagesc(red, clims2); %colorbar;

%% select ROIs with beads
Bd = input('Do you want to save ROIs with beads? press y or n?  ', 's');
Bd1 = Bd;
if Bd == 'y'

    disp('Merged RGB image shows start (R), middle (G) and end (B)');
    a(:,:,1) = red/max(red(:));
    a(:,:,2) = green/max(green(:));
    a(:,:,3) = 0;
    roisel = [];
    
    figure;
    imshow(a);
    title('Select regions to crop and save bead movies.');
    
    while Bd1 == 'y'
          h = imrect;
          position = wait(h);
          pos = getPosition(h);
          roisel = [roisel pos'];
          Bd1 = input('Do you wanna select another bead? press y or n?  ',...
                      's');
    end
    
    [dim noofbeads] = size(roisel);
    roisel = uint16(roisel);
end
%% calculate Xcorr for each ROU selected from averaged file
if exist('Xcorr.tif', 'file')
       delete('Xcorr.tif');
end
fitParams = struct( ...
            'curvatureThreshold',   1000, ...
            'intensityThreshold',   70, ...
            'windowSize',           7, ...
            'minWidth',             0, ...
            'maxWidth',             7, ...
            'angle',                -0 * pi/180 ... 
            );
        
%window = 6;
[M N]= size(red);
XcorrROI = cell(1,noofbeads);
ACROI = cell(1,noofbeads);
Gdata = cell(1,noofbeads);
offset = [];

if Bd == 'y'
   for bno = 1:noofbeads
       y1 = roisel(2,bno);
       y2 = roisel(2,bno)+roisel(4,bno);
       x1 = roisel(1,bno);
       x2 = roisel(1,bno)+roisel(3,bno);
       ACROI{1,bno} = normxcorr2(red(y1:y2,x1:x2), red(y1:y2,x1:x2));
       XcorrROI{1,bno} = normxcorr2(red(y1:y2,x1:x2), green(y1:y2,x1:x2));
      
       imwrite(uint16(XcorrROI{1,bno}*10000), 'Xcorr.tif',...
           'tif', 'Compression', 'none', 'WriteMode', 'append'); 
       
       % Green data Gaussian fit
       [n, c] = max(XcorrROI{1,bno}(:));
       [ym xm] = ind2sub(size(XcorrROI{1,bno}), c);
       Gdata{bno} = [xm, ym];
       Gdata{bno} = fit_spots2S(XcorrROI{1,bno}*100, 'gaussian',...
                                fitParams, Gdata{bno});
                                  
       % Red data Gaussian fit
       [n, c] = max(ACROI{1,bno}(:));
       [ym xm] = ind2sub(size(ACROI{1,bno}), c);
       ACXY = [xm, ym];
       ACXY = fit_spots2S(ACROI{1,bno}*100, 'gaussian',...
                          fitParams, ACXY);       
       % offset
       offset = [offset (Gdata{1,bno}-ACXY)];  
       %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
       %%%%%% Drift vector is from Red bead to Green bead %%%%%%%
       %%% So Drift correction is : Green dots - Drift vector %%%
       %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
   end   
end
   
%% X and Y Pixel Size
if nargin<2
   Xpixel      =input('Please input X pixel size (in nm)   - ');
   Ypixel      =input('Please input Y pixel size (in nm)   - ');
end
%% Offset calculation
offset = mean(reshape(offset, bno, 8),1);
% offset calculation of Z
disp('Use different Z Calibration Value for different color?');
colorjudge=input('Press s for same, d for different, n without Z value','s');
if strcmp(colorjudge,'d')
ZcalfileG = 'F:\Q. Peter Su\0-Matlab codes for STORM\Calibration Values Green.dat'
ZCalibrationG = importdata(ZcalfileG);
pG = polyfit (ZCalibrationG(:,2),ZCalibrationG(:,1),3);
z_offcalG = pG(1)*(Gdata{1,bno}(:,4)).^3+...
            pG(2)*(Gdata{1,bno}(:,4)).^2+...
            pG(3)*(Gdata{1,bno}(:,4)).^1+...
            pG(4);        
ZcalfileR = 'F:\Q. Peter Su\0-Matlab codes for STORM\Calibration Values Red.dat'
ZCalibrationR = importdata(ZcalfileR);
pR = polyfit (ZCalibrationR(:,2),ZCalibrationR(:,1),3);
z_offcalR = pR(1)*(ACXY(:,4)).^3+...
            pR(2)*(ACXY(:,4)).^2+...
            pR(3)*(ACXY(:,4)).^1+...
            pR(4);
end
if strcmp(colorjudge,'s')
ZcalfileG = 'F:\Q. Peter Su\0-Matlab codes for STORM\Calibration Values.dat'
ZCalibrationG = importdata(ZcalfileG);
pG = polyfit (ZCalibrationG(:,2),ZCalibrationG(:,1),3);
z_offcalG = pG(1)*(Gdata{1,bno}(:,4)).^3+...
            pG(2)*(Gdata{1,bno}(:,4)).^2+...
            pG(3)*(Gdata{1,bno}(:,4)).^1+...
            pG(4);        
ZcalfileR = 'F:\Q. Peter Su\0-Matlab codes for STORM\Calibration Values.dat'
ZCalibrationR = importdata(ZcalfileR);
pR = polyfit (ZCalibrationR(:,2),ZCalibrationR(:,1),3);
z_offcalR = pR(1)*(ACXY(:,4)).^3+...
            pR(2)*(ACXY(:,4)).^2+...
            pR(3)*(ACXY(:,4)).^1+...
            pR(4);
end
if strcmp(colorjudge,'n')
    z_offcalG=0*(Gdata{1,bno}(:,4));
    z_offcalR=0*(ACXY(:,4));
end
%% offset vector of X Y and Z   
                                      %%%%%%%%%%%%%%%%%%%
x_off = offset(:,1)*Xpixel/1000;      %%% x_off in um %%%
y_off = offset(:,2)*Ypixel/1000;      %%% y_off in um %%%
z_off = (z_offcalG-...                %%% z_off in um %%%
        z_offcalR)*0.001;             %%%%%%%%%%%%%%%%%%%
    
%% Driftave matrix
ti = [1:frames];
Driftave = zeros(frames, 5);
Driftave(:,1) = ti;
Driftave(:,2) = repmat(x_off, frames, 1);
Driftave(:,3) = repmat(y_off, frames, 1);
Driftave(:,4) = repmat(z_off, frames, 1);

save([fileG(1:end-4) '-Drift2Color.mat'],'Driftave');
