%% Identifies unamplified centrosomes in cells, tracks pairs of centrioles, and calculates the intracentrosomal distance  
% by Robert Hauschild 
% requires >=R2019a
% License: BSD-3-Clause
% uses Jean-Yves Tinevez's "simple tracker" https://www.mathworks.com/matlabcentral/fileexchange/34040-simpletracker
% centrioles need to be identified employing the "pixel classification" workflow of Ilastik https://www.ilastik.org/documentation/pixelclassification/pixelclassification.html 
% export pixel probabilities as "...._Probabilities.h5" 

clearvars; close all;
[currentPath,~,~]=fileparts(mfilename('fullpath'));
addpath(currentPath);

workingDir='D:\TempData\Mirka\'; % folder that contains the Ilastik segmentation files. Output is also saved to this folder. 
physPixelsize=0.1817350; % mu/pixel. All parameters are in [pixel]. Only the final result is converted to physical units [mu] 
useSmoothing=true; % tracks are smoothed using weighted linear least squares and a 1st degree polynomial model over 5 time frames 

% centriole detection parameters:
minSize=2;  % minimum size in pixel for a valid detection of a centriole
maxSize=25; % maximum size in pixel for a valid detection of a centriole
segThreshold=0.5; % Segmentation threshold applied to the pixel probability output of Ilastik

% Track filtering parameters:
minTrackDuration=10; % minimum number of frames a centriole is present
minTrackLength=100; % minimum distance 
MaxLinkingDistance=50; % link consecutive detection of a centriole if they are < MaxLinkingDistance apart
MaxGapClosing=2; % tracks are connected even if the detection is missing for MaxGapClosing-1 frames 
maxspeed=inf; % tracks are discarded if the speed at one time point exceeds maxspeed
frameinterval=1.0; % only relevant if maxspeed is set<inf
pixelsize=1.0; 

%Parameters for the identification of unamplified centrosomes
minAlignment=0.5; % 0.5 = pairing if both centrioles move in parallel 50% of the time
maxAverageDistance=40; % maximum average distance of centrioles in a centrosome [pixel]
distToThirdCentriol=10;  % it is a normal cell if a possible third centriole is at least which much away during the entire track duration


cd(workingDir);
allFiles=dir('*_Probabilities.h5');
cd(currentPath);
AllDistances=[];

for ff=1:length(allFiles)
    fname=fullfile(workingDir,allFiles(ff).name);
    pairedTracks=[];
    maxAlignment=[];
    pairing=[];
    meanDistance=[];
    distTo3=[];
    
    % get tracks
    [Tracks,duration] = goTracking(fname,minTrackDuration,minTrackLength,pixelsize,useSmoothing,frameinterval,maxspeed,segThreshold,minSize,maxSize,MaxLinkingDistance,MaxGapClosing);
    
    % pair tracks by checking for parallel movement
    for i=1:length(Tracks)
        u0=[cosd(Tracks{i}(:,5)), sind(Tracks{i}(:,5)),zeros(size(Tracks{i},1),1)];
        X10=[Tracks{i}(:,1), Tracks{i}(:,2)];
        maxAlignment(i)=-inf;
        pairing(i)=nan;
        meanDistance(i)=nan;
        framesU=Tracks{i}(:,3);
        for j=i+1:length(Tracks)
            v=[cosd(Tracks{j}(:,5)), sind(Tracks{j}(:,5)),zeros(size(Tracks{j},1),1)];
            X2=[Tracks{j}(:,1), Tracks{j}(:,2)];
            framesV=Tracks{j}(:,3);
            [~,idxU] = intersect(framesU,framesV,'stable');
            [~,idxV] = intersect(framesV,framesU,'stable');
            % keep only frames of coexistence
            u=u0(idxU,:);
            X1=X10(idxU,:);
            v=v(idxV,:);
            X2=X2(idxV,:);
            coexistFrames=framesV(idxV);

            alignment=nan(size(u,1),1);
            dx=alignment;
            for s=1:size(u,1)
                alignment(s)=dot(u(s,:),v(s,:),2)/(norm(u(s,:))*norm(v(s,:))); % cos(angle) of the displacement vectors of a pair of centrioles 
                dx(s)=sqrt((X1(s,1)-X2(s,1)).^2+(X1(s,2)-X2(s,2)).^2); % distance of the centrioles
            end
            totAlignment=nansum(alignment);
            if totAlignment>maxAlignment(i)
                maxAlignment(i)=totAlignment;
                pairing(i)=j;
                meanDistance(i)=nanmean(dx);
                pairedTracks{i}=[X1,X2,dx,coexistFrames];  % only the coexistent part;
            end
        end
        % check how far a possible third centriol is from the pair
        % = check if the cell is normal or amplified
        if i<=length(pairedTracks)
            distTo3(i)=inf;
            coexistFrames=pairedTracks{i}(:,6);
            X1= pairedTracks{i}(:,1:2);
            X2= pairedTracks{i}(:,3:4);
            for k=1:length(Tracks)
                if k~=i && k~=pairing(i)
                    [~,idx3]=intersect(Tracks{k}(:,3),coexistFrames);
                    X3=[Tracks{k}(idx3,1), Tracks{k}(idx3,2)];
                    for ii=1:size(X3,1)
                        Xmean=0.5*(X1(ii,:)+X2(ii,:));
                        d23=sqrt((Xmean(1)-X3(ii,1))^2+(Xmean(2)-X3(ii,2))^2);
                        if d23<distTo3(i)
                            distTo3(i)=d23;
                        end
                    end
                end
            end
        else
            distTo3(i)=0;
        end
    end
    idx=find(maxAlignment>minAlignment*(size(X1,1)-1) & meanDistance<maxAverageDistance & distTo3>distToThirdCentriol);
    goodPairing=[idx', pairing(idx)'];
    
    % Plot paired tracks for visual inspection
    n=0;
    for i=1:length(pairedTracks)
        if ismember(i,goodPairing(:,1))
            n=n+1;
            
            % comment this part out if only the numerical results are
            % needed
            figure;
            hold('on');
            for iii=1:length(Tracks)
                plot(Tracks{iii}(:,1),Tracks{iii}(:,2));
            end
            plot(pairedTracks{i}(:,1),pairedTracks{i}(:,2),'linewidth',2);
            plot(pairedTracks{i}(:,3),pairedTracks{i}(:,4),'linewidth',2);
            hold('off');
            set(gca,'Ydir','reverse');
            title(['track ' num2str(n)]);
            % 
            
            dx=pairedTracks{i}(:,5);
            AllDistances=vertcat(AllDistances,dx);
            writematrix(dx.*physPixelsize,fullfile(workingDir,[allFiles(ff).name(1:end-3) '.xls']),'Sheet',['Cell' num2str(n)],'Range','A1');
        end
    end
end
figure
histogram(AllDistances.*physPixelsize,20);
xlabel('distance (mu)');
title('Intracentrosomal Distance Normal Cells' );
writematrix(AllDistances.*physPixelsize,fullfile(workingDir,'AllIntracentrosomalDistances.xls'));
