% =========================================================================
% TrackCells v1.0.0
% Author Janet Y. Sheung
% =========================================================================
% Program to take as input videos and saves a MATLAB structure file containing
% all automatically tracked traces as x-y coordinates.
%
%
%% Load videos and write as frames
% 1. ask to select movie file, cd to location 

clear all;           
close all;

MinCellArea = 300;          % Minimum realistic stentor area, in pixels
MaxCellArea = 1500;  
Threshold = 0.0; % intensity threshold adjustment

formats = VideoReader.getFileFormats();
filterSpec = getFilterSpec(formats);
[filename,path] = uigetfile(filterSpec, 'Select One or Data Videos to Initiate Tracking','MultiSelect', 'on');
cd(path);

% 2. if assignmentoptimal code not in directory, prompt for location
VidDirContents = dir();
AOcheck = ~cellfun('isempty',regexp({VidDirContents.name},'assignmentoptimal.m','match'));
    if (max(AOcheck) == 0)
        [AOfilename,AOpath] = uigetfile('','assignmentoptimal.m not found, please select assignmentoptimal.m location'); 
        AOloc = strcat(AOpath,AOfilename);
    else
        AOloc = strcat(path,'assignmentoptimal.m');     
    end

% 3. Loop through selected videos
for VidNum = 1:size(filename,2)
    [dummy,name,ext] = fileparts(filename{VidNum});
    mkdir(name);
    cd(name);
    copyfile(AOloc);
    Vidloc = [path filename{VidNum}];
    v = VideoReader(Vidloc);
    i_Im = 1;
    while hasFrame(v)
        Im = readFrame(v);
%% identify cells in each frame
        GreyIm = rgb2gray(Im);
        level = graythresh(GreyIm);
        BW = imbinarize(GreyIm,level+Threshold);
        IM2 = imcomplement(BW);

% Take off the objects that are smaller or bigger than normal stentors
        MaskL = bwlabel(IM2);
        Areas = regionprops(MaskL, 'Area'); 
        Mask = zeros(size(MaskL));
        if length(Areas) > 1
            for i = 1:length(Areas)
                if (Areas(i).Area > MinCellArea) && (Areas(i).Area < MaxCellArea)
                    Mask(MaskL == i) = 1; 
                end
            end
        end
        Mask2 = bwlabel(Mask);
        s = regionprops(Mask2, 'Centroid');
% Cell outline overlayed with the original image
        Outline = bwlabel(double(bwperim(Mask)));
        Outline = imdilate(Outline, strel('disk',2));
        ImgR = Im(:,:,1);
        ImgR(logical(Outline)) = max(ImgR(:));              
        Im(:,:,1) = ImgR;        
        X{i_Im} = [];
        Y{i_Im} = [];
        for iObj = 1:size(s,1)  
            X{i_Im} = vertcat(X{i_Im},s(iObj).Centroid(1));  
            Y{i_Im} = vertcat(Y{i_Im},s(iObj).Centroid(2));
        end
        
        if i_Im <10
            imwrite(GreyIm,['00' num2str(i_Im) '.png'])
        elseif i_Im <100
            imwrite(GreyIm,['0' num2str(i_Im) '.png'])
        else 
            imwrite(GreyIm,[num2str(i_Im) '.png'])
        end
        i_Im = i_Im+1;
    end
    % save XY locations
    save('raw_stentor_detections.mat',  'X','Y')


%% 3 Link cell detections into traces using kalman filter
    load('raw_stentor_detections.mat')

    %get frame list
    f_list =  dir('*png');

    % define main variables for KALMAN FILTER! :P
    dt = 1;  %our sampling rate
    S_frame = 5 % find(cellfun(@length, X)>11,1); %starting frame

    %now, since we have multiple flies, we need a way to deal with a changing
    %number of estimates! this way seems more clear for a tutorial I think, but
    %there is probably a much more efficient way to do it.

    u = 0; % define acceleration magnitude to start
    HexAccel_noise_mag = 1; %process noise: the variability in how fast the Hexbug is speeding up (stdv of acceleration: meters/sec^2)
    tkn_x = 10;  %measurement noise in the horizontal direction (x axis).
    tkn_y = 10;  %measurement noise in the horizontal direction (y axis).
    Ez = [tkn_x 0; 0 tkn_y];
    Ex = [dt^4/4 0 dt^3/2 0; ...
        0 dt^4/4 0 dt^3/2; ...
        dt^3/2 0 dt^2 0; ...
        0 dt^3/2 0 dt^2].*HexAccel_noise_mag^2; % Ex convert the process noise (stdv) into covariance matrix
    P = Ex; % estimate of initial Hexbug position variance (covariance matrix)

    % Define update equations in 2-D! (Coefficent matrices): A physics based model for where we expect the HEXBUG to be [state transition (state + velocity)] + [input control (acceleration)]
    A = [1 0 dt 0; 0 1 0 dt; 0 0 1 0; 0 0 0 1]; %state update matrice
    B = [(dt^2/2); (dt^2/2); dt; dt];
    C = [1 0 0 0; 0 1 0 0];  %this is our measurement function C, that we apply to the state estimate Q to get our expect next/new measurement

    % initize result variables
    Q_loc_meas = []; % the fly detecions  extracted by the detection algo
    % initize estimation variables for two dimensions
    Q= [X{S_frame} Y{S_frame} zeros(length(X{S_frame}),1) zeros(length(X{S_frame}),1)]'
    Q_estimate = nan(4,500);
    Q_estimate(:,1:size(Q,2)) = Q;  %estimate of initial location estimation of where the flies are(what we are updating)
    Q_loc_estimateY = nan(500); %  position estimate
    Q_loc_estimateX= nan(500); %  position estimate
    P_estimate = P;  %covariance estimator
    strk_trks = zeros(1,500);  %counter of how many strikes a track has gotten
    nD = size(X{S_frame},1); %initize number of detections
    nF =  find(isnan(Q_estimate(1,:))==1,1)-1 ; %initize number of track estimates

%for each frame
    for t = S_frame:length(f_list)-1 

        % load the image
        img_tmp = double(imread(f_list(t).name));
        img = img_tmp(:,:,1);
        % make the given detections matrix
        Q_loc_meas = [X{t} Y{t}];

        % do the kalman filter
        % Predict next state of the flies with the last state and predicted motion.
        nD = size(X{t},1); %set new number of detections
        for F = 1:nF
            Q_estimate(:,F) = A * Q_estimate(:,F) + B * u;
        end

        %predict next covariance
        P = A * P* A' + Ex;
        % Kalman Gain
        K = P*C'*inv(C*P*C'+Ez);


        % now we assign the detections to estimated track positions
        %make the distance (cost) matrice between all pairs rows = tracks, coln =
        %detections
        est_dist = pdist([Q_estimate(1:2,1:nF)'; Q_loc_meas]);
        est_dist = squareform(est_dist); %make square
        est_dist = est_dist(1:nF,nF+1:end) ; %limit to just the tracks to detection distances

        [asgn, cost] = assignmentoptimal(est_dist); %do the assignment with hungarian algo
        asgn = asgn';

        % ok, now we check for tough situations and if it's tough, just go with estimate and ignore the data
        %make asgn = 0 for that tracking element

        %check 1: is the detection far from the observation? if so, reject it.
        rej = [];
        for F = 1:nF
            if asgn(F) > 0
                rej(F) =  est_dist(F,asgn(F)) < 50 ;
            else
                rej(F) = 0;
            end
        end
        asgn = asgn.*rej;


        %apply the assingment to the update
        k = 1;
        for F = 1:length(asgn)
            if asgn(F) > 0
                Q_estimate(:,k) = Q_estimate(:,k) + K * (Q_loc_meas(asgn(F),:)' - C * Q_estimate(:,k));
            end
            k = k + 1;
        end

        % update covariance estimation.
        P =  (eye(4)-K*C)*P;

        % Store data
        Q_loc_estimateX(t,1:nF) = Q_estimate(1,1:nF);
        Q_loc_estimateY(t,1:nF) = Q_estimate(2,1:nF); 

        %ok, now that we have our assignments and updates, lets find the new detections and
        %lost trackings

        %find the new detections. basically, anything that doesn't get assigned
        %is a new tracking
        new_trk = [];
        new_trk = Q_loc_meas(~ismember(1:size(Q_loc_meas,1),asgn),:)';
        if ~isempty(new_trk)
            Q_estimate(:,nF+1:nF+size(new_trk,2))=  [new_trk; zeros(2,size(new_trk,2))];
            nF = nF + size(new_trk,2);  % number of track estimates with new ones included
        end


        %give a strike to any tracking that didn't get matched up to a
        %detection
        no_trk_list =  find(asgn==0);
        if ~isempty(no_trk_list)
            strk_trks(no_trk_list) = strk_trks(no_trk_list) + 1;
        end

        %if a track has a strike greater than 6, delete the tracking. i.e.
        %make it nan first vid = 3
        bad_trks = find(strk_trks > 20);
        Q_estimate(:,bad_trks) = NaN;

        %%{
        clf
        img = imread(f_list(t).name);
        imshow(img);
        hold on;
        plot(X{t}(:),Y{t}(:),'or'); % the actual tracking
        T = size(Q_loc_estimateX,2);
        Ms = [3 5]; %marker sizes
        c_list = ['r' 'b' 'g' 'c' 'm' 'y']
        for Dc = 1:nF
            if ~isnan(Q_loc_estimateX(t,Dc))
                Sz = mod(Dc,2)+1; %pick marker size
                Cz = mod(Dc,6)+1; %pick color
                if t < 21
                    st = t-1;
                else
                    st = 19;
                end
               % tmX = Q_loc_estimateX(t-st:t,Dc);
               % tmY = Q_loc_estimateY(t-st:t,Dc);
                tmX = Q_loc_estimateX(1:t,Dc);
                tmY = Q_loc_estimateY(1:t,Dc);
                plot(tmX,tmY,'.-','markersize',Ms(Sz),'color',c_list(Cz),'linewidth',3)
                axis off
            end
        end
       % pause
        %}

        t

    end

%reviewing S_frame
    for t = 1:length(f_list)-1 %S_frame:length(f_list)
        clf;
        img = imread(f_list(t).name);
        imshow(img);
        hold on;
        plot(X{t}(:),Y{t}(:),'or'); % the actual tracking
        T = size(Q_loc_estimateX,2);
        Ms = [3 5]; %marker sizes
        c_list = ['r' 'b' 'g' 'c' 'm' 'y'];
        for Dc = 1:nF
            if ~isnan(Q_loc_estimateX(t,Dc))
                Sz = mod(Dc,2)+1; %pick marker size
                Cz = mod(Dc,6)+1; %pick color
                if t < 21
                    st = t-1;
                else
                    st = 19;
                end
                tmX = Q_loc_estimateX(1:t,Dc);
                tmY = Q_loc_estimateY(1:t,Dc);
                plot(tmX,tmY,'.-','markersize',Ms(Sz),'color',c_list(Cz),'linewidth',3)
                axis off
            end
        end
        t
    %    pause
    end


%%

save('position_estimates.mat', 'Q_loc_estimateX', 'Q_loc_estimateY')

print('TrackingResults','-dpng')

cd(path);
end
