function Data=TranslationCorrection(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;
if ~exist('Model','var')
  Model=Data;
  ModelOpt=true;
end

wasStruct = false;
if isstruct(Data)
    Data = struct2table(Data);
    wasStruct = true;
end

% Split up the table logTransformRoiPCAData into the ROI groups
uniqueRois = unique(Data.ROIName);
% TODO check if unique(Data.RoiName) = unique(Model.RoiName)
nRois = length(uniqueRois);
M = containers.Map(1:nRois, uniqueRois);

% Initialize a cell array to hold the subtables
subRoiTablesData = cell(numel(uniqueRois),  1);
subRoiTablesModel = cell(numel(uniqueRois),  1);

% 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

  meanRowTable = mean(subRoiPosValTable, "omitnan");
  meanRowTable.Properties.RowNames = {'Mean'};
  stdRowTable = std(subRoiPosValTable, 0, "omitnan");
  stdRowTable.Properties.RowNames = {'Std'};
  
  
  % TODO: is the following comment correct about what x=1 refers to?
  % set initial value prior optimisation to 1
  % and start the fminunc optimisation
  x=1; 
  options = optimoptions('fminunc','Algorithm','trust-region','SpecifyObjectiveGradient',true,'Display','none');
  x = fminunc(@(x) CostRelativePos(x, subRoiPosValTable, subRoiPosValTable), zeros(size(subRoiPosValTable,1),1), options);

  subRoiPosValTable = subRoiPosValTable - x;
  while max(abs(x))>.5 && ModelOpt
    x = fminunc(@(x) CostRelativePos(x, subRoiPosValTable, subRoiPosValTable),zeros(size(subRoiPosValTable,1),1),options);
    subRoiPosValTable=subRoiPosValTable-x;
  end
  
  % 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(subRoiPosValTable(patient_idx(kk), roi_idx(kk)));
    Data(kk, :) = dataRow;
  end
end

if wasStruct
    Data = table2struct(Data);
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;

% Standardize using Z-score normalization
standardized =(translatedData-table2array(mean(Model,'omitnan')))./table2array(std(Model,0,'omitnan'));

% Component wise square
sSquared=standardized.*standardized;

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

% Calculates the gradient of the cost function.
g=-2.*standardized./table2array(std(Model, 0, 'omitnan'));
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