function DataFinal = ImputePositionMatrix(Data)
%IMPUTEPOSITIONMATRIX Impute missing values in a position matrix using PCA and TSR regression.
%
%   This function takes a structure containing position data, reshapes the data into a
%   2D matrix, imputes missing values using a combination of PCA (Principal Component
%   Analysis) and TSR (Total Sum of Squares and Residuals) regression, reshapes the
%   imputed data back to its original dimensions, and updates the structure with the
%   imputed position data.
%
%   Syntax
%   Data = ImputePositionMatrix(Data)
%
%   Inputs
%   Data : A structure containing position data. The structure must have a field
%          named 'PosVal' which contains the position data matrix.
%
%   Outputs
%   Data : The input structure with the 'PosVal' field updated to contain the
%          imputed position data.
%
%   Example
%   --------
%   % Assume Data is a structure with position data
%   Data = ...; % Initialize Data structure
%   ImputedData = ImputePositionMatrix(Data);
%
%   See also
%   --------
%   pcambtsr : The function used for imputing missing values using PCA and TSR regression.
%
%   Notes
%   -----
%   This function assumes that the input structure contains a field named 'PosVal'
%   with the position data matrix. The position data matrix is reshaped into a
%   2D matrix for imputation and then reshaped back to its original dimensions.
%
%   Copyright (c) 2024, Odense University Hospital - Laboratory of
%   Radiophysics

Data = struct2table(Data);
DataReduced = Data(:, {'xcm', 'ycm', 'zcm', 'ROIName', 'PtNumber'});

% We are doing a Restructuring of the Data to the imputation and then back
% from the imputation. The reason for this is that we want to have a
% structure that is similar to this:
% Pt1 ROI1_x ROI1_y ROI1_z ROI2_x ROI2_y ROI2_z ...
%
% This serves the purpose that by that we have a table structure that
% represents the correlations between ROIs nicely which is exactly what we
% leverage during the imputation process.
DataRestructured = RestructureDataToImputation(DataReduced);

DataMat = table2array(DataRestructured(:, 2:end));
[DataMatImputed,~,~,~,~,cost]=pcambtsr(DataMat, round(size(DataRestructured,2)/9), 5000);
DataImputed = array2table([(1:height(DataRestructured))', DataMatImputed], 'VariableNames', DataRestructured.Properties.VariableNames);

DataImputedRestructured = RestructureDataFromImputation(DataImputed);
DataJoined = outerjoin(Data, DataImputedRestructured, ...
    "LeftKeys", ["PtNumber", "ROIName"], ...
    "RightKeys", ["PtNumber", "ROIName"], "MergeKeys", true);
DataFinal = removevars(DataJoined, {'xcm_Data', 'ycm_Data', 'zcm_Data'});
DataFinal = renamevars(DataFinal, {'xcm_DataImputedRestructured', 'ycm_DataImputedRestructured', 'zcm_DataImputedRestructured'}, ...
    {'xcm', 'ycm', 'zcm'});
DataFinal = movevars(DataFinal, "xcm", "Before", "lengthPCA1");
DataFinal = movevars(DataFinal, "ycm", "Before", "lengthPCA1");
DataFinal = movevars(DataFinal, "zcm", "Before", "lengthPCA1");
DataFinal = table2struct(DataFinal);
end

function DataReshaped = RestructureDataToImputation(Data)
  
    uniquePt = max(Data.PtNumber);
    uniqueRois = unique(Data.ROIName);

    reshapedData = [];

    for iPt = 1:uniquePt
        pt = unique(iPt);
        rowData = [];
        varNames = {};

        for iRoi = 1:length(uniqueRois)
            roi = uniqueRois{iRoi};
            idx = strcmp(Data.ROIName, roi) & Data.PtNumber == pt;

            if any(idx)
                x = Data.xcm(idx);
                y = Data.ycm(idx);
                z = Data.zcm(idx);
            else
                x = NaN;
                y = NaN;
                z = NaN;
            end

            rowData = [rowData, x, y, z];
            varNames = [varNames, ...
                [roi '_x'], [roi '_y'], [roi '_z']];
        end

        reshapedData = [reshapedData; rowData];
    end

    DataReshaped = array2table(reshapedData, 'VariableNames', matlab.lang.makeValidName(varNames));
    DataReshaped = addvars(DataReshaped, (1:uniquePt)', 'Before', 1, 'NewVariableNames', 'PtNumber');
end
        
function DataReshaped = RestructureDataFromImputation(Data)
    % Get all variable names
    varNames = Data.Properties.VariableNames;
    
    % Get list of ROIs from variable names (excluding PtNumber)
    roiVars = varNames(2:end);  % exclude PtNumber
    roiNames = unique(erase(roiVars, {'_x','_y','_z'}));  % unique ROI base names
    
    % Initialize cell array to hold rows
    DataReshaped = [];
    
    % Loop over each ROI and build rows
    for i = 1:length(roiNames)
        roi = roiNames{i};
    
        % Get column names for this ROI
        xCol = [roi '_x'];
        yCol = [roi '_y'];
        zCol = [roi '_z'];
    
        % Check if those columns exist
        if all(ismember({xCol, yCol, zCol}, varNames))
            % Extract data
            x = Data.(xCol);
            y = Data.(yCol);
            z = Data.(zCol);
            pt = Data.PtNumber;
    
            % Create table for this ROI
            tempTable = table(x, y, z, repmat({roi}, height(Data), 1), pt, ...
                'VariableNames', {'xcm', 'ycm', 'zcm', 'ROIName', 'PtNumber'});
    
            % Append to full table
            DataReshaped = [DataReshaped; tempTable];
        end
    end
    
    % Sort if needed (e.g. by PtNumber)
    DataReshaped = sortrows(DataReshaped, {'PtNumber', 'ROIName'});
end