function roiDataUpdated = RemoveModelOutliers(roiData, varargin)
%REMOVEMODELOUTLIERS Removes outlier ROIs based on z-score thresholding.
%
%   roiDataUpdated = RemoveModelOutliers(roiData) roiDataUpdated =
%   RemoveModelOutliers(roiData, 'threshold', value, 'variablesToConsider',
%   vars)
%
%   This function removes entire ROI entries from the input struct
%   `roiData` if any of the specified features exceed a z-score threshold.
%   The z-score is computed per ROIName group (i.e., within-group
%   normalization).
%
%   Inputs:
%       roiData - Struct array of ROI data with fields corresponding to
%                 geometric features (e.g., Volume, xcm, ycm, etc.)
%
%   Optional Name-Value Pairs:
%       'threshold' - Z-score threshold for identifying outliers (default =
%       3.28) 
%       'variablesToConsider' - String array of field names to check
%       for outliers. Defaults to ["Volume", "xcm", "ycm", "zcm", ...
%                "lengthPCA1", "lengthPCA2", "lengthPCA3", "Area", ...
%                "Solidity", "Compactness", "CMToSurface", "MaxLength"];
%       'additionalVariables' - String array of additional field names
%       beyond the default ones. The main difference between
%       variablesToConsider and additionalVariables is that
%       additionalVariables adds additional variables whereas
%       variablesToConsider replaces the default
%
%   Output:
%       roiDataUpdated - Struct array with outlier ROIs removed
%

%% Argparser
% Default values
defaultThreshold = 3.28;
defaultVariables = ["Volume", "xcm", "ycm", "zcm", ...
    "lengthPCA1", "lengthPCA2", "lengthPCA3", "Area", ...
    "Solidity", "Compactness", "CMToSurface", "MaxLength"];
defaultAdditions = [];

% Parse function inputs
p = inputParser;
addRequired(p, 'roiData', @isstruct);
addParameter(p, 'threshold', defaultThreshold, @isnumeric);
addParameter(p, 'variablesToConsider', defaultVariables, @(x) isstring(x) || iscellstr(x));
addParameter(p, 'additionalVariables', defaultAdditions, @(x) isstring(x) || iscellstr(x));
parse(p, roiData, varargin{:});

threshold = p.Results.threshold;
variablesToConsider = p.Results.variablesToConsider;
additionalVariables = p.Results.additionalVariables;

variablesToConsider = [variablesToConsider, additionalVariables];

%% Group by ROIName
roiDataTable = struct2table(roiData);

if any(contains(variablesToConsider, 'relPos'))
    roiDataTable = FlattenRelativePosition(roiDataTable, true);
    % replace relPos_ROIName with relPos_ROIName_x, relPos_ROIName_y and
    % relPost_ROIName_z
    variablesToConsider = variablesToConsider(~startsWith(variablesToConsider, "relPos"));
    variablesToConsider = [variablesToConsider, roiDataTable.Properties.VariableNames(startsWith(roiDataTable.Properties.VariableNames, 'relPos_'))];
end

roiDataTableReduced = roiDataTable(:, [variablesToConsider, "ROIName"]);
roiDataTableReduced.originalRowIndex = (1:height(roiDataTableReduced))'; % this is done to keep track of the original index in the datatable later on

subRoiTables = subTableByGroup(roiDataTableReduced);
rowsToRemove = false(height(roiDataTable), 1); % logical mask for removal

%% Find outliers per ROI group
for iRoi = 1:length(subRoiTables)
    % Standardize with z-scores if z-score is higher than threshold value
    % classify datapoint as outlier and replace with NaN. NaN value will
    % later be imputed.
    roiTable = subRoiTables{iRoi};
    
    % Note that in cases where all the data is the same the normalize
    % returns NaN (because std = 0 and dividing by zero is not so
    % nice..right?!).
    standardizedRoiData = normalize(roiTable(:, variablesToConsider));

    for iVar = 1:width(standardizedRoiData)
        % Find all cells with absolute value above z-score threshold
        outlierIndex = abs(standardizedRoiData.(iVar)) > threshold;

        % We do not want to impute any parameters at this stage. We only
        % remove outliers for the center of mass outliers.
        
        % Case 1: You only want the specific parameter that is causing the
        % issue to be set to nan % standardizedPCAData{indices, iVar} =
        % nan;

        % Case 2: You want the entire ROI row to be set to nan if we find
        % an outlier. This is the way it was originally handled in the
        % first version.
        if any(outlierIndex)
            % Mark original rows for removal
            ogRowIndices = roiTable.originalRowIndex(outlierIndex);
            % logical mask to remove all at once later. 
            % This is a bit faster and arguably more "elegant".
            rowsToRemove(ogRowIndices) = true; 
        end
    
        % Case 3: Impute the ROI row instead of removing it TODO: impute
        % this value differently with RSS maybe or pcaambr?
        % standardizedPCAData.(iVar) =
        % fillmissing(standardizedPCAData.(iVar), 'linear');
        
    end
end
roiDataTable(rowsToRemove, :) = [];
roiDataUpdated = table2struct(roiDataTable);
end

function [subRoiTables] = subTableByGroup(table)
    % Split up the table logTransformRoiPCAData into the ROI groups
    uniqueRois = unique(table.ROIName);
    
    % Initialize a cell array to hold the subtables
    subRoiTables = cell(numel(uniqueRois),  1);
    
    % Loop through each unique group
    for iRoi =  1:numel(uniqueRois)
        % Find rows in the original table that belong to the current group
        groupRows = ismember(table.ROIName, uniqueRois{iRoi});
        
        % Create a subtable for the current group
        subRoiTables{iRoi} = table(groupRows, :);
    end
end