% =========================================================================
% CleanTraces v1.0.0
% Author Janet Y. Sheung
% =========================================================================
% Program to display traces found by TrackCells.m one by one so the user
% can accept or reject each.
%
%

clear all;
path = uigetdir('Select data folder output by TrackCells.m. It will have the same name as the video file.'); % prompt to choose a folder
cd(path);

files = dir('*.png');
lastframe = strrep(files(end-1).name,'.png','');
load('position_estimates.mat'); % load the traces

% loop through columns
for itrace = 1:60 % go through each of first 60 columns
    X = Q_loc_estimateX(:,itrace); % copy X vector of itrace-th trace
    Y = Q_loc_estimateY(:,itrace); % copy Y vector of itrace-th trace
    
    % find where numbers start
    ValidRows = find(~isnan(X));
    FirstRow = min(ValidRows); % will be number or empty
    LastRow = max(ValidRows); % will be number or empty
    
    % diffimage first and last frames of trace
    if ~isempty(FirstRow)
        I1 = imread(files(FirstRow).name); 
        I2 = imread(files(LastRow).name);
        imshowpair(I1,I2); % display first and last frame of trace
        hold  on 
        plot(Q_loc_estimateX(:,itrace),Q_loc_estimateY(:,itrace),'or') % plot first column
        text(Q_loc_estimateX(FirstRow,itrace)+10,Q_loc_estimateY(FirstRow,itrace)+10,num2str(itrace));
        % overlay trace
        hold off
        
        % prompt 0 or 1 
        choice = input('Enter 1 to keep trace, 0 to reject trace.');
       
        if ~choice %reject trace
            Q_loc_estimateX(:,itrace) = NaN;
            Q_loc_estimateY(:,itrace) = NaN;
        end
    end
end

% plot overlay with all traces
c_list = ['r' 'b' 'g' 'c' 'm' 'y']
I = imread(files(end-1).name); % find last frame in your folder 
imshow(I); % display your last frame so you can plot tracks on top
hold  on % plot stuff on top.
% double click position_estimate variable to load Q_loc_estimateX and
% Q_loc_estimateY
for itrace = 1:60 % go through each of first 60 columns
    X = Q_loc_estimateX(:,itrace); % copy X vector of itrace-th trace
    Y = Q_loc_estimateY(:,itrace); % copy Y vector of itrace-th trace
    
    % find where numbers start
    ValidRows = find(~isnan(X));
    FirstRow = min(ValidRows);
  Cz = mod(itrace,6)+1; %pick color
plot(Q_loc_estimateX(:,itrace),Q_loc_estimateY(:,itrace),'.-','markersize',3,'color',c_list(Cz),'linewidth',3) % plot first column
text(Q_loc_estimateX(FirstRow,itrace)+10,Q_loc_estimateY(FirstRow,itrace)+10,num2str(itrace));

end
hold off

mkdir('CLEAN TRACES') % make cleaned folder
save('CLEAN TRACES/clean_traces.mat', 'Q_loc_estimateX', 'Q_loc_estimateY') %save
print('CLEAN TRACES/LabeledTraces','-dpng')

%% replace entire column with NaN then replot
Col = []; % list of columns you want to delete, for example [3 4 7]
c_list = ['r' 'b' 'g' 'c' 'm' 'y']
Q_loc_estimateX(:,Col) = NaN;
Q_loc_estimateY(:,Col) = NaN;
I = imread(files(end-1).name); % find last frame in your folder 
imshow(I); % display your last frame so you can plot tracks on top
hold  on % plot stuff on top.
% double click position_estimate variable to load Q_loc_estimateX and
% Q_loc_estimateY
for itrace = 1:100 % go through each of first 60 columns
    X = Q_loc_estimateX(:,itrace); % copy X vector of itrace-th trace
    Y = Q_loc_estimateY(:,itrace); % copy Y vector of itrace-th trace
    
    % find where numbers start
    ValidRows = find(~isnan(X));
    FirstRow = min(ValidRows);
  Cz = mod(itrace,6)+1; %pick color
plot(Q_loc_estimateX(:,itrace),Q_loc_estimateY(:,itrace),'.-','markersize',3,'color',c_list(Cz),'linewidth',3) % plot first column
text(Q_loc_estimateX(FirstRow,itrace)+10,Q_loc_estimateY(FirstRow,itrace)+10,num2str(itrace));

end
hold off
%% combine traces
Col1 = 1;
Col2 = 15;
ValidRows1 = find(~isnan(Q_loc_estimateX(:,Col1)));
ValidRows2 = find(~isnan(Q_loc_estimateX(:,Col2)));
End1 = max(ValidRows1)
Start2 = min(ValidRows2)
%% copy second trace into first trace
Q_loc_estimateX(ValidRows2,Col1) = Q_loc_estimateX(ValidRows2,Col2);
Q_loc_estimateY(ValidRows2,Col1) = Q_loc_estimateY(ValidRows2,Col2);
% erase second trace
Q_loc_estimateX(:,Col2) = NaN;
Q_loc_estimateY(:,Col2) = NaN;
%%    once you have cleaned X and Y variables, execute these two lines
mkdir('CLEAN TRACES') % make cleaned folder
save('CLEAN TRACES/clean_traces.mat', 'Q_loc_estimateX', 'Q_loc_estimateY') %save
print('CLEAN TRACES/LabeledTraces','-dpng')
%%    once you have dead X and Y, execute these two lines
%mkdir('DEAD') % make cleaned folder
%save('DEAD/dead_traces.mat', 'Q_loc_estimateX', 'Q_loc_estimateY') %save
%% replace zeros with NaNs
for col = 1:size(Q_loc_estimateX,2)
    col
    BadFrames = find(Q_loc_estimateX(:,col) == 0);
    Q_loc_estimateX(BadFrames,col) = NaN;
    Q_loc_estimateY(BadFrames,col) = NaN;
end