function Data=TranslationCorrectionInference(Data,Model)
% TRANSLATIONCORRECTION Corrects the position values of ROIs by aligning them with a reference model.
%
% SYNTAX:
% Data = TranslationCorrection(Data, Model)
%
% INPUTS:
% Data - A structure containing ROI data, including position values.
% Model - (Optional) A structure containing model data. If not provided, Data is used as the model.
%
% OUTPUTS:
% Data - The input Data structure with corrected position values.
%
% DESCRIPTION:
% This function corrects the position values of regions of interest (ROIs) in a dataset by
% aligning them with a reference model or dataset. It iteratively optimizes the translation
% vector for each spatial dimension (x, y, z) to minimize the cost function, which measures the
% difference between the test and model position values. The corrected position values are
% stored back into the original Data structure.
%
% The cost function is calculated using the CostRelativePos function, which normalizes the
% position differences by subtracting the mean and dividing by the standard deviation of the
% model data, and then computes the sum of squared differences between the normalized test and
% model positions.
%
% Copyright (c) 2024, Odense University Hospital - Laboratory of Radiation Physics

ModelOpt=false;

% Split up the table logTransformRoiPCAData into the ROI groups
uniqueRoisData = unique(Data.ROIName);
uniqueRoisModel = fieldnames(Model);

% TODO check if unique(Data.RoiName) = unique(Model.RoiName)
nRois = length(uniqueRoisData);
M = containers.Map(1:nRois, uniqueRoisData);

% merge center of masses into one Table column so it gets easier to handle
% in the loop containing fminunc
Data.PosVal = [Data.xcm, Data.ycm, Data.zcm];
Data.PosValTranslated = NaN(height(Data),3);

% Loop over the three spatial dimensions
% We are going to look at the center of masses of:
% x (i.e. left-right in patient coordinate system)
% y (i.e. anterior-posterior in patient coordinate system)
% z (i.e. caudal-cranial in patient coordinate system)
for i=1:3

  % save the data and only the data for the current selected spatial
  % dimension in a table form such that we obtain a NxM table where N is
  % the number of patients and M is the number of ROIs in the Model.
  % For patients where the ROI is not existent fill the PatientxROI cell
  % with NaN
  [roi_idx, roiName] = findgroups(Data.ROIName);
  [patient_idx, patientID] = findgroups(Data.PatientID);

  initSubRoiPosValTable = [repmat({NaN(numel(patientID),1)},1,numel(roiName))];
  subRoiPosValTable = table(initSubRoiPosValTable{:},'VariableNames', roiName.');
  subRoiPosValTable.Properties.RowNames = patientID;

  for iRow = 1:size(Data,1)
    subRoiPosValTable{patient_idx(iRow),roi_idx(iRow)} = Data(iRow, :).PosVal(i);
  end

  
  % We want to collect the Mean and Std value (i.e. x_cm, y_cm or zcm, depends on i)
  % for all Rois found in the to-be-inferred sample from the Model
  % object so we can use them in the optimisation function CostRelativePos.
  modelPosValMean = zeros(length(roiName), 0);
  modelPosValStd = zeros(length(roiName), 0);
  for iRoiColumn = 1:length(roiName)
      modelPosValMean(iRoiColumn) = Model.(roiName{iRoiColumn}).posVal.Mean(i);
      modelPosValStd(iRoiColumn) = Model.(roiName{iRoiColumn}).posVal.Std(i);
  end
  meanRow = num2cell(modelPosValMean);
  meanRowTable = cell2table(meanRow, 'VariableNames', roiName.');
  meanRowTable.Properties.RowNames = {'Mean'};
  stdRow = num2cell(modelPosValStd);
  stdRowTable = cell2table(stdRow, 'VariableNames', roiName.');
  stdRowTable.Properties.RowNames = {'Std'};
  
  % Start the fminunc optimisation
  options = optimoptions('fminunc','Algorithm','trust-region','SpecifyObjectiveGradient',true,'Display','none');
  % options = optimoptions('fminunc','PlotFcn','optimplotx','Display','iter');
  x = fminunc(@(x) CostRelativePos(x, subRoiPosValTable, [meanRowTable; stdRowTable]), zeros(size(subRoiPosValTable,1),1), options);
  posValCorrected = subRoiPosValTable - x;
  
  % Now we have to restructure the data returned from fminunc into the
  % original structure. I am really sorry for all of this back and forth.
  % Sometimes I beat Matlab. Sometimes Matlab beats me.
  for kk = 1:height(Data)
    dataRow = Data(kk, :);
    dataRow.PosValTranslated(:,i,:) = table2array(posValCorrected(patient_idx(kk), roi_idx(kk)));
    Data(kk, :) = dataRow;
  end
end
end

function [f,g]=CostRelativePos(V,Test,Model)
% COSTRELATIVEPOS Calculates the cost function and its gradient for aligning
% position values between a test dataset and a model dataset.
%
% SYNTAX:
% [f, g] = CostRelativePos(V, Test, Model)
%
% INPUTS:
% V - A vector representing the translation applied to the test dataset.
% Test - The test dataset containing position values to be aligned.
% Model - The model dataset containing reference position values.
%
% OUTPUTS:
% f - The cost function value, representing the sum of squared normalized
%     differences between the test and model position values.
% g - The gradient of the cost function, indicating the direction of change
%     needed to minimize the cost.
%
% DESCRIPTION:
% This function calculates the cost function for aligning position values
% between a test dataset and a model dataset using a Z-score normalization
% approach. It computes both the cost function value f and its gradient g,
% which are used in an optimization process to minimize the cost and align the
% position values as closely as possible.
%
% The cost function is calculated as the sum of the squared normalized
% differences between the test and model position values, and the gradient
% represents the direction in which the test values need to change to reduce
% the cost function.
%
% The function is designed to be used in an optimization algorithm, such as
% 'fminunc' in MATLAB, where the goal is to find the translation vector V
% that minimizes the cost function.
%
% Copyright (c) 2024, Odense University Hospital - Laboratory of
% Radiophysics

% Applying translation to dataset
translatedData=Test-V*ones(1,size(Test,2));

% Standardize using Z-score normalization
standardized =(translatedData-table2array(Model('Mean',:)))./(table2array(Model('Std',:)));

% Component wise square
sSquared=standardized.*standardized;

f=sum(sSquared, "all", 'omitnan');

% Calculates the gradient of the cost function.
g=-2.*standardized./table2array(Model('Std',:));
g=sum(g,2,'omitnan');

% Fminunc expects cost and gradient of cost to be returned as data matrix
f = table2array(f);
g = table2array(g);
end