"""Group Mean Calculator

Description:

    This file contains the relevant scripts for calculating the group mean for the subjects present in the UK Biobank.
    This is a standalone scrip, intended to be used only once during the project. Hence, it is not integrated into the larger utils packages.

Usage:

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

        from utils.group_mean import function_name

"""

import os
import numpy as np
from fsl.data.image import Image
from fsl.utils.image.resample import resampleToPixdims
from utils.preprocessor import directory_reader


def mean_calculator(data_directory, mean_type='dMRI'):
    """ Population mean calculator

    This function calculates the mean across all subjects for a given data type.

    Args:
        data_directory (str): A string containing the address of the required directory.
        mean_type (str): String flag indicating which data type the algorithm should compute the mean for.

    """
    subDirectoryList, number_of_subjects = directory_reader(
        os.path.join(os.path.expanduser("~"), data_directory))

    subject_number = len(os.listdir(os.path.join(
        os.path.expanduser("~"), data_directory)))

    volume_sum = None
    counter = None
    header = None

    for directory in subDirectoryList:
        if mean_type == 'dMRI':
            data_path = "dMRI/autoptx_preproc/tractsNormSummed.nii.gz"
        elif mean_type == 'fMRI':
            data_path = "fMRI/rfMRI_25.dr/dr_stage2.nii.gz"
        else:
            raise ValueError('mean_type can be either dMRI or fMRI!')

        subject_path = os.path.join(os.path.expanduser(
            "~"), data_directory, directory, data_path)
        if mean_type == 'dMRI':
            volume = Image(subject_path).data
        else:
            volume = Image(subject_path).data[:, :, :, 0]

        if volume_sum is None:
            volume_sum = np.zeros(volume.shape)

        if counter is None:
            counter = 0

        if header is None:
            header = Image(subject_path).header

        volume_sum = np.add(volume_sum, volume)
        counter += 1

        print("Added volume {}/{}, --{}%".format(counter,
                                                 subject_number, counter/subject_number*100))

    volume_mean = np.divide(volume_sum, counter)

    volume_mean_image = Image(volume_mean, header=header)

    if mean_type == 'dMRI':
        volume_mean_image.save('utils/mean_tractsNormSummed.nii.gz')
        volume_mean_donwsampled, xform = resampleToPixdims(
            Image('utils/mean_tractsNormSummed.nii.gz'), (2, 2, 2))
        header_downsampled = Image(volume_mean_donwsampled, header=Image(
            'utils/mean_tractsNormSummed.nii.gz').header, xform=xform).header
        volume_mean_downsampled_image = Image(
            volume_mean_donwsampled, header=header_downsampled)
        volume_mean_downsampled_image.save(
            'utils/mean_tractsNormSummed_downsampled.nii.gz')
    else:
        volume_mean_image.save('utils/mean_dr_stage2.nii.gz')


if __name__ == '__main__':
    data_directory = "/well/win-biobank/projects/imaging/data/data3/subjectsAll/"
    mean_type = 'dMRI'
    # mean_type = 'fMRI'

    mean_calculator(data_directory, mean_type)