function [StudyOverview] = CollectDataForModel(data, Modality, roiList, mapCell, mergeMapCell, includePatientInfo, includeScanInfo)
% COLLECTDATAFORMODEL Collects and organizes data for modeling from DICOM files.
%
%   This function processes DICOM files located at the specified dataPath to
%   collect and organize data for modeling. It supports different modalities and
%   allows for the inclusion of patient and scan information based on the provided
%   flags. The function performs sanity checks to ensure the integrity of the
%   data and calculates basic features from the specified ROIs. It also warns
%   the user if any ROIs are not found in the RTStruct.
%
%   Syntax
%   [StudyOverview] = CollectDataForModel(dataPath, Modality, roiList, includePatientInfo, includeScanInfo)
%
%   Inputs
%   dataPath - Path to the directory containing DICOM files.
%   Modality - Modality of the medical imaging (e.g., 'MR', 'CT').
%   roiList - List of regions of interest (ROIs) to be analyzed. If not
%   defined we will use all ROI names mentioned in column A of the mapCell cell that has the information coming from e.g. an excel spreadsheet.
%   mapCell - cell-wise information coming from e.g. an Excel-Ark containing mappings
%   of RoiNames. Look at example below for how the structure of this cell
%   should look like
%   mergeMapCell - cell information coming from e.g. an Excel-Ark containing
%   ROI names to use for additional merged ROIs. Look at example below for
%   file structure notes
%   includePatientInfo - Boolean flag indicating whether to include patient information.
%   includeScanInfo - Boolean flag indicating whether to include scan information.
%
%   Outputs
%   StudyOverview - A structure containing an overview of the study, including
%                   patient information, scan information, and ROI information.
%
%   Example
%   dataPath = 'C:\Users\UDA2UH\Desktop\Prostate_MRL_Testing_Cohort\GoldSegmentation';
%   Modality = 'MR';
%   roiList = {'ROI1', 'ROI2', 'ROI3'};
%   mapCell = readcell('Mapping.xlsx'); This file can contain multiple rows.
%   Column A contains the "desired" name for a number of ROIs. Columns B -
%   end contain any occuring Roi name in the data that you want to map to
%   the "desired" name. There is also an example excel spreadsheet in the
%   git repository.
%   mergeMapCell = readcell('Merge.xlsx'); This file can contain multiple rows.
%   Column A contains the "desired" name for the newly created structure.
%   Columns B - end contain the names of the structures you want to use for
%   the merge process. Note that the names should contain the
%   "corrected"/"desired" name as the merging happens after the mapping of
%   ROI names.
%   includePatientInfo = true;
%   includeScanInfo = true;
%   [StudyOverview] = CollectDataForModel(dataPath, Modality, roiList, includePatientInfo, includeScanInfo);
%
%   See also
%   SortDicomFiles_OUH : Function to sort DICOM files.
%   LoadDicomData_OUH : Function to load DICOM data.
%   DicomLinks_OUH : Function to link DICOM files.
%
%   Notes
%   The function may take a long time to execute, especially for large datasets. 
%   The function includes placeholders for future enhancements, 
%   such as handling duplicate structures and adding missing tumor information.
%
%   Copyright (c) 2024, Odense University Hospital
    nin = nargin;
    % Check if boolOne and boolTwo are provided
    if nin < 7
        includeScanInfo = false; % Default value for includeScanInfo
    end
    if nin < 6
        includePatientInfo = false; % Default value for includePatientInfo
    end

    % Load DicomData from the dataPath
    % This step can take quite long.
    % So alternatively we can also pass an already loaded files variable.
    if iscell(data)
        files = data;
    elseif ischar(data) || isstring(data)
        files = SortDicomFiles_OUH(data);
    end
    % load('Files_HN.mat', 'files')
    % save('Files_Backup.mat','files')
    
    % TODO: Check for duplicate patient/study. How?
    
    % Sanity Checks 
    % Check if the lengths of the two image modality are equal
    % Modalities checks CT/MR vs RTStruct length
    if length(files{5,2}) + length(files{3,2}) + length(files{8,2}) ~= length(files{4,2})
        warning('Number of RTStructs and CT/MR files are not the same in the given dataPath')
    end
    % TODO change nStudy to RTStruct
    
    nStudy = length(files{4,2});
    
    StudyOverview = struct('StudyInstanceUID', {}, 'PatientInformation', {}, 'ScanInformation', {}, 'ROIInformation', {}, 'FileInformation', {});
    
    for iStudy = 1:nStudy
    fprintf('Starting data analysis for row %i in filelist.\n', iStudy)     
        % The "RTStruct"-variable contains the Dicom Series of the RTStruct.
        % It is structured as a 1x3 cell with the individual cells containing:
        % Cell 1: A 1x1 cell containing a struct with the RTStruct header (i.e.
        %   fieldnames are Dicom header tags and field values are Dicom header
        %   values).
        % Cell 2: A Nx2 cell (where N is the number of ROIs defined in the
        %   RTStruct). In {M,1} the ROIName of the M-th ROI is written. {M,
        %   2} contains a Sx1 cell where the DataPoints for the ROI definition 
        %   per slice is defined and S is the number of slices where it was
        %   defined (i.e. number of (3006, 0016) Contour Image Sequences for
        %   the ROIName/ROINumber). The content of each cell is a double-array
        %   with the length T where T refers to tag 3006, 0046 and is the
        %   number of contour Points for the slice.
        rtstructIndex = find(strcmp(files(:,1), 'RTSTRUCT'));
        RTStruct=LoadDicomData_OUH(files,rtstructIndex,iStudy);
        indexScan = DicomLinks_OUH(files,rtstructIndex,iStudy);
        
        % Find the affiliated scan with the RTStruct in the
        % SortLinkDicom overview
        nonEmptyRows = ~cellfun(@isempty, indexScan(:,2));
        notRTStructRow = ~strcmp(indexScan(:,1), 'RTSTRUCT');
        rowsToKeep = nonEmptyRows & notRTStructRow;
        indexScan = indexScan(rowsToKeep, :);

        % The "Scan"-variable contains the Dicom Series of the ImageScan
        % It is structured as a 1x5 cell with the individual cells containing:
        % Cell 1: A cell with the first row populated and the length equal to 
        %   the amount of Dicom imagefiles in the series. 
        %   Each element of the row contains a 1x1 struct.
        %   The struct's fieldnames are the names of the DicomTags in the
        %   Header.
        % Cell 2: The 3D image matrix 
        % Cell 3 - 5: Irrelevant for this use case
        try
            if strcmp(Modality, 'MR')
              mrIndex = find(strcmp(files(:,1), 'MR'));
              Scan=LoadDicomData_OUH(files,mrIndex,indexScan{1,2});
            elseif strcmp(Modality, 'CT')
              ctIndex = find(strcmp(files(:,1), 'CT'));
              Scan=LoadDicomData_OUH(files,ctIndex,indexScan{1,2});
            end
        catch ME
            warning('Study with non-matching Modality found. You might have tried to analyse a CT image even though you defined MR as the modality or vice-versa. Sneaky...sneaky. But unfortunately not possible.')
            continue;
            % rethrow(ME)
        end
        
        % Check if scan image and RTStruct are affiliated with each other.
        % If not throw error.
        if ~strcmp(Scan{1,1}{1,1}.StudyInstanceUID, RTStruct{1,1}{1,1}.StudyInstanceUID)
            assert(false, ['Scan and RTStruct are not related to the same study. ' ...
                'Check if the Scan and RTStruct were imported in the same order ' ...
                '(i.e. compare Study Instance ID of Scan and RTStruct)'])
        end
        
        % Collect Patient Specific Data
        if includePatientInfo
          PatientData = struct( ...
              'PatientID', {''}, ...
              'PatientSex', {''}, ...
              'PatientAge', {''}, ...
              'PatientWeight', {''}, ...
              'PregnancyStatus', {''});
          patientParameterFields = fieldnames(PatientData);
          
          tic;
          for iParameter = 1:numel(patientParameterFields)
              fieldName = patientParameterFields{iParameter};
              try
                  PatientData.(fieldName) = Scan{1,1}{1,1}.(fieldName);
              catch
                  warning('Error. Did not find PatientData fieldName %s.\n', fieldName) 
                  PatientData.(fieldName) = NaN;
              end
          end
          elapsedTime = toc; % Stop the timer and get the elapsed time
          fprintf('Time for PatientParameter: %.2f seconds\n', elapsedTime);
        end

            % Define base structure (common fields for all modalities)
            ScanData = struct(...
                'StudyInstanceUID', '', ...
                'SeriesInstanceUID', '', ...
                'Modality', '', ...
                'Manufacturer', '', ...
                'SeriesDescription', '', ...
                'StudyDate', '', ...
                'ManufacturerModelName', '', ...
                'PatientPosition', '', ...
                'ImageOrientationPatient', '' ...
            );
            
            % Add modality-specific fields
            switch Modality
                case 'CT'
                    ScanData.KVP = '';
            
                case 'MR'
                    ScanData.MRAcquisitionType = '';
                    ScanData.SliceThickness = '';
                    ScanData.RepetitionTime = '';
                    ScanData.EchoTime = '';
                    ScanData.NumberOfAverages = '';
                    ScanData.ImagingFrequency = '';
                    ScanData.MagneticFieldStrength = '';
                    ScanData.SpacingBetweenSlices = '';
                    ScanData.NumberOfPhaseEncodingSteps = '';
                    ScanData.EchoTrainLength = '';
                    ScanData.ReceiveCoilName = '';
                    ScanData.FlipAngle = '';
            end
                    
          
          scanParameterFields = fieldnames(ScanData);
          
          tic;
          for iParameter = 1:numel(scanParameterFields)
              fieldName = scanParameterFields{iParameter};
              try
                  ScanData.(fieldName) = Scan{1,1}{1,1}.(fieldName);
              catch
                  warning('Did not find PatientData fieldName %s.\n', fieldName) 
                  ScanData.(fieldName) = NaN;
              end
          end
          elapsedTime = toc; % Stop the timer and get the elapsed time
          fprintf('Time for ScanParameter: %.2f seconds\n', elapsedTime);

    
        
        % Start of pre-processing for calculate basic features from ROIs
        
        % Check for roi Names and if there are any mappings of roiNames
        % that need to be applied
        if nin > 4
          RTStruct{2}(:, 1) = correctRoiNames(RTStruct{2}(:, 1), mapCell);
        end

        % Check for roi Names and if there is an intersection between the
        % roiList array or not. If not use all the ROIs defined in the
        % RTStruct
        % DEPRECATED: This was for the version where you had to give a
        % manually typed roiList as input to the function.
        % if isempty(roiList)
        %   overviewRois = RTStruct{2}(:, 1);
        % else
        %   overviewRois = intersect(roiList, RTStruct{2}(:, 1));
        % end
        
        % if roiList is empty take the names of the provided values in the
        % mapping excel spreadsheet.
        if isempty(roiList)
          mappingList = mapCell;
          overviewRois = mappingList(:,1);
        else
          overviewRois = intersect(roiList, RTStruct{2}(:, 1));
        end

        notFoundRois = setdiff(roiList, RTStruct{2}(:, 1));
        if ~isempty(notFoundRois)
          notFoundRoisStr = strjoin(notFoundRois, ', ');
          warning('Could not find the ROI(s) "%s" in the RTStruct', notFoundRoisStr);
        end
        ROIBasicFeatures = struct();
        % ROIAdvancedFeatures = struct();
    
        % TODO: check for duplicate structures
        % Loop through the cell array to add fields to the structure
        for i = 1:numel(overviewRois)
            try
                validFieldName = matlab.lang.makeValidName(overviewRois{i});
                ROIBasicFeatures.(validFieldName) = NaN;
            catch e
                fprintf(1,'The identifier was:\n%s',e.identifier);
                fprintf(1,'There was an error! The message was:\n%s',e.message);
                fprintf('Skipping...\n')
            end
        end
        
        for iRois = 1:numel(overviewRois)
    
            % Calculate basic features from ROIs
            tic; % Start the timer
            
            matchingRoi = cellfun(@(x) strcmp(x, overviewRois{iRois}), RTStruct{2}(:, 1), 'UniformOutput', true);
            if any(matchingRoi)
              % if any(strcmp(RTStruct{iRois}, overviewRois)
              indexRoi = find(matchingRoi);
              validFieldName = matlab.lang.makeValidName(overviewRois{iRois});
              mask=StructMask_OUH(RTStruct,Scan,indexRoi);
              
              try
                  ROIBasicFeatures.(validFieldName) = CalcBasicFeatures(Scan,mask);
                  % TODO: add Advanced features with Matlabs radiomatics.
                  % ROIAdvancedFeatures.(validFieldName) = CalcAdvancedFeatures(Scan, mask);
              catch e
                  fprintf(1,'The identifier was:\n%s',e.identifier);
                  fprintf(1,'There was an error! The message was:\n%s',e.message);
                  ROIBasicFeatures.(validFieldName) = NaN;
              end
              elapsedTime = toc; % Stop the timer and get the elapsed time
              fprintf('Time for Basic Feature Calculations %d (%s): %.2f seconds\n', iRois, overviewRois{iRois}, elapsedTime);
            else
              fprintf('No structure found for %s\n', overviewRois{iRois});
            end
        end

        if ~isempty(mergeMapCell)
          mergeRoiCell = mergeMapCell;
          for iMerge = 1:size(mergeRoiCell, 1)
            roisToMerge = mergeRoiCell(iMerge, 2:end);
            newRoiName = mergeRoiCell{iMerge, 1};
            [~,Locb] = ismember(mergeRoiCell(iMerge, 2:end), RTStruct{2}(:, 1));
            fprintf('Merging ROIs into %s:\n', newRoiName);
            fprintf('%s\n', string(roisToMerge));
            if all(Locb)
                mask = StructMask_OUH(RTStruct,Scan, Locb);
                validFieldName = matlab.lang.makeValidName(newRoiName);
    
                try
                    ROIBasicFeatures.(validFieldName) = CalcBasicFeatures(Scan,mask);
                catch e %e is an MException struct
                    fprintf(1,'The identifier was:\n%s',e.identifier);
                    fprintf(1,'There was an error! The message was:\n%s',e.message);
                    ROIBasicFeatures.(validFieldName) = NaN;
                end
                elapsedTime = toc; % Stop the timer and get the elapsed time
                fprintf('Time for Basic Feature Calculations %s: %.2f seconds\n', validFieldName, elapsedTime);
            else
                fprintf('Not all ROIs available required for merging. Skipping...\n');
            end
          end
        end
        
    
        % Collect data into one overview structure
        StudyOverview(iStudy).StudyInstanceUID = Scan{1,1}{1,1}.StudyInstanceUID;
        StudyOverview(iStudy).PatientInformation = PatientData;
        StudyOverview(iStudy).ScanInformation = ScanData;
        StudyOverview(iStudy).ROIInformation = ROIBasicFeatures;
        StudyOverview(iStudy).FileInformation = {'Scan', cellfun(@(x) x.Filename, Scan{1,1}, 'UniformOutput', false); ...
            'RTStruct', RTStruct{1,1}{1,1}.Filename};
        
        %TODO: potentially add missing tumour information (i.e. stage, volume, classification)
    end
end

function [roiData] = correctRoiNames(roiData, mapCell)
% CORRECTROINAMES
% Based on a mapping datastructure obtained via a excel spreadsheet
% roiNames in the provided data are corrected.
% (e.g. Prostata -> Prostate).
% The excel spreadsheet needs to be structures as follows:
% Each row is ROI specific with the first column containing the "desired"
% ROI name, while the other columns contain the "existent but not desired"
% ROI names. The "existent but not desired" ROI names are changed into the
% "desired" ROI name.
% NOTE: As opposed to the inference ROI name correction via the
% Damerau-Levenshtein distance, this method is not as robust to tiny errors
% in the mapping. This function ONLY corrects what you tell it to correct.
% 
% E.g.:
% Prostate_AI	Prostata	Prostate	Prosttate	Prostata_AI
% FemoralHead_AI	FemoralHead	Femoral_Head		
% Sigmoid_AI	Sigmo	Sigmi	Sigmoaid	
%
% This will ONLY correct Prostata, Prostate, Prosttate, Prostata_AI to
% Prostate_AI. Same goes for FemoralHead_AI and Sigmoid_AI.
% Thus you have to be aware of all the deviations in the namegiving of your
% data yourself!


  mapROITable = mapCell;
  wrongRoiNames = mapROITable(1:end,2:end);
  wrongRoiNames(cellfun(@(x) any(ismissing(x)), wrongRoiNames)) = {''};
  
  for iRow = 1:length(roiData)
    [Lia,Locb] = ismember(roiData{iRow}, wrongRoiNames);
    if Lia
      rowIndex = mod(Locb, length(wrongRoiNames));
      if rowIndex == 0
          rowIndex = size(wrongRoiNames,1);
      end
      roiData{iRow} = mapROITable{rowIndex,1};
    end
  end
end