from pre_processing import pre_process
from co_registration import co_register
from template_creation import create_template_from_images
import helper_functions as hf
import sys
import os
from shutil import copyfile


def pre_process_brain_tiffs(tif_path, result_path, mouse_id):
    """
    Utility function pre-processing all tiffs in folder, outputting pre-processed nifties in result_path.
    :param tif_path: The tiffs directory.
    :param result_path: The output directory.
    :param mouse_id: The mouse id.
    :return: Returns nothing.
    """
    # Get all images and the enclosing tiff folders.
    images, folders = hf.images_in_folders(tif_path)

    # Ensure that the number of images matches the naming of the folders
    if len(images) == int(folders[-1][-2:]):
        print('Number of images matches the naming on the folders')
    else:
        sys.exit('Mismatch between number of images and naming of last TIF folder.')

    # Make output directory.
    hf.make_dirs([result_path])

    # Loop through each image and pre-process.
    for image_nr, image_path in enumerate(images, 1):
        image_nr_beauty = str(image_nr).zfill(2)
        output = result_path + mouse_id + '_S' + image_nr_beauty + '.nii'
        if not os.path.exists(output):
            pre_process(image_path, output)


def main(preprocess=False, coreg=False, make_template=False):
    """
    Main function performing pre-processing of tiffs, 3-dimensional reconstruction to reference and population-based
    template creation.
    :param preprocess: Boolean [Default: False]. Will perform pre-processing of tiffs if set to true.
    :param coreg: Boolean [Default: False]. Will perform brain reconstruction based on the sequential slices if set to
    true.
    :param make_template: Boolean [Default: False]. Will perform population-based template creation based on the
    reconstructed brain volumes if set to true.
    :return: Returns nothing.
    """
    # Working directory
    sandbox = '/run/media/frederik/S_T1/frederik_filip_mqr375/auto-seg/' # Should be removed for final submission.

    # Get list of mice ids.
    sandbox_mice = sandbox + 'mice/'
    mice = os.listdir(sandbox_mice)

    if coreg or preprocess:
        # Loop through mice ids.
        for mouse in sorted(mice):
            # Current mouse directory and subdirectories.
            work_dir = sandbox_mice + mouse + '/'
            raw_dir = work_dir + 'TIF/'
            pre_processed_dir = work_dir + 'preprocessed/'
            out_dir = work_dir + 'volumes/'

            # Make volumes directory.
            hf.make_dirs(out_dir)

            # Define raw volume and initial transform path (if present).
            initial_vol_trans = out_dir + 'initial_transform_' + mouse + '.txt'
            raw_vol_path = out_dir + '00_hv_raw_' + mouse + '.nii'

            # Pre process tiff images.
            if preprocess:
                pre_process_brain_tiffs(raw_dir, pre_processed_dir, mouse)

            # Create raw volume if not already present.
            if not os.path.exists(raw_vol_path):
                hf.make_volume(hf.files_in_dir(pre_processed_dir, '.nii'), raw_vol_path)

            # Check that an initial transform is present.
            if not os.path.exists(initial_vol_trans):
                sys.exit('An initial moving transform was not found for at least brain with id : ' + mouse +
                         '\nMake sure it is named and placed exactly: ' + initial_vol_trans)

            # -- This call should be manually modified depending on first or second iteration -- #
            # Perform co-registration to reference (Allen in first iteration. DAPI template in second iteration).
            # Manually created initial transform is only used in the first iteration.
            # First iteration had 0 non-linear steps (parameter nl=0).
            if coreg:
                mouse_coreg = co_register(sandbox + 'second-iteration_average_allen_0_15.nii.gz',
                                          pre_processed_dir, raw_vol_path, '', out_dir, mouse,
                                          print_log=True, nl=2)
                copyfile(mouse_coreg, sandbox + 'volumes/' + mouse + '.nii')

    # Population based template creation using symmetric modelling.
    if make_template:
        create_template_from_images(
            sandbox + 'volumes/',
            sandbox + 'template/',
            symmetric=True
        )


if __name__ == '__main__':
    main(preprocess=True, coreg=True, make_template=True)