function [ROI_mask, varargout] = ROI_detection(CO, varargin)
% 
% ROI DETECTION Determine the ROI of a complex-valued dataset.
%
%   [ROI_mask, CO_inversion] = ROI_detection(CO) returns a logical matrix 
%   of the ROI.
%
%   INPUTS:  CO:          Calibration Object dataset, normalized [x,y,f,z]      
%
%   ------------------------------------------------------------------
%
%   ROI DETECTION uses the cluster ROI method by default. The user can use 
%   the following optional inputs, written in a 'key', value pairing during
%   the function call:
%
%   OPTIONAL INPUTS:  'maxima_ROI', 1:              Select maxima method to estimate PSF ROI; default is cluster method.
%                     'probe_size', [x,y]:          Size of scattering probe [x,y] in mm. Default is 10x10.
%                     'plot_ROI_summation', 1:      Plot automated ROI against frequency summation.
%
%   ------------------------------------------------------------------
%
%   OUTPUTS:            ROI_mask:           Cell matrix containing logical ROI mask [f,z] 
%                       CO_inversion:       Initial magnitude CO dataset with inversion performed  
%
%   ------------------------------------------------------------------
%
%   Example 1
%   ---------
%   Determine ROI with ROI cluster method
%       [ROI_mask] = ROI_detection(CO);
%
%   Example 2
%   ---------
%   Determine ROI with ROI maxima method
%       [ROI_mask] = ROI_detection(CO, 'maxima_ROI', 1);
%
%   ------------------------------------------------------------------
%
%   Written by Jessica Trac, McMaster University, EMVi Lab
%   Created: Jul 9 2018, for PSF_classification.m validation
%
%   ------------------------------------------------------------------
%

%% (0) SETTING OPTIONAL PARAMETERS 
% ---------------------------------------------------------------------- 
% From user input, where an optional parameter can be set using a 'key',
% value pairing (ie. to evaluate magnitude SNR, pass in 'SNR_mag', x where
% x is the desired user-threshold for magnitude SNR classification).
% ---------------------------------------------------------------------- 
 
fprintf('Detecting ROI... \n');

var_input = struct('maxima_ROI', 0, 'probe_size', [10,10],'plot_ROI_summation', 0); % Default values for each optional parameter
var_input_names = fieldnames(var_input); % Optional parameter names

for pair = reshape(varargin, 2, []) % 1x2 cell of parameter name, value
    var_name = pair{1}; % Parameter name, as specified by user in function call
    var_value = pair{2}; % Parameter value, as specified by user in function call
    if any(strcmp(var_name, var_input_names))
        var_input.(var_name) = var_value; % Override default value
    else
        error('%s is not a recognized parameter name - please re-enter', var_name); 
    end
end

flag = 0; % Flag to indicate any issues with algorithm to user

%% (1) SIZE CONSIDERATIONS
% ---------------------------------------------------------------------- 

[num_x, num_y, num_f, num_z] = size(CO);
center_xy = [round(num_x/2), round(num_y/2)]; % Determine center-point of image in x,y (x = columns, y = rows)

%% (2) INVERSION
% ---------------------------------------------------------------------- 
% Inversion is performed if mean(ROI) < mean(background). Here, mean(ROI) 
% uses the user-defined PSF estimate, derived from the 'padding' input.
% ---------------------------------------------------------------------- 

padding = ceil(num_x*0.05); % Take 5% of image range to set initial PSF estimate location to center 10% of image (ie. center point is [50,50] for 100x100 image - will take [45:55, 45:55] as estimate of PSF location)
assignin('base','padding',padding);

PSF_init_col = (center_xy(1,1) - padding):(center_xy(1,1) + padding); % The 'x-direction' is an index for columns
PSF_init_row = (center_xy(1,2) - padding):(center_xy(1,2) + padding); % The 'y-direction' is an index for rows
PSF_init_location = zeros(num_x, num_y);
PSF_init_location(PSF_init_row,PSF_init_col) = 1; % Initial PSF estimate, represented by 1s (center of image + user-defined padding)

CO_inversion = zeros(num_x,num_y,num_f,num_z); % New dataset with inversion considerations

for f = 1:num_f
    for z = 1:num_z

        temp_CO_mag = abs(CO(:,:,f,z)); % Magnitude of CO at specific frequency and z-plane (temporary use)
        
        temp_CO_mag_ROI = NaN(num_x,num_y); 
        temp_CO_mag_ROI(PSF_init_location == 1) = temp_CO_mag(PSF_init_location == 1); % Populate ROI with signal value, rest of matrix is NaN

        temp_CO_exclusion = temp_CO_mag; 
        temp_CO_exclusion(PSF_init_location == 1) = NaN; % Populate exclusion zone with signal value in non-ROI region, ROI is NaN

        if mean2(temp_CO_mag_ROI(~isnan(temp_CO_mag_ROI))) < mean2(temp_CO_exclusion(~isnan(temp_CO_exclusion))) 
            CO_inversion(:,:,f,z) = (abs(CO(:,:,f,z))-1)*(-1); % Perform inversion on magnitude of CO
        else
            CO_inversion(:,:,f,z) = abs(CO(:,:,f,z)); % Do not perform inversion - taking magnitude of CO
        end
        
    end
end

%% (3A) CLUSTER: ROI 
% ---------------------------------------------------------------------- 
% This is the default ROI detection method. It detects the PSF ROI by
% searching for a center cluster of points with a value higher than the
% baseline 3dB loss from 'maxima'. The center cluster must be within 25% of
% the center of the image (ie. middle 50% of image) - the algorithm will
% iterate in increments of 5%, up to 50% (ie. 100% of image) until a center
% cluster is detected.
% ---------------------------------------------------------------------- 

sum_CO_f = zeros(num_x, num_y, num_z);
center_cluster_region_percent = 0.25:0.05:0.50;
percent_step_size = 0.05; 
circle_size = round(num_x*percent_step_size); 

if ~var_input.maxima_ROI
    for z = 1:num_z

        sum_CO_f(:,:,z) = sum(CO_inversion(:,:,:,z),3); % Summation across all frequencies (taking magnitude of complex dataset) for each z-plane

        baseline_threshold = (max(max(sum_CO_f(:,:,z))) - min(min(sum_CO_f(:,:,z))))*(10^-0.3); % -3dB = 10^-0.3 from maximum
        threshold = max(max(sum_CO_f(:,:,z))) - baseline_threshold;

        ROI_points = threshold < sum_CO_f(:,:,z); % All points above 3dB threshold
        ROI_clusters = bwconncomp(ROI_points); % Identify all clusters above 3dB threshold
        ROI_clusters_centroids = regionprops(ROI_clusters,'Centroid'); % Calculate center-points of each cluster
        ROI_clusters_centroids = cell2mat(struct2cell(ROI_clusters_centroids)');

        clear ROI_clusters_centroids_region % Need to clear for each z-plane in order for conditional below to be valid at the start
        m = 0; 
        
        while ~exist('ROI_clusters_centroids_region','var') % Iterate until a cluster is found within the region 

            m = m + 1;
            x_lower = center_xy(1,1) - floor(num_x*center_cluster_region_percent(m)); % Updating search box for center cluster region
            y_lower = center_xy(1,2) - floor(num_y*center_cluster_region_percent(m));

            x_upper = center_xy(1,1) + floor(num_x*center_cluster_region_percent(m));
            y_upper = center_xy(1,2) + floor(num_y*center_cluster_region_percent(m));

            c = 1;
            for k = 1:length(ROI_clusters_centroids(:,1))
                if (ROI_clusters_centroids(k,1) > x_lower && ROI_clusters_centroids(k,1) < x_upper) && (ROI_clusters_centroids(k,2) > y_lower && ROI_clusters_centroids(k,2) < y_upper)
                    ROI_clusters_centroids_region(c,:) = [ROI_clusters_centroids(k,1),ROI_clusters_centroids(k,2)]; % Holds all centroids that are within center cluster region search box
                    c = c + 1;
                end
            end
        end
        
        if m > 1 % Notifies user of successful search region (ie. if 50%, indicates that ROI is very off-centered)
            fprintf('Important: Center cluster for PSF was found within %d%% of center point of image! \n', center_cluster_region_percent(m)*100); 
        end
            
        center_cluster_i = knnsearch(ROI_clusters_centroids_region,center_xy); 
        center_cluster_i = find(ROI_clusters_centroids(:,1) == ROI_clusters_centroids_region(center_cluster_i,1)); % Finding index from original set of coordinates (looking at column 1 only)

        ROI_cluster_linear_index = ROI_clusters.PixelIdxList{center_cluster_i}; % Extract linear indices of center cluster (now ROI) from ROI_clusters struct
        [row_index_cluster, col_index_cluster] = ind2sub([num_x, num_y], ROI_cluster_linear_index); % Convert linear indices to their indices for row and column
        ROI_cluster = [row_index_cluster,col_index_cluster]; 
        
        determined_ROI = zeros(num_x,num_y); 
        num_pts_cluster = length(ROI_cluster(:,1));
        
        % Increasing a 1-pt ROI to a 5-pt 
        if num_pts_cluster == 1 
           point_x = ROI_cluster(1,1);
           point_y = ROI_cluster(1,2);
           to_add = [point_x + 1, point_y; point_x - 1, point_y; point_x, point_y + 1; point_x, point_y - 1];
           ROI_cluster = [ROI_cluster;to_add];
           num_pts_cluster = length(ROI_cluster(:,1));
        end
        
        for k = 1:num_pts_cluster % # of rows
            determined_ROI(ROI_cluster(k,1),ROI_cluster(k,2)) = 1; % Creating a logical matrix with ROI = 1, non-ROI = 0 [x,y]
        end
        
        % Determining a bounding circle to define a circular exclusion zone
        center_coords(z,:) = ROI_clusters_centroids(center_cluster_i,:); % Holds centroid coordinates (each z-plane)
        x_center = center_coords(z,1);
        y_center = center_coords(z,2);

        outer_boundary_cell = bwboundaries(determined_ROI);
        outer_boundary = outer_boundary_cell{1};
        boundary_distances = sqrt((outer_boundary(:,2) - x_center).^2 + (outer_boundary(:,1) - y_center).^2);

        [~, max_radius_i] = max(boundary_distances); % Determines maximum radius out of all boundary coordinates, and finds index of coordinates
        max_point_coord = outer_boundary(max_radius_i,:); % Boundaries for max radius
        
        radius_init(z,1) = ceil(sqrt((max_point_coord(1,2) - x_center)^2 + (max_point_coord(1,1) - y_center)^2)); % Use centroid and max_boundary coords to determine radius of circular mask
        
        % Check size of radius (to avoid using a very abnormal cluster)
        iteration_limit_init = num_x - x_center; % left
        if x_center < iteration_limit_init 
            iteration_limit_init = x_center; % right
        end
        
        iteration_limit_init_y = num_y - y_center; % top
        if y_center < iteration_limit_init_y 
            iteration_limit_init_y = y_center; % bottom
        end
        
        if iteration_limit_init_y < iteration_limit_init
            iteration_limit_init = iteration_limit_init_y;
        end

        if radius_init(z,1) > iteration_limit_init % If cluster is an abnormal shape, could lead to radius that goes over image dimensions - a flag to switch to more fullproof maxima method
            flag = 1;
            fprintf('Warning: could not find suitable ROI cluster and defaulted to maxima ROI method. Dataset may be very noisy.\n\n');
            break % Switching to maxima method
        end
        
        [xx,yy] = ndgrid((1:num_x) - y_center,(1:num_y) - x_center);
        circle_ez_logical = (xx.^2 + yy.^2) < radius_init(z,1)^2;

        for f = 1:num_f % Mapping ROI and Exclusion Zone to each frequency
            ROI_mask{f,z} = determined_ROI; % Complete matrix of logical mask for ROI (does not vary w frequency, just recopied)
            Ez_circle_initial{f,z} = circle_ez_logical;
        end
    end

end

%% (3B) MAXIMA: ROI
% ---------------------------------------------------------------------- 
% This is the alternative method of ROI detection, which can be selected by
% the user for use. It detects the maxima of the frequency-summed PSF and
% uses that as a center-point for a box, whose size is that of the
% scattering probe. 
% ---------------------------------------------------------------------- 

if var_input.maxima_ROI || flag 
    
    if flag % % Executing maxima method because cluster method had flag raised
        clearvars determined_ROI center_coords radius_init ROI_mask Ez_circle_initial % Clear variables that may have been used by cluster method
    end
    
    probe_size = var_input.probe_size; % as specified by user, otherwise default 10x10 mm [row,col > Y, X]
    probe_size = floor(probe_size/2); % 1px = 2mm
    probe_step_size = floor(probe_size/2); 
    
    for z = 1:num_z
        sum_CO_f(:,:,z) = sum(CO_inversion(:,:,:,z),3); % Summation across all frequencies (taking magnitude of complex dataset) for each z-plane
        CO_initial_estimate(:,:,z) = NaN([num_x num_y]); % Initialize CO-sized matrix with NaN values
        CO_initial_estimate(PSF_init_row,PSF_init_col,z) = sum_CO_f(PSF_init_row,PSF_init_col,z); % Holds values within initial PSF estimate
        sum_CO_f_max = find(CO_initial_estimate(:,:,z) == max(max(CO_initial_estimate(:,:,z))));
        [X_index_maxima, Y_index_maxima] = ind2sub([num_x, num_y], sum_CO_f_max); % X and Y coordinates of summation maxima
        
        ROI_X_maxima = (X_index_maxima - probe_step_size(1,2)):(X_index_maxima + probe_step_size(1,2)); % X values
        ROI_Y_maxima = (Y_index_maxima - probe_step_size(1,1)):(Y_index_maxima + probe_step_size(1,1)); % Y values
        determined_ROI = zeros(num_x,num_y); % Creating a logical matrix with ROI = 1, non-ROI = 0
        determined_ROI(ROI_X_maxima, ROI_Y_maxima) = 1;
        
        % Determining a bounding circle to define a circular exclusion zone
        center_coords(z,:) = [Y_index_maxima, X_index_maxima]; % Holds centroid coordinates (each z-plane)
        x_center = center_coords(z,1);
        y_center = center_coords(z,2);

        outer_boundary_cell = bwboundaries(determined_ROI);
        outer_boundary = outer_boundary_cell{1};
        boundary_distances = sqrt((outer_boundary(:,2) - x_center).^2 + (outer_boundary(:,1) - y_center).^2);

        [~, max_radius_i] = max(boundary_distances); % Determines maximum radius out of all boundary coordinates, and finds index of coordinates
        max_point_coord = outer_boundary(max_radius_i,:); % Boundaries for max radius

        radius_init(z,1) = ceil(sqrt((max_point_coord(1,2) - x_center)^2 + (max_point_coord(1,1) - y_center)^2)); % use centroid and max_boundary coords to determine radius of circular mask
        
        [xx,yy] = ndgrid((1:num_x) - y_center,(1:num_y) - x_center);
        circle_ez_logical = (xx.^2 + yy.^2) < radius_init(z,1)^2;
        
        for f = 1:num_f
            ROI_mask{f,z} = determined_ROI; % Complete matrix of logical mask for ROI (does not vary w frequency, just recopied)
            Ez_circle_initial{f,z} = circle_ez_logical; 
        end
    end

end


%% (4) PLOT ROI AGAINST FREQUENCY SUMMATION (OPTIONAL)
% ---------------------------------------------------------------------- 

if var_input.plot_ROI_summation % Plot frequency summation of PSF with automated ROI 
    figure; 
    for z = 1:num_z
        ROI_xy_z = cell2mat(ROI_mask(1,z));
        
        if num_z > 1 % Plot all z-planes in one figure
            subplot(round(num_z/2),2,z);
        end

        imagesc(sum_CO_f(:,:,z)); colorbar; 
        hold on;
        mask_ROI = cat(3, zeros(num_x, num_y), zeros(num_x, num_y), zeros(num_x, num_y));
        h = imagesc(mask_ROI); 
        viscircles(center_coords(z,:),radius_init(z,1),'Color','w', 'LineWidth', 0.1);% hold off; 
        hold off;
        set(h, 'AlphaData', ~ROI_xy_z*0.35); % Automated ROI will appear as bright set of pixels
        title(['Frequency Summation; ROI, z-plane: ' num2str(z)]);
        set(gca,'YDir','normal');
    end
end

varargout{1,:} = CO_inversion; % Initial magnitude CO dataset with inversion performed