function [Model]=BuildModel(logTransformRoiPCAData, varargin)
% BUILDMODEL Constructs a model for ROI PCA data processing and analysis.
%
% SYNTAX
% [Model] = BuildModel(logTransformRoiPCAData)
%
% INPUTS
% logTransformRoiPCAData - A table containing log-transformed ROI PCA data obtained from the CalculateParameterValuesForInference method.
% varargin - Optional parameters:
%   - 'averageDirectionOutput' (string): Path to a .mat file containing the average direction data.
%   - 'modelOutput' (string): Path to save the output model.
%   - 'threshold' (numeric): Value that defines the Z-score threshold. All
%   ROIs that have for any parameter a Z-score higher than this value will
%   be removed from the model building. Defaults to 3.28.
%
% OUTPUTS
% Model - A struct containing the processed model data, including PCA analysis results.
%
% DESCRIPTION
% This function processes log-transformed ROI PCA data to remove outliers, apply
% translation corrections, and perform PCA analysis. The output is a struct that
% includes the final model data, PCA coefficients, scores, eigenvalues, explained
% variance, mean PCA, and t-squared values for each ROI. The function also saves
% the final model data to a .mat file for further analysis.
%
% The processing steps include:
% 1. Grouping the data by ROI and standardizing each group using z-scores.
% 2. Identifying and replacing outliers with NaN values.
% 3. Applying translation corrections to the standardized data.
% 4. Imputing missing values in the position matrix.
% 5. Performing PCA analysis on the processed data for each ROI.
%
% The function also handles the saving of intermediate and final results to disk.
%
% EXAMPLE
% % Assuming 'logTransformRoiPCAData' is a table containing log-transformed ROI PCA data.
% Model = BuildModel(logTransformRoiPCAData);
%
% NOTES
% - The function uses a hardcoded threshold for outlier detection based on z-scores.
% - The PCA analysis includes the calculation of eigenvalues, explained variance, and t-squared values.
% - The final model data is saved to 'ModelData_MRLProstateGoldSeg.mat'.
%
% Copyright (c) 2024, Odense University Hospital - Laboratory of Radiation Physics

% Set up input parser
inputParam = inputParser;
inputParam.addRequired('logTransformRoiPCAData', @(x) istable(x));
inputParam.addOptional('averageDirectionOutput', '', @(x) ischar(x) || isstring(x) || isstruct(x));
inputParam.addOptional('modelOutput', '', @(x) ischar(x) || isstring(x));
inputParam.addOptional('threshold', 3.28, @(x) isnumeric(x));
inputParam.addOptional('verbose', false, @(x) islogical(x))
inputParam.parse(logTransformRoiPCAData, varargin{:});

averageDirectionOutput = inputParam.Results.averageDirectionOutput;
modelOutput = inputParam.Results.modelOutput;
threshold = inputParam.Results.threshold;
verbose = inputParam.Results.verbose;

% Set-up Model struct
flattenedRoiData = table();
Model = struct();

subRoiTables = subTableByGroup(logTransformRoiPCAData);

%% Remove model outliers
% This code is very similar to RemoveModelOutliers. However we have the
% extra step that we store the mean and std values of the normalization in
% a struct. We do this because we need the mean and standard deviation
% later during the inference run.
for i = 1:length(subRoiTables)
    % Standardize with z-scores
    % Done on the log-transformed ROI PCA Data grouped by ROI now !!
    % if z-score is higher than threshold value classify datapoint as outlier 
    % and replace with NaN. NaN value will later be imputed.
    % TODO: How does z-score relate to p-value and how to set threshold?
    % One-tailed or two-tailed for p-value?
    roiTable = subRoiTables{i};
    
    % 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?!). But we still want to keep the mean and std values for
    % inference runs later.
    [standardizedPCAData, meanData, meanStd] = normalize(roiTable(:, 4:end-2), 'zscore');
    
    standardizedPCAData = standardizeMissing(standardizedPCAData, NaN);
    standardizedPCAData = fillmissing(standardizedPCAData, 'constant', 0);

    standardizeValues.(roiTable.ROIName{1}).Mean = meanData;
    standardizeValues.(roiTable.ROIName{1}).Std = meanStd;
    islandsLambdaValues.(roiTable.ROIName{1}) = unique(roiTable.IslandsLambda);

    % % PosVal Means for Translation Correction during Inference. We want to
    % % have the center of mass values specifically (and prior to
    % % standardization).
    % posVal.(roiTable.ROIName{1}).Mean = [meanData.xcm; meanData.ycm; meanData.zcm];
    % posVal.(roiTable.ROIName{1}).Std = [meanStd.xcm; meanStd.ycm; meanStd.zcm];
    
    % DEPRECATED: Because we do throw away model outliers already at an
    % earlier stage
    % % threshold = 3.28; % deperecated. set as kwargs inputs
    % for iVar = 1:length(standardizedPCAData.Properties.VariableNames)
    %     % Find all cells with absolute value above z-score threshold
    %     indices = find(abs(standardizedPCAData.(iVar)) > threshold);
    % 
    %     % we do not want to impute any parameters at this stage. This will
    %     % happen later with TranslationCorrection and ImputePositionMatrix.
    %     % We only remove outliers for the center of mass outliers.
    %     Lia = ismember(string(standardizedPCAData.Properties.VariableNames{iVar}), ["xcm", "ycm", "zcm"]);
    %     if isa(indices, 'double') && any(Lia)
    %       % 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.
    %       standardizedPCAData{indices, :} = nan;
    % 
    %       % TODO: impute this value differently with RSS maybe or pcaambr?
    %       % standardizedPCAData.(iVar) = fillmissing(standardizedPCAData.(iVar), 'linear');
    %     end
    % 
    % end
    logTransformStandardizedRoiPCAData = [roiTable(:, 1:3), standardizedPCAData];
    % logTransformStandardizedRoiPCAData.PosVal = [logTransformStandardizedRoiPCAData.xcm, logTransformStandardizedRoiPCAData.ycm, logTransformStandardizedRoiPCAData.zcm];
    
    flattenedRoiData = [flattenedRoiData; logTransformStandardizedRoiPCAData];
end

% G = findgroups(flattenedRoiData.ROIName);
% roiGroupedPosVal = splitapply(@(x) table(x, 'VariableNames', {'PosVal'}), flattenedRoiData.PosVal, G);

% TranslationCorrection creates a new column called PosValTranslated, which
% takes PosVal as starting point and optimizes the position of the center
% of mass. We do this to adjust for slight offsets of the patient on the
% scan couch.
%
% DEPRECATED: We switched from xcm, ycm, zcm to using relative
% position
% translatedPCAData = TranslationCorrection(flattenedRoiData);
% 
% finalModelData = ImputePositionMatrix(translatedPCAData);
% 
% finalModelData.xcm = finalModelData.PosValTranslated(:,1,:);
% finalModelData.ycm = finalModelData.PosValTranslated(:,2,:);
% finalModelData.zcm = finalModelData.PosValTranslated(:,3,:);
%
% subModel = subTableByGroup(finalModelData);
subModel = subTableByGroup(flattenedRoiData);

for iModel = 1:length(subModel)
    model = subModel{iModel};
    roiName = matlab.lang.makeValidName(model.ROIName{1});
    
    % % Isolation Tree
    % [Mdl, ~, scores] = iforest(model(:,4:20), 'NumLearners', 500);
    % ModelContent = struct('Model', Mdl, 'Scores', scores, 'standardizeScales', standardizeValues.(roiName));

    % PosVal Means for Translation Correction during Inference
    % posVal.Mean = [mean(model.xcm); mean(model.ycm); mean(model.zcm)];
    % posVal.Std = [std(model.xcm); std(model.ycm); std(model.zcm)];

    % In a lot of cases we won't find any islands or the exact amount of
    % islands for a given ROI. Thus the IslandsLikelihood parameter will be
    % identical for all occurances of that given ROI. As we normalize 
    % (Z-score) the data earlier this leads to an issue where the Z-score
    % is NaN for these cases. As the following PCA calculation can't handle
    % an entire column (aka dimension) to be NaN, we set the values to 0.
    % NOTE: This could be considered an adequate strategy in general. So we
    % might want to apply this procedure to other parameters except just
    % IslandsLikelihood. However, at the moment I don't see a reason to do
    % this because I can't see a scenario where any of the other values
    % will be NaN at this stage.
    if any(isnan(model.IslandsLikelihood))
        model.IslandsLikelihood = fillmissing(model.IslandsLikelihood, 'constant', 0);
    end
    if any(isnan(model.IslandsLambda))
        model.IslandsLambda = fillmissing(model.IslandsLambda, 'constant', 1e-5);
    end

    % PCA
    mainVars = {'Volume', 'lengthPCA1', 'lengthPCA2', 'lengthPCA3', 'Area', ...
        'Compactness', 'CMToSurface', 'MaxLength', 'PCAalpha', 'PCAbeta', 'PCAgamma', ...
        'RatioPCA21', 'RatioPCA31', 'Solidity'};
    [coeff, score, latent, tsquared, explained, mu] = pca(table2array(model(:,mainVars))); % PCA to create eigenvectors for geometrical/shape features

    availableFieldNames = fieldnames(model);
    relPosVars = availableFieldNames(startsWith(availableFieldNames, "relPos_"))';
    [coeffRelPos, scoreRelPos, latentRelPos, tsquaredRelPos, explainedRelPos, muRelPos] = pca(table2array(model(:,relPosVars))); % PCA to create eigenvectors for relative position features
    relPosContent = struct('Eigenvectors', coeffRelPos, 'Score', scoreRelPos, 'Eigenvalues', latentRelPos, 'ExplainedVariance', explainedRelPos, 'MeanPCA', muRelPos, 'TSquared', tsquaredRelPos);

    ModelContent = struct('Eigenvectors', coeff, 'Score', score, 'Eigenvalues', latent, 'ExplainedVariance', explained, 'MeanPCA', mu, 'TSquared', tsquared, 'standardizeScales', standardizeValues.(roiName), 'IslandLambda', islandsLambdaValues.(roiName), 'RelativePosition', relPosContent);

    Model.(roiName) = ModelContent;

    notConsideredVars = {'IslandsLikelihood', 'IslandsLambda', 'Islands'}; % no PCA calculation for these
    if ~isequal(length(ModelContent.Eigenvectors)+length(ModelContent.RelativePosition.Eigenvectors)+length(notConsideredVars), ... 
            length(ModelContent.MeanPCA)+length(ModelContent.RelativePosition.MeanPCA)+length(notConsideredVars), ...
            width(ModelContent.standardizeScales.Mean), width(ModelContent.standardizeScales.Std)) & verbose
        warning('Model Contents for ROI %s do not have the same size. There is something rotten in the state of Denmark...', roiName)
    end
end
if ~isempty(averageDirectionOutput)
    if ~isstruct(averageDirectionOutput)
      load(averageDirectionOutput, 'averageDirection');
    else
      averageDirection = averageDirectionOutput;
    end
    Model.averageDirection = averageDirection;
end
if ~isempty(modelOutput)
  save(modelOutput,'Model')
end

disp('Model building done!')  
end

function find95VarianceExplained(arrayPCAData, explained)
    % Plot the explained variance ratio
    figure;
    plot(1:size(arrayPCAData,  2), cumsum(explained)/sum(explained), 'o-', 'LineWidth',  2);
    hold on;
    
    % Add horizontal line for  95% cut-off threshold
    plot([0, size(arrayPCAData,  2)], [0.95,  0.95], '--', 'Color', 'red');
    text(size(arrayPCAData,  2),  0.95, '95%', 'Color', 'red', 'FontSize', 12);
    
    % Set plot properties
    xlabel('Number of Components');
    ylabel('Cumulative variance (%)');
    
    grid on;
    grid minor;
    grid minor;
    
    % Adjust the x-axis ticks
    xticks(1:size(arrayPCAData,  2));
    
    % Show the figure
    hold off;

    % Save the figure
    dir = "C:\Users\UDA2UH\Downloads\internal-research-tools-master\PCA_Outlier_Model\RunInference\figures\explained_variance";
    fullFileName = sprintf('%s\\explained_variance_%s.png', dir, char(roiParameterGrouped{1,1}));
    saveas(gcf, fullFileName, 'png');
    % scatter3(score(:,1), score(:,2), score(:,3));
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 i =  1:numel(uniqueRois)
        % Find rows in the original table that belong to the current group
        groupRows = ismember(table.ROIName, uniqueRois{i});
        
        % Create a subtable for the current group
        subRoiTables{i} = table(groupRows, :);
    end
end