%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
function cluster_correct(img,thresh,cext,cdef)
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%% 
%%%%% Extract cluster information from from stat map
%%%%%
%%%%% Written by: Michael T Rubens
%%%%% Gazzaley Lab, Department of Neurology, UC San Francisco
%%%%% May 30th, 2008
%%%%% ---------------------------------------------------------------------
%%%%% ---------------------------------------------------------------
%%%%% ---------------------------------------------
%%%%% ---------------------
%%%%%
%%%%% This script will take any image and correct based on a cluster extent
%%%%% and voxelwise threshold. The threshold entered will be applied to
%%%%% positive and negative values separately, and separate pos and neg
%%%%% corrected images will be output. This script requres a license for
%%%%% the image processing toolbox.
%%%%%
%%%%%   Functions callled: spm_vol - get header info from image file
%%%%%                      spm_read_vols - get data and coords from ^
%%%%%                      spm_write_vol
%%%%%                      spm_matrix
%%%%%                      bwlabeln (image processing toolbox)
%%%%%
%%%%%    GNU General Public License
%%%%%    ============================================================
%%%%%    This program is free software: you can redistribute it and/or modify
%%%%%    it under the terms of the GNU General Public License as published by
%%%%%    the Free Software Foundation.
%%%%%
%%%%%    This program is distributed in the hope that it will be useful,
%%%%%    but WITHOUT ANY WARRANTY; without even the implied warranty of
%%%%%    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
%%%%%    GNU General Public License for more details.
%%%%%
%%%%%    A full copy of the GNU General Public License can be found here: 
%%%%%    <http://www.gnu.org/licenses/>.
%%%%%    ============================================================
%%%%%
%%%%% ---------------------------------------------------------------------
%%%%% Thanks to: J Zack Chadick, Ted Zanto & Jonathan Kalkstein
%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%% Critical Input Parameters


Cluster_ext = 25; % Cluster extent, only report clusters with at least this # of contiguous voxels

Tval            = 2.1; % Threshold, only look for clusters where all voxels are T >= Tval
                             % as well as T<= -Tval
                             % this does not have to be a T value, can be anything.

Cluster_def = 6; % Look at CONN variable in 'help bwlabeln'
                          %N-voxel connected neigborhood ( 26 | 18 | 6 ) 
                          % 26 is max - voxel connected at face, edge and vertex
                          % 18        - connected at face and edge
                          % 6         - connected at face only (nearest-neighbor)
                    
autopilot       = 0;%1 = autorun, 0 = select which images to run
                           % if on, it will automatically run all images in the
                           % specified directory

%% select images for analysis

%Set default values if not specified
if ~exist('thresh','var')
    thresh = Tval;
end
if ~exist('cext','var')
    cext = Cluster_ext;
end
if ~exist('cdef','var')
    cdef = Cluster_def;
end

if ~nargin
    fprintf('Choose directory containing images to Analyze\n');
    % master_dir = uigetdir('/Volumes/Gazzlab/Data/delay_length/fMRI_DATA/processed/group_stats/RFX_analysis_allsubs','Choose directory containing images to zero');
    master_dir = uigetdir('/Volumes/Gazzlab/Data/AID/AID_fMRI/AID_fMRI_data/group_stats','Choose directory containing images to zero');
    if master_dir == 0
        disp('You need to select at least 1 image!')
        return
    end
    
    files1 = dir(master_dir); files = {}; fnames1 = {};
    for a=1:length(files1)
        if length(files1(a).name)>4 && strcmpi(files1(a).name(end-3:end),'.img') && ~files1(a).isdir && (~strcmpi(files1(a).name,'.') && ~strcmpi(files1(a).name,'..') && ~strcmpi(files1(a).name,'.DS_Store')) % && ~isnan(str2double(scan_info(a).name))
            fnames1{end+1} = files1(a).name;
        end
    end
    fnames1 = sortrows(fnames1');
    %% Select subject files to run, den run
    if autopilot
        fnames = fnames1;
    elseif exist('mask_this_image','var')
        fnames = cellstr(mask_this_image);
    else
        Sel = listdlg('ListString',fnames1,'Name','Select all subject files to process','PromptString','Select all subject files to process');
        fnames = fnames1(Sel);
    end
else
    [tmp1 tmp2] = fileparts(img); cd(tmp1)
    fnames = {img};
    clear tmp1 tmp2 img
end
%% Preceedings
current_dir = pwd;
start_time = datestr(now);tic

warning('off','MATLAB:pack:InvalidInvocationLocation')
%% Begin cluster analysis

fprintf('\n\n\n\n=================================================================================================\n')
fprintf('=================================================================================================\n\n')
fprintf('                             ... Beginning Cluster Analysis...\n\n');
fprintf('=================================================================================================\n')
fprintf('=================================================================================================\n\n')

%% Load mask and apply it to image

for xx = 1:length(fnames)
    fprintf('\n\nAnalyzing Image: %s\n', fnames{xx})
    Shdr = spm_vol(fullfile(master_dir,fnames{xx})); 
    [Simg] = spm_read_vols(Shdr); %Stat image
    %Remove nans
    Simg(find(isnan(Simg))) = 0;
    
    % 2 images, one positive only, one negative only
    fprintf('\nSeparating data into + and - maps...\n');
    Simg_pos = Simg;
    Simg_neg = Simg;
    Simg_pos(find(Simg_pos<0)) = 0;
    Simg_neg(find(Simg_neg>0)) = 0;
    
    fprintf('Thresholding stat image...\n');
    Simg_pos(find(Simg_pos<thresh)) = 0;
    Simg_neg(find(abs(Simg_neg)<thresh)) = 0;
    
    % Make binary
    fprintf('Binarizing maps...\n');
    Simg_pos(find(Simg_pos)) = 1;
    Simg_neg(find(Simg_neg)) = 1;
    
    %get_cluster_matrix
    fprintf('Generating cluster maps from data...\n');
    [pos_clust pclusts] = bwlabeln(Simg_pos,cdef);
    [neg_clust nclusts] = bwlabeln(Simg_neg,cdef);
    
    %Chips method of extracting out only clusters of size cext
    fprintf('Refining cluster maps to specifications..   0%%');
    
    pos_clust = pos_clust(:);
    neg_clust = neg_clust(:);
    
    plist = [0:pclusts];
    nlist = [0:nclusts];
    
    pcnts = arrayfun(@(balls)sum(pos_clust==balls),plist);
    fprintf('\b\b50%%')
    ncnts = arrayfun(@(balls)sum(neg_clust==balls),nlist);
    fprintf('\b\b\b100%%\n')
    
    prem  = unique(bsxfun(@times,plist,(pcnts<cext)));
    nrem  = unique(bsxfun(@times,nlist,(ncnts<cext)));
    
    positive_clusters = bsxfun(@times,pos_clust,~ismember(pos_clust,prem));
    negative_clusters = bsxfun(@times,neg_clust,~ismember(neg_clust,nrem));
    
    %Write new volume, if specified:
    %Create new image, thresholded by a cluster extent calculated at threshold thresh
    nimg = zeros(Shdr.dim);
    indx = find(positive_clusters|negative_clusters);
    nimg(indx) = Simg(indx);

    tstr = num2str(Tval);

    %Write new image
    fprintf('\nWriting Thresholded Image...\t');
    Shdr.fname = fullfile(master_dir,[strrep(fnames{xx}(1:end-4),'__','_')... % double underscore was filenaming error made earlier, corrected here
    [strrep(sprintf('_Cluster-%d_Thresh-%s',cext,tstr),'.','p') '.img']]);
    spm_write_vol(Shdr,nimg);
    fprintf('Complete!\n');
    fprintf('\n\nSuccessfully completed cluster correction of %s @ %s\n\n',fnames{xx},datestr(now));
    
end
fprintf('\n\n\n========================================================================\n')
fprintf('Finished cluster correction on %d images!\nStarted: %s\nEnded: %s\n\n',length(fnames),start_time,datestr(now));toc
disp('========================================================================')
fprintf('========================================================================\n\n\n')


end
