function join_axon()
%JOIN_AXON opens a GUI that takes data saved by suite2p and clusters the
%ROIs based on the correlation coefficient. The user can then save the data
%in a similar format to the original produced by suite2p.
%
% Requires the "distinguishable_colors" function, by Tim Holy.
%
%   Bill Connelly, 2017.

% We perform filtering prior to correlation. This Produces significant
% imporovement  in correlation between ROIs belonging to the same cell.
% Window Size is arbitrary, and sets the filter properties: larger = more
% aggressive filtering. We found running average filter behaved better than
% a typical low pass.

% Once the clusters are set they are saved in a cell array called
% dat.clusters, where each cell is a cluster of ROIs belonging (hopefully)
% to the same cell, and each element in the cell is the index of the ROI.
% This index can be used to access data.stat and dat.Fcell.

% SYNTAX:
%   > join_axon(); %function does not return anything

% USAGE:
%   Click on the help button for help.

% You can write your own script to iterate through ROIs belonging to the
% same cluster. For instance, to produce a weighted average for the F
% signal (weighted by number of pixels in each ROI) one could write:
%
% n_datapoints = dat.ops.Nframes(1);
% raw_f = zeros(length(dat.clusters), n_datapoints); %preallocate
% 
% for c = 1:length(dat.clusters)                 %iterate over all clusters
%     total_weight = 0;
%     sum_f = zeros(1, n_datapoints);            %preallocate
%     for r = 1:length(dat.clusters{c})              %iterate over each ROI in cluster
%        total_weight = total_weight + dat.stat(dat.clusters{c}(r)).npix;
%        sum_f = sum_f + dat.stat(dat.clusters{c}(r)).npix * dat.Fcell{1}(dat.clusters{c}(r),:);
%     end
%     raw_f(c,:) = sum_f./total_weight;          % Weighted average F value for this cluster.
% end

% Programmers confession:
% This function is written with a large number of variables which are
% global between this function and the nested subfunctions. This, in my
% opinion, makes for much more managable callback function management.
% However, if this offends you, I understand.

%% FUNCTION BEGINS

%Default variables. Change these if you're constantly using other values.
%These are just ones that worked well for my data.

thresh = 0.8;    %Threshold correlation coefficient. All ROIs with a
                 %correlation coefficient higher than this are considered to
                 %belong to the same cluster.
                
windowSize = 20; %The window for the running average filter, in frames.

%Declare some variables to make the global within the scope of this
%function

filename = [];
path = [];
stats = [];
n_rois = [];
fcell = [];
corr_mat = [];
fcell_filt = [];
data = [];
clusters = {};


%% GUI CREATION
% Figure is created with three axis. One is the image of the sampled plane
% with ROIs on it, ROIs that have been clustered at the same color.
% Top right figure shows the flourescent signal for all clustered ROIs
% Bottom right shows the flourescent signal for the ROI that you have clicked
% on.

% Finally there is a slider that allows you to set the correlation
% coefficient threshold. 

corr_fig = figure('Position', [50 150 1800 600]);

corr_axes = axes('Position', [0.01 0.28 0.25 0.7], 'XTickMode','manual',...
    'YTickMode', 'manual');

cellvector = 1;   %Array to hold to current displayed ROIs
target_clust = 1; %Scalar for the current cluster

df_axes = axes('Position', [0.3, 0.6, 0.68, 0.35]);
title('F signal for all ROI(s) in cluster');
ylabel('Z score');

single_trace_axes = axes('Position', [0.3 0.18 0.68 0.35]);
title('F signal for selected ROI(s)');
ylabel('Filtered F Signal');

slider1_handle=uicontrol(corr_fig,'Style','slider','Max',1,'Min',0,...
    'Value',thresh,'SliderStep',[0.01 0.05], 'Units','normalized',...
    'Position',[.4 .1 .2 .025], 'Callback', @slider_callback);

r_string = uicontrol(corr_fig,'Style','text','Units','normalized',...
    'Position',[.42 .05 .14 .04], 'String', ['Set R threshold. ', num2str(thresh)]);

roi_list = uicontrol(corr_fig, 'Style', 'listbox',...
    'Position',  [10, 50, 100, 60], 'Callback', @show_single_clusters);

roi_txt = uicontrol(corr_fig, 'Style', 'text', 'String', 'View Cluster #',...
    'Position', [10 120 100 20]);

filt_value = uicontrol(corr_fig, 'Style', 'edit', 'String', windowSize,...
    'Position', [130 90 100 20], 'Callback', @filter_change_callback);

filt_txt = uicontrol(corr_fig, 'Style', 'text', 'String', 'Set Filter Width in Frames',...
    'Position', [130 120 100 30]);

open_button = uicontrol(corr_fig, 'Style', 'pushbutton', 'String', 'Open data',...
    'Position', [130 50 80 30], 'Callback', @open_callback);

help_button = uicontrol(corr_fig, 'Style', 'pushbutton', 'String', 'Help!',...
    'Position', [215 50 80 30], 'Callback', @help_callback);

save_button = uicontrol(corr_fig, 'Style', 'pushbutton', 'string', 'Save',...
    'Position', [10 10 80 30], 'Callback', @save_callback);

%% CALLBACK AND HELPER FUNCTIONS BEOW HERE


function open_callback(~,~)
    % Loads the *_proc.mat file
    % Updates several global variables, most importantly fcell, data and stats.
    
    [filename, path] = uigetfile('*.mat','Select the results file', 'G:\temp\results');

    if ~filename    %If the use clicks cancel
        return;
    end
    
    data = load([path filename], 'dat');
    n_rois = sum([data.dat.stat.iscell]);

    % Create only cell matrix, that is, select cells that are dat.stat.iscell
    % Also collapse across different experiments
    n_exp = size(data.dat.Fcell, 2);

    fcell = [];
    for e = 1:n_exp
        fcell = [fcell data.dat.Fcell{1,e}(logical([data.dat.stat.iscell]),:)];
    end

    %Create stats of just cells
    stats = data.dat.stat(1,logical([data.dat.stat.iscell]));        
    
    %Filter the data and performing clustering
    filt_data();
    
    %Show the clusters
    draw_clust();
    
end

function slider_callback(src, ~)
    %Updates the GUI when the R threshold slider is changed
    %Updates the global variable thresh.
    
    thresh = get(src, 'Value');
    clust();
    draw_clust();
    set(r_string, 'String', ['Set R threshold. R = ', num2str(thresh)])
end    


function clust()
    %Clusters ROIs based on their correlation coefficient.
    %It updates the global variable "clusters".
       
    in_cluster = []; %vector to store cells already in cluster
    clusters = {}; %cell array. Each cell is a cluster, and each element in the cell is a ROI.
    n_clusters = 0;

    for c = 1:n_rois
        if ~ismember(c, in_cluster)
            n_clusters = n_clusters + 1;
            clusters{n_clusters} = [];
            for c2 = c:n_rois
                if ~ismember(c2, in_cluster)
                    if corr_mat(c,c2) > thresh
                        clusters{n_clusters} = horzcat(clusters{n_clusters}, c2);
                        in_cluster = [in_cluster, c2];

                    end
                end        
            end
       end    
    end
    
    %Repopulate View Cluster listbox
    set(roi_list,'Value',1); 
    set(roi_list, 'String', { num2str([1:length(clusters)]'), 'All'});
end

function draw_clust()
    % Show Clusters.
    % Does not update any global variables.

    colors = distinguishable_colors(length(clusters), [0 0 0; 1 0 0]);

    im = zeros(data.dat.ylim(2), data.dat.xlim(2), 3);
    im(:,:,1) = imscale(data.dat.mimg(:,:,2));
    for c = 1:length(clusters)
        for r = 1:length(clusters{c})
            for p = 1:length(stats(clusters{c}(r)).ypix)
                im(stats(clusters{c}(r)).ypix(p), stats(clusters{c}(r)).xpix(p), :) = colors(c, :);
            end
        end
    end
    axes(corr_axes)
    imhandle = imshow(im);
    set(imhandle, 'ButtonDownFcn', @imClickCallback);
    
end   

function imClickCallback(src , ~)
    %Show F trace for clicked on ROI(s)
    %It needs to set the global variable cellvector inorder for the
    %currently selected ROI to continue to be displayed even after the GUI
    %updates.
    
    if nargin > 0 % If truely called as a callback.
        axesHandle  = get(src,'Parent');
        coordinates = get(axesHandle,'CurrentPoint');
        coord = coordinates(1,1:2);
        coord = round(coord);   


        xmatch = zeros(length(stats), 1);
        ymatch = zeros(length(stats), 1);
        for r = 1:length(stats)
           xmatch(r) = any(coord(1)==stats(r).xpix);
           ymatch(r) = any(coord(2)==stats(r).ypix);
        end

        cellvector = xmatch & ymatch;
    end
   
    axes(single_trace_axes);
    plot(fcell_filt(windowSize:end, cellvector));
    title('F signal for selected ROI(s)');
    ylabel('Filtered Raw F Signal');

end


function save_callback(~, ~)
    %Save a new *.mat file with similar filename as before. Chops out all
    %non cell data (i.e. stat.icell is false).
    
    dat = data.dat;
    dat.stat =  stats; %replace with just iscell stats
    
    for exp = 1:size(data.dat.Fcell, 2) %replace Fcell and FcellNeu with just iscell data
       dat.Fcell{1,exp} = dat.Fcell{1,exp}(logical([data.dat.stat.iscell]),:);
       dat.FcellNeu{1,exp} = dat.FcellNeu{1,exp}(logical([data.dat.stat.iscell]),:);
    end
    
    %I don't know what they are, but replacing them with just iscell data
    dat.cl.rands_orig = data.dat.cl.rands_orig(1, logical([data.dat.stat.iscell]));
    dat.cl.rands = data.dat.cl.rands(1, logical([data.dat.stat.iscell]));
    
    dat.clusters = clusters;
    
    save([path filename(1:end-4) '_clust.mat'], 'dat');
end


function show_single_clusters(src, ~)
    %Show the F trace (as Z score) for all the ROIs in a cluster
    %Needs to set the global variable target_clust in order to display the
    %same cluster even after GUI update.
    
    if nargin > 0 %If truely called as a callback function
        im = zeros(data.dat.ylim(2), data.dat.xlim(2), 3);
        im(:,:,1) = imscale(data.dat.mimg(:,:,2));
        if src.Value < length(src.String) %If you're not on the last element 
            target_clust = src.Value;

            axes(df_axes)

            for r = 1:length(clusters{target_clust})
                plot( zscore( fcell_filt(windowSize:end,clusters{target_clust}(r)) ) );
                hold on
                for p = 1:length(stats(clusters{target_clust}(r)).ypix)
                    im(stats(clusters{target_clust}(r)).ypix(p), stats(clusters{target_clust}(r)).xpix(p), 2:3) = 0.5;
                    im(stats(clusters{target_clust}(r)).ypix(p), stats(clusters{target_clust}(r)).xpix(p), 1) = im(stats(clusters{target_clust}(r)).ypix(p), stats(clusters{target_clust}(r)).xpix(p), 1)/2;
                end
            end
            hold off
            title('F signal for all ROI(s) in cluster');
            ylabel('Z score');

            axes(corr_axes);
            imhandle = imshow(im);
            set(imhandle, 'ButtonDownFcn', @imClickCallback);
        else
            draw_clust();
            axes(df_axes);       
            cla;
        end
    else
        axes(df_axes)
        for r = 1:length(clusters{target_clust})
            plot( zscore( fcell_filt(windowSize:end,clusters{target_clust}(r)) ) );
            hold on
        end
        hold off
    end
end

function filter_change_callback(src, ~)
    %Parse the filter window input, and then refilter and redraw.
    %Updates the global variable windowSize
    
    filt_length = str2double(src.String);
    if length(filt_length) > 1 %someone trys to enter an array
        errordlg('Please enter a single value')
    elseif isempty(filt_length) %someone enters a string
        errordlg('Please enter a number')
    elseif mod(filt_length,1)~=0 %someone enters a floating point
        errordlg('Please enter an integer')
    elseif filt_length < 1 %someone enters a value less than one
        errordlg('Please enter a value greater than zero')
    elseif isnan(filt_length) %random catch all
        errorlg('Please enter a positive integer')
    else
        windowSize = filt_length;
        filt_data();
        draw_clust();
        imClickCallback();
        show_single_clusters();
    end
end


function filt_data()
    %Filter the data using a running average filter
    %Updates the global variables fcell_filt and corr_mat
    
    b = (1/windowSize)*ones(1,windowSize);
    fcell_filt = filter(b,1,fcell');
    corr_mat = corrcoef(fcell_filt);
    clust();
end

function help_callback(~, ~)
    msg{1} = 'Press the "open" button to open an already processed data file';
    msg{2} = 'i.e a file that ends in ...plane1_proc.mat';
    
    msg{4} = 'The figure on the left will show you how how the different ROIs';
    msg{5} = 'have been clustered. With each cluster being its own color.';
    msg{6} = 'Investigate whether each cluster is appropriate by clicking on a';
    msg{7} = 'cluster in the "View Cluster #" listbox. The ROIs that make up that';
    msg{8} = 'cluster will be draw blue on the image and the waveforms for all ROIs';
    msg{9} = 'in the culster will be shown in the "F signal for all ROI(s) in cluster"';
    msg{10} = 'graph. This graph displays the waveform normalized by as a z score,';
    msg{11} = 'which keeps the waveforms at approximately the same size.';
    
    msg{13} = 'If ROIs that you believe belong to the same dendrite/axon/cell are put';
    msg{14} = 'in different clusters, try reducing the R threshold slider, or';
    msg{15} = 'increasing the filter width (see below). If ROIs that belong to';
    msg{16} = 'different dendrites/axons/cells are being placed in the same cluster';
    msg{17} = 'try increase the R threshold slider, or decreasing the filter width';
  
    msg{19} = 'To go back to viewing all ROIs, click "All" in the "View Cluster #"';
    msg{20} = 'listbox.';
    
    msg{22} = 'To view the waveform for an individual ROI, click on it in the image';
    msg{23} = 'and it will appear in the "F signal for selected ROI(s)".';
    
    msg{25} = 'Better performance will be acheived if the filter is set appropriately.';
    msg{26} = 'Simply enter an integer into the "Set Filter Width in Frames" box, and';
    msg{27} = 'press enter. The filter should be set such that fast noise is minimized,';
    msg{28} = 'but calcium transients are largely unaltered';

    msgbox(msg, 'Help');        
        
end


end