data_evaluation_utils.py

"""Data Evaluation Functions

Description:

    This folder contains several functions which, either on their own or included in larger pieces of software, perform data evaluation tasks.

Usage:

    To use content from this folder, import the functions and instantiate them as you wish to use them:

        from utils.data_evaluation_utils import function_name

TODO: Might be worth adding some information on uncertaintiy estimation, later down the line

"""

import os
import pickle
import numpy as np
import torch
import logging
import h5py
import utils.data_utils as data_utils
from utils.common_utils import create_folder
import pandas as pd
from fsl.data.image import Image
from fsl.utils.image.roi import roi
import itertools
from scipy.spatial.distance import cosine
from scipy.stats import pearsonr, spearmanr

log = logging.getLogger(__name__)

def evaluate_data(trained_model_path,
                     data_directory,
                     mapping_data_file,
                     mapping_targets_file,
                     data_list,
                     prediction_output_path,
                     prediction_output_database_name,
                     prediction_output_statistics_name,
                     brain_mask_path,
                     dmri_mean_mask_path,
                     rsfmri_mean_mask_path,
                     regression_factors,
                     mean_regression_flag,
                     mean_regression_all_flag, 
                     mean_subtraction_flag,
                     scale_volumes_flag,
                     normalize_flag,
                     minus_one_scaling_flag,
                     negative_flag, 
                     outlier_flag,
                     shrinkage_flag,
                     hard_shrinkage_flag,
                     crop_flag,
                     device=0, 
                     exit_on_error=False,
                     output_database_flag=False,
                     cross_domain_x2x_flag=False,
                     cross_domain_y2y_flag=False,
                     mode='evaluate'):
                     
    """Model Evaluator

    This function generates the rsfMRI arrays for the given inputs

    Args:
        trained_model_path (str): Path to the location of the trained model
        data_directory (str): Path to input data directory
        mapping_data_file (str): Path to the input file
        mapping_targets_file (str): Path to the target file
        data_list (str): Path to a .txt file containing the input files for consideration
        prediction_output_path (str): Output prediction path
        prediction_output_database_name (str): Name of the output database
        prediction_output_statistics_name (str): Name of the output statistics database
        brain_mask_path (str): Path to the MNI brain mask file
        dmri_mean_mask_path (str): Path to the dualreg subject mean mask
        rsfmri_mean_mask_path (str): Path to the summed tract mean mask
        regression_factors (str): Path to the linear regression weights file
        mean_regression_flag (bool): Flag indicating if the volumes should be de-meaned by regression using the mean_mask_path
        mean_regression_all_flag (bool): Flag indicating if both the input and target volumes should be regressed. If False, only targets are regressed.
        mean_subtraction_flag (bool): Flag indicating if the targets should be de-meaned by subtraction using the mean_mask_path
        scale_volumes_flag (bool): Flag indicating if the volumes should be scaled.
        normalize_flag (bool): Flag signaling if the volume should be normalized ([0,1] if True) or scaled to [-1,1] if False.
        minus_one_scaling_flag (bool): Flag signaling if the volume should be scaled to [-1,1] if True
        negative_flag (bool): Flag indicating if all the negative values should be 0-ed. 
        outlier_flag (bool): Flag indicating if outliers should be set to the min/max values.
        shrinkage_flag (bool): Flag indicating if shrinkage should be applied.
        hard_shrinkage_flag (bool): Flag indicating if hard shrinkage should be applied. If False, soft shrinkage is applied. 
        crop_flag (bool): Flag indicating if the volumes should be cropped from 91x109x91 to 72x90x77 to reduce storage space and speed-up training
        device (str/int): Device type used for training (int - GPU id, str- CPU)
        mode (str): Current run mode or phase
        exit_on_error (bool): Flag that triggers the raising of an exception
        output_database_flag (bool): Flag indicating if the output maps should be saved to hdf5 database
        cross_domain_x2x_flag (bool): Flag indicating if cross-domain training is occuring between the inputs
        cross_domain_y2y_flag (bool): Flag indicating if cross-domain training is occuring between the targets

    Raises:
        FileNotFoundError: Error in reading the provided file!
        Exception: Error code execution!
    """

    log.info(
        "Started Evaluation. Check tensorboard for plots (if a LogWriter is provided)")

    with open(data_list) as data_list_file:
        volumes_to_be_used = data_list_file.read().splitlines()

    # Test if cuda is available and attempt to run on GPU

    cuda_available = torch.cuda.is_available()
    if type(device) == int:
        if cuda_available:
            model = torch.load(trained_model_path)
            torch.cuda.empty_cache()
            model.cuda(device)
        else:
            log.warning(
                "CUDA not available. Switching to CPU. Investigate behaviour!")
            device = 'cpu'

    if (type(device) == str) or not cuda_available:
        model = torch.load(trained_model_path,
                           map_location=torch.device(device))

    model.eval()

    # Create the prediction path folder if this is not available

    create_folder(prediction_output_path)

    # Initiate the evaluation

    log.info("rsfMRI Generation Started")
    if cross_domain_y2y_flag == True:
        # If doing y2y autoencoder, then we load the targets as inputs. In all other cases (x2x & x2y) we load the inputs as inputs.
        file_paths, volumes_to_be_used = data_utils.load_file_paths(data_directory, data_list, mapping_data_file=mapping_targets_file, mapping_targets_file=mapping_targets_file)
    elif cross_domain_x2x_flag == True:
        file_paths, volumes_to_be_used = data_utils.load_file_paths(data_directory, data_list, mapping_data_file, mapping_targets_file=mapping_data_file)
    else:
        file_paths, volumes_to_be_used = data_utils.load_file_paths(data_directory, data_list, mapping_data_file, mapping_targets_file)

    if output_database_flag == True:
        output_database_path = os.path.join(prediction_output_path, prediction_output_database_name)
        if os.path.exists(output_database_path):
            os.remove(output_database_path)
        output_database_handle = h5py.File(output_database_path, 'w')

    output_statistics = {}
    output_statistics_path = os.path.join(prediction_output_path, prediction_output_statistics_name)

    with torch.no_grad():

        for volume_index, file_path in enumerate(file_paths):
            try:
                print("Mapping Volume {}/{}".format(volume_index+1, len(file_paths)))
                # Generate volume & header

                subject = volumes_to_be_used[volume_index]

                predicted_complete_volume, predicted_volume, header, xform = _generate_volume_map(file_path,
                                                                                                subject,
                                                                                                model,
                                                                                                device,
                                                                                                cuda_available,
                                                                                                brain_mask_path,
                                                                                                dmri_mean_mask_path,
                                                                                                rsfmri_mean_mask_path,
                                                                                                regression_factors,
                                                                                                mean_regression_flag,
                                                                                                mean_regression_all_flag, 
                                                                                                mean_subtraction_flag,
                                                                                                scale_volumes_flag,
                                                                                                normalize_flag,
                                                                                                minus_one_scaling_flag,
                                                                                                negative_flag, 
                                                                                                outlier_flag,
                                                                                                shrinkage_flag,
                                                                                                hard_shrinkage_flag,
                                                                                                crop_flag,
                                                                                                cross_domain_x2x_flag,
                                                                                                cross_domain_y2y_flag)

                if crop_flag == True:
                    predicted_volume = roi(Image(predicted_volume, header=header), ((-9,82),(-10,99),(0,91))).data

                target_volume = _generate_target_volume(file_path,
                                                        subject,
                                                        dmri_mean_mask_path,
                                                        rsfmri_mean_mask_path,
                                                        regression_factors,
                                                        mean_regression_flag,
                                                        mean_regression_all_flag, 
                                                        mean_subtraction_flag,
                                                        crop_flag, 
                                                        cross_domain_x2x_flag
                                                        )

                mse, mae, cel, pearson_r, pearson_p, spearman_r, spearman_p, regression_w, regression_b = _statistics_calculator(predicted_volume, target_volume)
                output_statistics[subject] = [mse, mae, cel, pearson_r, pearson_p, spearman_r, spearman_p, regression_w, regression_b]

                if output_database_flag == True:
                    group = output_database_handle.create_group(subject)
                    group.create_dataset('predicted_complete_volume', data=predicted_complete_volume)
                    group.create_dataset('predicted_volume', data=predicted_volume)
                    group.create_dataset('header', data=header)
                    group.create_dataset('xform', data=xform)

                log.info("Processed: " + volumes_to_be_used[volume_index] + " " + str(
                    volume_index + 1) + " out of " + str(len(volumes_to_be_used)))

                print("Mapped Volumes saved in: ", prediction_output_path)

            except FileNotFoundError as exception_expression:
                log.error("Error in reading the provided file!")
                log.exception(exception_expression)
                if exit_on_error:
                    raise(exception_expression)

            except Exception as exception_expression:
                log.error("Error code execution!")
                log.exception(exception_expression)
                if exit_on_error:
                    raise(exception_expression)

        output_statistics_df = pd.DataFrame.from_dict(output_statistics, orient='index', columns=['mse', 'mae', 'cel', 'pearson_r', 'pearson_p', 'spearman_r', 'spearman_p', 'regression_w', 'regression_b'])     
        output_statistics_df.to_csv(output_statistics_path)

    log.info("Output Data Generation Complete")

    if output_database_flag == True:
        output_database_handle.close()


def evaluate_mapping(trained_model_path,
                     data_directory,
                     mapping_data_file,
                     mapping_targets_file,
                     data_list,
                     prediction_output_path,
                     brain_mask_path,
                     dmri_mean_mask_path,
                     rsfmri_mean_mask_path,
                     regression_factors,
                     mean_regression_flag,
                     mean_regression_all_flag, 
                     mean_subtraction_flag,
                     scale_volumes_flag,
                     normalize_flag,
                     minus_one_scaling_flag,
                     negative_flag, 
                     outlier_flag,
                     shrinkage_flag,
                     hard_shrinkage_flag,
                     crop_flag,
                     device=0, 
                     exit_on_error=False,
                     cross_domain_x2x_flag=False,
                     cross_domain_y2y_flag=False,
                     mode='evaluate'):
    """Model Evaluator

    This function generates the rsfMRI map for an input running on on a single axis or path

    Args:
        trained_model_path (str): Path to the location of the trained model
        data_directory (str): Path to input data directory
        mapping_data_file (str): Path to the input file
        mapping_targets_file (str): Path to the target file
        data_list (str): Path to a .txt file containing the input files for consideration
        prediction_output_path (str): Output prediction path
        brain_mask_path (str): Path to the MNI brain mask file
        dmri_mean_mask_path (str): Path to the dualreg subject mean mask
        rsfmri_mean_mask_path (str): Path to the summed tract mean mask
        regression_factors (str): Path to the linear regression weights file
        mean_regression_flag (bool): Flag indicating if the volumes should be de-meaned by regression using the mean_mask_path
        mean_regression_all_flag (bool): Flag indicating if both the input and target volumes should be regressed. If False, only targets are regressed.
        mean_subtraction_flag (bool): Flag indicating if the targets should be de-meaned by subtraction using the mean_mask_path
        scale_volumes_flag (bool): Flag indicating if the volumes should be scaled.
        normalize_flag (bool): Flag signaling if the volume should be normalized ([0,1] if True) or scaled to [-1,1] if False.
        minus_one_scaling_flag (bool): Flag signaling if the volume should be scaled to [-1,1] if True
        negative_flag (bool): Flag indicating if all the negative values should be 0-ed. 
        outlier_flag (bool): Flag indicating if outliers should be set to the min/max values.
        shrinkage_flag (bool): Flag indicating if shrinkage should be applied.
        hard_shrinkage_flag (bool): Flag indicating if hard shrinkage should be applied. If False, soft shrinkage is applied. 
        crop_flag (bool): Flag indicating if the volumes should be cropped from 91x109x91 to 72x90x77 to reduce storage space and speed-up training
        device (str/int): Device type used for training (int - GPU id, str- CPU)
        mode (str): Current run mode or phase
        exit_on_error (bool): Flag that triggers the raising of an exception
        cross_domain_x2x_flag (bool): Flag indicating if cross-domain training is occuring between the inputs
        cross_domain_y2y_flag (bool): Flag indicating if cross-domain training is occuring between the targets

    Raises:
        FileNotFoundError: Error in reading the provided file!
        Exception: Error code execution!
    """

    log.info(
        "Started Evaluation. Check tensorboard for plots (if a LogWriter is provided)")

    with open(data_list) as data_list_file:
        volumes_to_be_used = data_list_file.read().splitlines()

    # Test if cuda is available and attempt to run on GPU

    cuda_available = torch.cuda.is_available()
    if type(device) == int:
        if cuda_available:
            model = torch.load(trained_model_path)
            torch.cuda.empty_cache()
            model.cuda(device)
        else:
            log.warning(
                "CUDA not available. Switching to CPU. Investigate behaviour!")
            device = 'cpu'

    if (type(device) == str) or not cuda_available:
        model = torch.load(trained_model_path,
                           map_location=torch.device(device))

    model.eval()

    # Create the prediction path folder if this is not available

    create_folder(prediction_output_path)

    # Initiate the evaluation

    log.info("rsfMRI Generation Started")
    if cross_domain_y2y_flag == True:
        # If doing y2y autoencoder, then we load the targets as inputs. In all other cases (x2x & x2y) we load the inputs as inputs.
        file_paths, volumes_to_be_used = data_utils.load_file_paths(data_directory, data_list, mapping_data_file=mapping_targets_file)
    else:
        file_paths, volumes_to_be_used = data_utils.load_file_paths(data_directory, data_list, mapping_data_file)

    with torch.no_grad():

        for volume_index, file_path in enumerate(file_paths):
            try:
                print("Mapping Volume {}/{}".format(volume_index+1, len(file_paths)))
                # Generate volume & header

                subject = volumes_to_be_used[volume_index]

                predicted_complete_volume, predicted_volume, header, xform = _generate_volume_map(file_path,
                                                                                                subject,
                                                                                                model,
                                                                                                device,
                                                                                                cuda_available,
                                                                                                brain_mask_path,
                                                                                                dmri_mean_mask_path,
                                                                                                rsfmri_mean_mask_path,
                                                                                                regression_factors,
                                                                                                mean_regression_flag,
                                                                                                mean_regression_all_flag, 
                                                                                                mean_subtraction_flag,
                                                                                                scale_volumes_flag,
                                                                                                normalize_flag,
                                                                                                minus_one_scaling_flag,
                                                                                                negative_flag, 
                                                                                                outlier_flag,
                                                                                                shrinkage_flag,
                                                                                                hard_shrinkage_flag,
                                                                                                crop_flag,
                                                                                                cross_domain_x2x_flag,
                                                                                                cross_domain_y2y_flag)

                if crop_flag == False:
                    if cross_domain_y2y_flag == True:
                        output_nifti_image = Image(predicted_volume, header=header)
                    else:
                        output_nifti_image = Image(predicted_volume, header=header, xform=xform)
                elif crop_flag == True:
                    output_nifti_image = Image(predicted_volume, header=header)
                    output_nifti_image = roi(output_nifti_image, ((-9,82),(-10,99),(0,91)))

                output_nifti_path = os.path.join(
                    prediction_output_path, volumes_to_be_used[volume_index])

                if '.nii' not in output_nifti_path:
                    output_nifti_path += '.nii.gz'

                output_nifti_image.save(output_nifti_path)

                if mean_regression_flag == True:
                    if crop_flag == False:
                        if cross_domain_y2y_flag == True:
                            output_nifti_image = Image(predicted_complete_volume, header=header)
                        else:
                            output_complete_nifti_image = Image(predicted_complete_volume, header=header, xform=xform)
                    elif crop_flag == True:
                        output_complete_nifti_image = Image(predicted_complete_volume, header=header)
                        output_complete_nifti_image = roi(output_complete_nifti_image, ((-9,82),(-10,99),(0,91)))

                    output_complete_nifti_path = os.path.join(
                        prediction_output_path, volumes_to_be_used[volume_index]) + '_complete'

                    if '.nii' not in output_complete_nifti_path:
                        output_complete_nifti_path += '.nii.gz'

                    output_complete_nifti_image.save(
                        output_complete_nifti_path)

                log.info("Processed: " + volumes_to_be_used[volume_index] + " " + str(
                    volume_index + 1) + " out of " + str(len(volumes_to_be_used)))

                print("Mapped Volumes saved in: ", prediction_output_path)

            except FileNotFoundError as exception_expression:
                log.error("Error in reading the provided file!")
                log.exception(exception_expression)
                if exit_on_error:
                    raise(exception_expression)

            except Exception as exception_expression:
                log.error("Error code execution!")
                log.exception(exception_expression)
                if exit_on_error:
                    raise(exception_expression)

    log.info("rsfMRI Generation Complete")


def _generate_volume_map(file_path,
                         subject,
                         model,
                         device,
                         cuda_available,
                         brain_mask_path,
                         dmri_mean_mask_path,
                         rsfmri_mean_mask_path,
                         regression_factors,
                         mean_regression_flag,
                         mean_regression_all_flag, 
                         mean_subtraction_flag,
                         scale_volumes_flag,
                         normalize_flag,
                         minus_one_scaling_flag,
                         negative_flag, 
                         outlier_flag,
                         shrinkage_flag,
                         hard_shrinkage_flag,
                         crop_flag,
                         cross_domain_x2x_flag,
                         cross_domain_y2y_flag
                         ):
    """Output Volume Generator

    This function uses the trained model to generate a new volume

    Args:
        file_path (str): Path to the desired file
        subject (str): Subject ID of the subject volume to be regressed
        model (class): BrainMapper model class
        device (str/int): Device type used for training (int - GPU id, str- CPU)
        cuda_available (bool): Flag indicating if a cuda-enabled GPU is present
        brain_mask_path (str): Path to the MNI brain mask file
        dmri_mean_mask_path (str): Path to the group mean volume
        rsfmri_mean_mask_path (str): Path to the dualreg subject mean mask
        regression_factors (str): Path to the linear regression weights file
        mean_regression_flag (bool): Flag indicating if the volumes should be de-meaned by regression using the mean_mask_path
        mean_regression_all_flag (bool): Flag indicating if both the input and target volumes should be regressed. If False, only targets are regressed.
        mean_subtraction_flag (bool): Flag indicating if the targets should be de-meaned by subtraction using the mean_mask_path
        scale_volumes_flag (bool): Flag indicating if the volumes should be scaled.
        normalize_flag (bool): Flag signaling if the volume should be normalized ([0,1] if True) or scaled to [-1,1] if False.
        minus_one_scaling_flag (bool): Flag signaling if the volume should be scaled to [-1,1] if True
        negative_flag (bool): Flag indicating if all the negative values should be 0-ed. 
        outlier_flag (bool): Flag indicating if outliers should be set to the min/max values.
        shrinkage_flag (bool): Flag indicating if shrinkage should be applied.
        hard_shrinkage_flag (bool): Flag indicating if hard shrinkage should be applied. If False, soft shrinkage is applied. 
        crop_flag (bool): Flag indicating if the volumes should be cropped from 91x109x91 to 72x90x77 to reduce storage space and speed-up training
        cross_domain_x2x_flag (bool): Flag indicating if cross-domain training is occuring between the inputs
        cross_domain_y2y_flag (bool): Flag indicating if cross-domain training is occuring between the targets

    Returns
        predicted_volume (np.array): Array containing the information regarding the generated volume
        header (class): 'nibabel.nifti1.Nifti1Header' class object, containing volume metadata
    """

    volume, header, xform = data_utils.load_and_preprocess_evaluation(file_path, crop_flag, cross_domain_y2y_flag)

    if mean_regression_flag == True:
        if mean_regression_all_flag == True:
            volume = _regress_input(volume, subject, dmri_mean_mask_path, rsfmri_mean_mask_path, regression_factors, crop_flag, cross_domain_y2y_flag)
            scaling_parameters = [-0.0626, 0.1146, -14.18, 16.9475]
        else:
            scaling_parameters = [0.0, 0.2, -14.18, 16.9475]
    elif mean_subtraction_flag == True:
        scaling_parameters = [0.0, 0.2, 0.0, 10.0]

    print('volume range:', np.min(volume), np.max(volume))

    if scale_volumes_flag == True:
        volume = _scale_input(volume, scaling_parameters, normalize_flag, minus_one_scaling_flag, negative_flag, outlier_flag, shrinkage_flag, hard_shrinkage_flag, cross_domain_y2y_flag)

    if len(volume.shape) == 5:
        volume = volume
    else:
        volume = volume[np.newaxis, np.newaxis, :, :, :]

    volume = torch.tensor(volume).type(torch.FloatTensor)

    if cuda_available and (type(device) == int):
        volume = volume.cuda(device)

    output = model(volume)
    output = (output.cpu().numpy()).astype('float32')
    output = np.squeeze(output)

    print('output range:', np.min(output), np.max(output))

    output = _rescale_output(output, scaling_parameters, normalize_flag, minus_one_scaling_flag, negative_flag, shrinkage_flag, hard_shrinkage_flag, cross_domain_x2x_flag)

    print('output rescaled:', np.min(output), np.max(output))

    if crop_flag == False:
        MNI152_T1_2mm_brain_mask = Image(brain_mask_path).data
    elif crop_flag == True:
        MNI152_T1_2mm_brain_mask = roi(Image(brain_mask_path),((9,81),(10,100),(0,77))).data

    if mean_regression_flag == True or mean_subtraction_flag == True:

        if cross_domain_x2x_flag == True:
            if crop_flag == False:
                mean_mask = Image(dmri_mean_mask_path).data
            elif crop_flag == True:
                mean_mask = roi(Image(dmri_mean_mask_path),((9,81),(10,100),(0,77))).data
        else:
            if crop_flag == False:
                mean_mask = Image(rsfmri_mean_mask_path).data[:, :, :, 0]
            elif crop_flag == True:
                mean_mask = roi(Image(rsfmri_mean_mask_path),((9,81),(10,100),(0,77))).data[:, :, :, 0]

        if mean_regression_flag == True:
            if cross_domain_x2x_flag == True:
                weight = pd.read_pickle(regression_factors).loc[subject]['w_dMRI']
            else:
                weight = pd.read_pickle(regression_factors).loc[subject]['w_rsfMRI']
            predicted_complete_volume = np.add(output, np.multiply(weight, mean_mask))

        if mean_subtraction_flag == True:
            predicted_complete_volume = np.add(output, mean_mask)

        print('predicted_complete_volume', np.min(
            predicted_complete_volume), np.max(predicted_complete_volume))

        predicted_complete_volume = np.multiply(
            predicted_complete_volume, MNI152_T1_2mm_brain_mask)

        print('predicted_complete_volume masked:', np.min(
            predicted_complete_volume), np.max(predicted_complete_volume))

    else:
        predicted_complete_volume = None

    predicted_volume = np.multiply(output, MNI152_T1_2mm_brain_mask)

    print('predicted_volume masked:', np.min(
        predicted_volume), np.max(predicted_volume))

    return predicted_complete_volume, predicted_volume, header, xform


def _scale_input(volume, scaling_parameters, normalize_flag, minus_one_scaling_flag, negative_flag, outlier_flag, shrinkage_flag, hard_shrinkage_flag, cross_domain_y2y_flag):
    """Input Scaling

    This function reads the scaling factors from the saved file and then scales the data.

    Args:
        volume (np.array): Numpy array representing the un-scalled volume. 
        scaling_parameters (list): List of scaling parameters.
        normalize_flag (bool): Flag signaling if the volume should be normalized ([0,1] if True) or scaled to [-1,1] if False.
        minus_one_scaling_flag (bool): Flag signaling if the volume should be scaled to [-1,1] if True
        negative_flag (bool): Flag indicating if all the negative values should be 0-ed. 
        outlier_flag (bool): Flag indicating if outliers should be set to the min/max values.
        shrinkage_flag (bool): Flag indicating if shrinkage should be applied.
        hard_shrinkage_flag (bool): Flag indicating if hard shrinkage should be applied. If False, soft shrinkage is applied. 
        cross_domain_y2y_flag (bool): Flag indicating if cross-domain training is occuring between the targets

    Returns:
        scaled_volume (np.array): Scaled volume
    """

    if cross_domain_y2y_flag == True:
        _, _, min_value, max_value = scaling_parameters
    else:
        min_value, max_value, _, _ = scaling_parameters

    if shrinkage_flag == True:
        if cross_domain_y2y_flag == True:
            lambd = 3.0
        else:
            lambd = 0.003 # Hard coded, equivalent to tht 1p and 99p values across the whole population in UKBB

        if hard_shrinkage_flag == True:
            volume = _hard_shrinkage(volume, lambd)
        elif hard_shrinkage_flag == False:
            volume = _soft_shrinkage(volume, lambd)
            min_value += lambd
            max_value -= lambd

    if negative_flag == True:
        volume[volume < 0.0] = 0.0
        min_value = 0.0

    if outlier_flag == True:
        volume[volume > max_value] = max_value
        volume[volume < min_value] = min_value

    if normalize_flag == True:
        # Normalization to [0, 1]
        scaled_volume = np.divide(np.subtract(volume, min_value), np.subtract(max_value, min_value))
    elif minus_one_scaling_flag == True:
        # Scaling between [-1, 1]
        scaled_volume = np.add(-1.0, np.multiply(2.0, np.divide(np.subtract(volume, min_value), np.subtract(max_value, min_value))))
    # Else, no scaling occus, but the other flags can still hold true if the scaling flag is true! 

    return scaled_volume


def _regress_input(volume, subject, dmri_mean_mask_path, rsfmri_mean_mask_path, regression_factors, crop_flag, cross_domain_y2y_flag):
    """ Inputn Regression

    This function regresse the group mean from the input volume using the saved regression weights.

    TODO: This function repressents only a temporary solution. For deployment, a NN needs to be trained which predicts the relevant scaling factors.

    Args:
        volume (np.array): Unregressed volume
        subject (str): Subject ID of the subject volume to be regressed
        dmri_mean_mask_path (str): Path to the group mean volume
        rsfmri_mean_mask_path (str): Path to the target group mean volume
        regression_factors (str): Path to the linear regression weights file
        crop_flag (bool): Flag indicating if the volumes should be cropped from 91x109x91 to 72x90x77 to reduce storage space and speed-up training
        cross_domain_y2y_flag (bool): Flag indicating if cross-domain training is occuring between the targets

    Returns:
        regressed_volume (np.array): Linear regressed volume

    """

    if cross_domain_y2y_flag == True:
        weight = pd.read_pickle(regression_factors).loc[subject]['w_rsfMRI'] 
        if crop_flag == False:
            group_mean = Image(rsfmri_mean_mask_path).data[:, :, :, 0]
        elif crop_flag == True:
            group_mean = roi(Image(rsfmri_mean_mask_path),((9,81),(10,100),(0,77))).data[:, :, :, 0]
    else:
        weight = pd.read_pickle(regression_factors).loc[subject]['w_dMRI']
        if crop_flag == False:
            group_mean = Image(dmri_mean_mask_path).data
        elif crop_flag == True:
            group_mean = roi(Image(dmri_mean_mask_path),((9,81),(10,100),(0,77))).data

    regressed_volume = np.subtract(volume, np.multiply(weight, group_mean))

    return regressed_volume


def _rescale_output(volume, scaling_parameters, normalize_flag, minus_one_scaling_flag, negative_flag, shrinkage_flag, hard_shrinkage_flag, cross_domain_x2x_flag):
    """Output Rescaling

    This function reads the scaling factors from the saved file and then scales the data.

    Args:
        volume (np.array): Unscalled volume
        scaling_parameters (list): List of scaling parameters.
        normalize_flag (bool): Flag signaling if the volume should be normalized ([0,1] if True) or scaled to [-1,1] if False.
        minus_one_scaling_flag (bool): Flag signaling if the volume should be scaled to [-1,1] if True
        negative_flag (bool): Flag indicating if all the negative values should be 0-ed. 
        shrinkage_flag (bool): Flag indicating if shrinkage should be applied.
        hard_shrinkage_flag (bool): Flag indicating if hard shrinkage should be applied. If False, soft shrinkage is applied. 
        cross_domain_x2x_flag (bool): Flag indicating if cross-domain training is occuring between the inputs

    Returns:
        rescaled_volume (np.array): Rescaled volume
    """

    if cross_domain_x2x_flag == True:
        min_value, max_value, _, _ = scaling_parameters
    else:
        _, _, min_value, max_value = scaling_parameters

    if shrinkage_flag == True:
        if cross_domain_x2x_flag == True:
            lambd = 0.003
        else:
            lambd = 3.0
        
        if hard_shrinkage_flag == True:
            pass
        elif hard_shrinkage_flag == False:
            min_value += lambd
            max_value -= lambd

    if negative_flag == True:
        min_value = 0.0

    if normalize_flag == True:
        # Normalization to [0, 1]
        rescaled_volume = np.add(np.multiply(volume, np.subtract(max_value, min_value)), min_value)
    elif minus_one_scaling_flag == True:
        # Scaling between [-1, 1]
        rescaled_volume = np.add(np.multiply(np.divide(np.add(volume, 1), 2), np.subtract(max_value, min_value)), min_value)
    # Else, no rescaling occus, but the other flags can still hold true if the scaling flag is true! 

    return rescaled_volume


def _hard_shrinkage(volume, lambd):
    """ Hard Shrinkage

    This function performs a hard shrinkage on the volumes.
    volume = { x , x > lambd | x < -lambd
                0 , x e [-lambd, lambd]
                }

    Args:
        volume (np.array): Unshrunken volume
        lambd (float): Threshold parameter
    
    Returns:
        volume (np.array) : Hard shrunk volume
    """

    volume[np.where(np.logical_and(volume>-lambd, volume<lambd))] = 0

    return volume


def _soft_shrinkage(volume, lambd):
    """ Soft Shrinkage

    This function performs a soft shrinkage on the volumes.
    volume = { x + lambd , x < -lambd
                0         , x e [-lambd, lambd]
                x - lambd , x > lambd
                }

    Args:
        volume (np.array): Unshrunken volume
        lambd (float): Threshold parameter
    
    Returns:
        volume (np.array) : Soft shrunk volume
    """

    volume[np.where(np.logical_and(volume>=-lambd, volume<=lambd))] = 0.0
    volume[volume < -lambd] = volume[volume < -lambd] + lambd
    volume[volume > lambd] = volume[volume > lambd] - lambd

    return volume


def _generate_target_volume(file_path,
                            subject,
                            dmri_mean_mask_path,
                            rsfmri_mean_mask_path,
                            regression_factors,
                            mean_regression_flag,
                            mean_regression_all_flag, 
                            mean_subtraction_flag,
                            crop_flag, 
                            cross_domain_x2x_flag
                            ):
    """Target Volume Generator

    This function loads and preprocesses a target volume for comparing with the network predicted volumes

    Args:
        file_path (str): Path to the desired file
        subject (str): Subject ID of the subject volume to be regressed
        dmri_mean_mask_path (str): Path to the group mean volume
        rsfmri_mean_mask_path (str): Path to the dualreg subject mean mask
        regression_factors (str): Path to the linear regression weights file
        mean_regression_flag (bool): Flag indicating if the volumes should be de-meaned by regression using the mean_mask_path
        mean_regression_all_flag (bool): Flag indicating if both the input and target volumes should be regressed. If False, only targets are regressed.
        mean_subtraction_flag (bool): Flag indicating if the targets should be de-meaned by subtraction using the mean_mask_path
        crop_flag (bool): Flag indicating if the volumes should be cropped from 91x109x91 to 72x90x77 to reduce storage space and speed-up training
        cross_domain_x2x_flag (bool): Flag indicating if cross-domain training is occuring between the inputs

    Returns:
        volume (np.array): Array containing the information regarding the target volume
    """

    volume = data_utils.load_and_preprocess_targets(file_path, cross_domain_x2x_flag, crop_flag=False)

    if mean_regression_flag == True:
        volume = _regress_target(volume, subject, dmri_mean_mask_path, rsfmri_mean_mask_path, regression_factors, crop_flag, cross_domain_x2x_flag, mean_regression_all_flag)
    elif mean_subtraction_flag == True:
        volume = _subtract_target(volume, subject, dmri_mean_mask_path, rsfmri_mean_mask_path, crop_flag, cross_domain_x2x_flag)

    return volume


def _regress_target(volume, subject, dmri_mean_mask_path, rsfmri_mean_mask_path, regression_factors, crop_flag, cross_domain_x2x_flag, mean_regression_all_flag):
    """ Target Regression

    This function regresse the group mean from the target volume using the saved regression weights.

    TODO: This function repressents only a temporary solution. For deployment, a NN needs to be trained which predicts the relevant scaling factors.

    Args:
        volume (np.array): Unregressed volume
        subject (str): Subject ID of the subject volume to be regressed
        dmri_mean_mask_path (str): Path to the group mean volume
        rsfmri_mean_mask_path (str): Path to the target group mean volume
        regression_factors (str): Path to the linear regression weights file
        crop_flag (bool): Flag indicating if the volumes should be cropped from 91x109x91 to 72x90x77 to reduce storage space and speed-up training
        cross_domain_x2x_flag (bool): Flag indicating if cross-domain training is occuring between the targets
        mean_regression_all_flag (bool): Flag indicating if both the input and target volumes should be regressed. If False, only targets are regressed.

    Returns:
        regressed_volume (np.array): Linear regressed volume

    """

    if cross_domain_x2x_flag == True:
        if mean_regression_all_flag == True:
            weight = pd.read_pickle(regression_factors).loc[subject]['w_dMRI']
            if crop_flag == False:
                group_mean = Image(dmri_mean_mask_path).data
            elif crop_flag == True:
                group_mean = roi(Image(dmri_mean_mask_path),((9,81),(10,100),(0,77))).data
            
            regressed_volume = np.subtract(volume, np.multiply(weight, group_mean))
        else:
            regressed_volume = volume
    else:
        weight = pd.read_pickle(regression_factors).loc[subject]['w_rsfMRI'] 
        if crop_flag == False:
            group_mean = Image(rsfmri_mean_mask_path).data[:, :, :, 0]
        elif crop_flag == True:
            group_mean = roi(Image(rsfmri_mean_mask_path),((9,81),(10,100),(0,77))).data[:, :, :, 0]

        regressed_volume = np.subtract(volume, np.multiply(weight, group_mean))

    return regressed_volume


def _subtract_target(volume, subject, dmri_mean_mask_path, rsfmri_mean_mask_path, crop_flag, cross_domain_x2x_flag):
    """ Target Subtraction

    This function subtracts the group mean from the target volume using the saved regression weights.

    TODO: This function repressents only a temporary solution. For deployment, a NN needs to be trained which predicts the relevant scaling factors.

    Args:
        volume (np.array): Unregressed volume
        subject (str): Subject ID of the subject volume to be regressed
        dmri_mean_mask_path (str): Path to the group mean volume
        rsfmri_mean_mask_path (str): Path to the target group mean volume
        crop_flag (bool): Flag indicating if the volumes should be cropped from 91x109x91 to 72x90x77 to reduce storage space and speed-up training
        cross_domain_x2x_flag (bool): Flag indicating if cross-domain training is occuring between the targets

    Returns:
        regressed_volume (np.array): Linear regressed volume

    """

    if cross_domain_x2x_flag == True:
        subtracted_volume = volume
    else:
        if crop_flag == False:
            group_mean = Image(rsfmri_mean_mask_path).data[:, :, :, 0]
        elif crop_flag == True:
            group_mean = roi(Image(rsfmri_mean_mask_path),((9,81),(10,100),(0,77))).data[:, :, :, 0]

        subtracted_volume = np.subtract(volume, group_mean)

    return subtracted_volume


def _statistics_calculator(volume, target):
    """ Training statistics calculator

    This function calculates the MSE, MAE, CEL, Pearson R and P and linear regression W and B for a predicted volume and it's target ground truth.

    Args:
        volume (np.array): Predicted volume
        target (np.array): Ground truth volume
    
    Returns:
        mse (np.float64): The mean squared error between the prediction and the ground truth; The closer to 0, the better
        mae (np.float64): The mean absolut error between the prediction and the ground truth; The closer to 0, the better
        cel (np.float64): The cosine distance between the prediction and the ground truth; The closer to 0, the better
        pearson_r (np.float64): Pearson’s correlation coefficient; The closer to 1, the better
        pearson_p (np.float64): Two-tailed p-value for Pearson’s correlation coefficient; the closer to 0, the better
        spearman_r (np.float64): Spearman correlation coefficient; The closer to 1, the better
        spearman_p (np.float64): Two-tailed p-value for Spearman's correlation coefficient; the closer to 0, the better
        regression_w (np.float64): Slope of the linear regression line; The closer to 1 the better
        regression_b (np.float64): Intersect of the linear regression line; The closer to 0, the better
    """

    x = np.reshape(volume, -1)
    y = np.reshape(target, -1)

    mse = np.square(np.subtract(x,y)).mean()
    mae = np.abs(np.subtract(x,y)).mean()
    cel = np.mean(cosine(x, y))
    pearson_r, pearson_p = pearsonr(x,y)
    spearman_r, spearman_p = spearmanr(x,y)

    x_matrix = np.vstack((np.ones(len(x)), x)).T
    regression_b, regression_w = np.linalg.inv(x_matrix.T.dot(x_matrix)).dot(x_matrix.T).dot(y)

    return mse, mae, cel, pearson_r, pearson_p, spearman_r, spearman_p, regression_w, regression_b


def _pearson_correlation(volume, target):
    """Calculate Pearson Correlation Coefficient

    This function calculates the pearson correlation coefficient between a predicted volume and the target volume

    Args:
        volume (np.array): The predicted volume
        target (np.array): The target volume

    Returns:
        r (np.float32): The Pearson Correlation Coefficient
    """

    r = np.sum(np.multiply(np.subtract(volume, volume.mean()), np.subtract(target, target.mean()))) / np.sqrt(np.multiply(
        np.sum(np.power(np.subtract(volume, volume.mean()), 2)), np.sum(np.power(np.subtract(target, target.mean()), 2))))

    return r