import os
from nipype.interfaces import utility as niu
import nipype.pipeline.engine as pe  # pypeline engine
from nipype.interfaces.utility import Function
from nipype.interfaces.base import (
    traits, TraitedSpec, BaseInterfaceInputSpec,
    File, BaseInterface
)
from nipype.interfaces.io import DataSink
from rabies.utils import run_command, flatten_list, ResampleMask
from .registration import run_antsRegistration
from .preprocess_visual_QC import PlotOverlap,template_masking

def init_commonspace_reg_wf(opts, commonspace_reg_opts, inherit_unbiased, output_folder, transforms_datasink, nthreads_modelbuild, output_datasinks, joinsource_list, name='commonspace_reg_wf'):
    # commonspace_wf_head_start
    """
    This workflow handles the alignment of all MRI sessions to a common space. This is conducted first by generating
    a dataset-specific unbiased template from the input structural images, thereby aligning the different MRI 
    sessions. Through a set of iterations, images are registered to a consensus average generated from the overlap 
    of all scans at the previous iteration. Registrations are increasingly stringent, executing 2 iterations each 
    for a rigid, then affine and finally non-linear template generation. The final iteration provides the individual 
    transforms to align each scan to the unbiased template. This template generation process creates a robust target 
    for the alignment of MRI sessions sharing the same acquisition properties, which will minimize registration 
    inconsistencies between sessions, as opposed to the direct registration to an external template. The algorithm is 
    implemented in https://github.com/CoBrALab/optimized_antsMultivariateTemplateConstruction.
    After generating the unbiased template, the template itself is registered with a non-linear registration to the
    reference atlas in common space, providing transforms to common space and the associated brain parcellations.


    References:
        Avants, B. B., Tustison, N. J., Song, G., Cook, P. A., Klein, A., & Gee, J. C. (2011). A reproducible evaluation 
        of ANTs similarity metric performance in brain image registration. NeuroImage, 54(3), 2033–2044.

    Command line interface parameters:
        --commonspace_reg COMMONSPACE_REG
                                Specify registration options for the commonspace registration.
                                * masking: Combine masks derived from the inhomogeneity correction step to support 
                                registration during the generation of the unbiased template, and then during template 
                                registration.
                                *** Specify 'true' or 'false'. 
                                * brain_extraction: conducts brain extraction prior to template registration based on the 
                                combined masks from inhomogeneity correction. This will enhance brain edge-matching, but 
                                requires good quality masks. This should be selected along the 'masking' option.
                                *** Specify 'true' or 'false'. 
                                * keep_mask_after_extract: If using brain_extraction, use the mask to compute the registration metric
                                within the mask only. Choose to prevent stretching of the images beyond the limit of the brain mask
                                (e.g. if the moving and target images don't have the same brain coverage).
                                *** Specify 'true' or 'false'.
                                * template_registration: Specify a registration script for the alignment of the 
                                dataset-generated unbiased template to the commonspace atlas.
                                *** Rigid: conducts only rigid registration.
                                *** Affine: conducts Rigid then Affine registration.
                                *** SyN: conducts Rigid, Affine then non-linear registration.
                                *** no_reg: skip registration.
                                * fast_commonspace: Skip the generation of a dataset-generated unbiased template, and 
                                instead, register each scan independently directly onto the commonspace atlas, using the 
                                template_registration. This option can be faster, but may decrease the quality of 
                                alignment between subjects. 
                                *** Specify 'true' or 'false'. 
                                (default: masking=false,brain_extraction=false,keep_mask_after_extract=false,template_registration=SyN,fast_commonspace=false)
                        
    Workflow:
        parameters
            opts: command line interface parameters (technically all the parser args, including winsorization and brain_extraction etc are included here)
            commonspace_reg_opts: the opts that are specific to this run of commonspace_reg (there can be different parameters, for instance for the robust_inho_cor workflows)
            output_folder: specify a folder to execute the workflow and store important outputs
            transforms_datasink: datasink node where the transforms are stored
            nthreads_modelbuild: set the maximum number of parallel threads to launch
            output_datasinks: whether to generate a datasink from the outputs of the workflow
            joinsource_list: names for the iterable nodes to join before unbiased template generation

        inputs
            moving_image_list: list of files corresponding to the images from different MRI sessions
            moving_mask_list: mask files overlapping with the moving images, inherited from the inhomogeneity 
                correction step. These masks are used for --commonspace_masking and --brain_extraction
            template_anat: the target structural template to register the unbiased template
            template_mask: the brain mask of the structural template

        outputs
            unbiased_template: the generated unbiased template
            unbiased_mask: brain mask resampled over the unbiased template
            anatspace_mask: the atlas brain mask resampled to an associated MRI session in anat space
            to_atlas_affine: affine transform for registration to the atlas
            to_atlas_warp: non-linear transform for registration to the atlas
            to_atlas_inverse_warp: inverse of the non-linear transform for registration to the atlas
            anat_to_unbiased_affine: affine transform from anat space to the unbiased template
            anat_to_unbiased_warp: non-linear transform from anat space to the unbiased template
            anat_to_unbiased_inverse_warp: inverse of the non-linear transform from anat space to the unbiased template
            anat_to_commonspace_transform_list: ordered list of the transforms to apply to move from the
                anat space to the common space
            anat_to_commonspace_inverse_list: list defining whether the inverse of affine transforms should
                be applied for anat_to_commonspace_transform_list
            commonspace_to_anat_transform_list: ordered list of the transforms to apply to move from the
                common space to the anat space
            commonspace_to_anat_inverse_list: list defining whether the inverse of affine transforms should
                be applied for commonspace_to_anat_transform_list
    """
    # commonspace_wf_head_end

    # this iterable node inherits the iterations from the main_wf, generating an iteration for each session to recover its outputs from commonspace registration
    inputnode = pe.Node(niu.IdentityInterface(fields=['moving_image', 'moving_mask']),
                                        name="inputnode")

    template_inputnode = pe.Node(niu.IdentityInterface(fields=['template_anat', 'template_mask']),
                                        name="template_inputnode")
    outputnode = pe.Node(niu.IdentityInterface(fields=['unbiased_template', 'unbiased_mask', 'anatspace_mask', 'to_atlas_affine', 'to_atlas_warp', 'to_atlas_inverse_warp',
                                                       'anat_to_unbiased_affine', 'anat_to_unbiased_warp', 'anat_to_unbiased_inverse_warp',
                                                       'anat_to_commonspace_transform_list','anat_to_commonspace_inverse_list',
                                                       'commonspace_to_anat_transform_list','commonspace_to_anat_inverse_list']),
                                        name="outputnode")
    workflow = pe.Workflow(name=name)

    modelbuild_stages = commonspace_reg_opts['stages']
    commonspace_masking = commonspace_reg_opts['masking']
    brain_extraction = commonspace_reg_opts['brain_extraction']
    keep_mask_after_extract = commonspace_reg_opts['keep_mask_after_extract']
    template_reg = commonspace_reg_opts['template_registration']
    fast_commonspace = commonspace_reg_opts['fast_commonspace']
    winsorize_lower_bound = commonspace_reg_opts['winsorize_lower_bound']
    winsorize_upper_bound = commonspace_reg_opts['winsorize_upper_bound']

    if fast_commonspace or inherit_unbiased:
        # without modelbuild, the inputs iterables are not merged
        source_join_common_reg = pe.Node(niu.IdentityInterface(fields=['file_list0', 'file_list1']),
                                            name="fast_commonreg_buffer")
        merged_join_common_reg = source_join_common_reg
    else:
        workflow, source_join_common_reg, merged_join_common_reg = join_iterables(workflow=workflow, joinsource_list=joinsource_list, node_prefix='commonspace_reg', num_inputs=2)

    if inherit_unbiased:
        # Identity node to relate inherited files to existing workflow
        atlas_reg = pe.Node(niu.IdentityInterface(fields=['moving_image', 'moving_mask', 'fixed_image', 'fixed_mask', 'affine', 'warp',
                                                    'inverse_warp', 'warped_image']),
                                            name="atlas_reg_inherited")
    else:
        atlas_reg = pe.Node(Function(input_names=['reg_method', 'brain_extraction', 'keep_mask_after_extract', 'moving_image', 'moving_mask', 'fixed_image', 'fixed_mask', 'winsorize_lower_bound', 'winsorize_upper_bound','rabies_data_type'],
                                        output_names=['affine', 'warp',
                                                    'inverse_warp', 'warped_image'],
                                        function=run_antsRegistration),
                            name='atlas_reg', mem_gb=2*opts.scale_min_memory, n_procs=int(os.environ['RABIES_ITK_NUM_THREADS']))

        # don't use brain extraction without a moving mask
        if brain_extraction:
            if not commonspace_masking:
                brain_extraction=False

        atlas_reg.inputs.winsorize_lower_bound = winsorize_lower_bound
        atlas_reg.inputs.winsorize_upper_bound = winsorize_upper_bound
        atlas_reg.inputs.brain_extraction = brain_extraction
        atlas_reg.inputs.keep_mask_after_extract = keep_mask_after_extract
        atlas_reg.inputs.rabies_data_type = opts.data_type
        atlas_reg.plugin_args = {
            'qsub_args': f'-pe smp {str(3*opts.min_proc)}', 'overwrite': True}
        atlas_reg.inputs.reg_method = template_reg

    prep_commonspace_transform_node = pe.Node(Function(input_names=['anatspace_ref', 'atlas_mask', 'anat_to_unbiased_affine',
                                                               'anat_to_unbiased_warp','anat_to_unbiased_inverse_warp',
                                                               'to_atlas_affine','to_atlas_warp','to_atlas_inverse_warp',
                                                               'fast_commonspace'],
                                               output_names=[
                                                   'anatspace_mask', 'anat_to_commonspace_transform_list','anat_to_commonspace_inverse_list','commonspace_to_anat_transform_list','commonspace_to_anat_inverse_list'],
                                               function=prep_commonspace_transform),
                                      name='prep_commonspace_transform')
    prep_commonspace_transform_node.inputs.atlas_mask = str(opts.brain_mask)
    prep_commonspace_transform_node.inputs.fast_commonspace = fast_commonspace

    workflow.connect([
        (inputnode, source_join_common_reg, [
            ("moving_image", "file_list0"),
            ("moving_mask", "file_list1"),
            ]),
        (template_inputnode, atlas_reg, [
            ("template_anat", "fixed_image"),
            ("template_mask", "fixed_mask"),
            ]),
        (atlas_reg, prep_commonspace_transform_node, [
            ("affine", "to_atlas_affine"),
            ("warp", "to_atlas_warp"),
            ("inverse_warp", "to_atlas_inverse_warp"),
            ]),
        (atlas_reg, outputnode, [
            ("affine", "to_atlas_affine"),
            ("warp", "to_atlas_warp"),
            ("inverse_warp", "to_atlas_inverse_warp"),
            ]),
        (prep_commonspace_transform_node, outputnode, [
            ('anat_to_commonspace_transform_list', 'anat_to_commonspace_transform_list'),
            ('anat_to_commonspace_inverse_list', 'anat_to_commonspace_inverse_list'),
            ('commonspace_to_anat_transform_list', 'commonspace_to_anat_transform_list'),
            ('commonspace_to_anat_inverse_list', 'commonspace_to_anat_inverse_list'),
            ('anatspace_mask', 'anatspace_mask'),
            ]),
        ])

    if fast_commonspace:
        prep_commonspace_transform_node.inputs.anat_to_unbiased_affine = None
        prep_commonspace_transform_node.inputs.anat_to_unbiased_warp = None
        prep_commonspace_transform_node.inputs.anat_to_unbiased_inverse_warp = None

        PlotOverlap_Anat2Atlas_node = pe.Node(
            PlotOverlap(), name='PlotOverlap_Anat2Atlas')
        PlotOverlap_Anat2Atlas_node.inputs.out_dir = output_folder+f'/preprocess_QC_report/{name}.Anat2Atlas/'

        workflow.connect([
            (merged_join_common_reg, atlas_reg, [
                ("file_list0", "moving_image"),
                ]),
            (merged_join_common_reg, prep_commonspace_transform_node, [
                ("file_list0", "anatspace_ref"),
                ]),
            (inputnode, PlotOverlap_Anat2Atlas_node, [
                ("moving_image", "name_source"),
                ]),
            (template_inputnode, PlotOverlap_Anat2Atlas_node,[
                ("template_anat", "fixed"),
                ]),
            (atlas_reg, PlotOverlap_Anat2Atlas_node, [
                ("warped_image", "moving"),
                ]),
            ])
        if commonspace_masking:
            workflow.connect([
                (merged_join_common_reg, atlas_reg, [
                    ("file_list1", "moving_mask"),
                    ]),
                ])
        if output_datasinks:
            workflow.connect([
                (atlas_reg, transforms_datasink, [
                    ("affine", "anat_to_atlas_affine"),
                    ("warp", "anat_to_atlas_warp"),
                    ("inverse_warp", "anat_to_atlas_inverse_warp"),
                    ]),
                ])

    else:

        if inherit_unbiased:
            # Identity node to relate inherited files to existing workflow
            generate_template = pe.Node(niu.IdentityInterface(fields=['moving_image_list', 'moving_mask_list', 'template_anat', 
                                                                      'unbiased_template', 'unbiased_mask', 'affine_list', 'warp_list', 
                                                                      'inverse_warp_list', 'warped_image_list']),
                                                name="generate_template_inherited")
        else:
            # setup a node to select the proper files associated with a given input scan for commonspace registration
            commonspace_selectfiles = pe.Node(Function(input_names=['filename', 'anat_list', 'affine_list', 'warp_list', 'inverse_warp_list', 'warped_anat_list'],
                                                    output_names=[
                                                        'anatspace_ref', 'warped_anat', 'anat_to_unbiased_affine','anat_to_unbiased_warp','anat_to_unbiased_inverse_warp'],
                                                    function=select_commonspace_outputs),
                                            name='commonspace_selectfiles')

            generate_template_outputs = f'{output_folder}/main_wf/{name}/generate_template'
            # the plugin is set here instead of as a node input to avoid re-running the nipype workflow if plugin is changed 
            os.environ['MODELBUILD_PLUGIN'] = opts.plugin
            # each modelbuild thread will be running a registration, and each registration uses RABIES_ITK_NUM_THREADS threads
            os.environ['MODELBUILD_NCORES'] = str(int(nthreads_modelbuild//int(os.environ['RABIES_ITK_NUM_THREADS'])))
            generate_template = pe.Node(GenerateTemplate(stages=modelbuild_stages, masking=commonspace_masking, output_folder=generate_template_outputs, winsorize_lower_bound = winsorize_lower_bound, winsorize_upper_bound = winsorize_upper_bound,
                                                ),
                                        name='generate_template', n_procs=nthreads_modelbuild, mem_gb=1*nthreads_modelbuild*opts.scale_min_memory)

        PlotOverlap_Anat2Unbiased_node = pe.Node(
            PlotOverlap(), name='PlotOverlap_Anat2Unbiased')
        PlotOverlap_Anat2Unbiased_node.inputs.out_dir = output_folder+f'/preprocess_QC_report/{name}.Anat2Unbiased/'
        PlotOverlap_Unbiased2Atlas_node = pe.Node(
            PlotOverlap(), name='PlotOverlap_Unbiased2Atlas')
        PlotOverlap_Unbiased2Atlas_node.inputs.out_dir = output_folder+f'/preprocess_QC_report/{name}.Unbiased2Atlas'
        PlotOverlap_Unbiased2Atlas_node.inputs.name_source = ''

        def prep_atlas_to_unbiased_transforms(to_atlas_affine,to_atlas_inverse_warp):
            transform_list=[to_atlas_affine,to_atlas_inverse_warp]
            inverse_list=[1,0]
            return transform_list,inverse_list
        atlas_to_unbiased_transforms_node = pe.Node(Function(input_names=['to_atlas_affine','to_atlas_inverse_warp'],
                                                output_names=[
                                                    'transform_list','inverse_list'],
                                                function=prep_atlas_to_unbiased_transforms),
                                        name='atlas_to_unbiased_transforms')

        resample_unbiased_mask_node = pe.Node(ResampleMask(), name='resample_unbiased_mask')
        resample_unbiased_mask_node.inputs.name_suffix = 'unbiased_mask_resampled'

        if brain_extraction and commonspace_masking:
            template_masking_node = pe.Node(Function(input_names=['template', 'mask', 'out_dir', 'figure_format'],
                                            function=template_masking),
                                    name='template_masking')
            template_masking_node.inputs.out_dir = output_folder+f'/preprocess_QC_report/{name}.unbiased_template_masking/'
            template_masking_node.inputs.figure_format = opts.figure_format

            workflow.connect([
                (generate_template, template_masking_node, [
                    ("unbiased_template", "template"),
                    ("unbiased_mask", "mask"),
                    ]),
                ])

        workflow.connect([
            (template_inputnode, generate_template, [
                ("template_anat", "template_anat"),
                ]),
            (merged_join_common_reg, generate_template, [
                ("file_list0", "moving_image_list"),
                ]),
            (template_inputnode, PlotOverlap_Unbiased2Atlas_node,[
                ("template_anat", "fixed"),
                ]),
            (inputnode, PlotOverlap_Anat2Unbiased_node, [
                ("moving_image", "name_source"),
                ]),
            (generate_template, atlas_reg, [
                ("unbiased_template", "moving_image"),
                ]),
            (generate_template, outputnode, [
                ("unbiased_template", "unbiased_template"),
                ]),
            (atlas_reg, atlas_to_unbiased_transforms_node, [
                ("affine", "to_atlas_affine"),
                ("inverse_warp", "to_atlas_inverse_warp"),
                ]),
            (atlas_to_unbiased_transforms_node, resample_unbiased_mask_node, [
                ("transform_list", "transforms"),
                ("inverse_list", "inverses"),
                ]),
            (generate_template, resample_unbiased_mask_node, [
                ("unbiased_template", "ref_file"),
                ]),
            (template_inputnode, resample_unbiased_mask_node, [
                ("template_mask", "mask_file"),
                ("template_mask", "name_source"),
                ]),
            (resample_unbiased_mask_node, outputnode, [
                ("resampled_file", "unbiased_mask"),
                ]),
            (generate_template, PlotOverlap_Anat2Unbiased_node, [
                ("unbiased_template", "fixed"),
                ]),
            (atlas_reg, PlotOverlap_Unbiased2Atlas_node, [
                ("warped_image", "moving"),
                ]),
            ])
        if output_datasinks:
            unbiased_template_datasink = pe.Node(DataSink(base_directory=output_folder,
                                                container="unbiased_template_datasink"),
                                        name="unbiased_template_datasink")

            workflow.connect([
                (outputnode, unbiased_template_datasink, [
                    ("unbiased_template", "unbiased_template"),
                    ]),
                (atlas_reg, unbiased_template_datasink, [
                    ("warped_image", "warped_unbiased_template"),
                    ]),
                ])

        if commonspace_masking:
            workflow.connect([
                (merged_join_common_reg, generate_template, [
                    ("file_list1", "moving_mask_list"),
                    ]),
                (generate_template, atlas_reg, [
                    ("unbiased_mask", "moving_mask"),
                    ]),
                ])

        if output_datasinks:
            workflow.connect([
                (atlas_reg, transforms_datasink, [
                    ("affine", "unbiased_to_atlas_affine"),
                    ("warp", "unbiased_to_atlas_warp"),
                    ("inverse_warp", "unbiased_to_atlas_inverse_warp"),
                    ]),
                ])

        if inherit_unbiased:
            inherit_unbiased_inputnode = pe.Node(niu.IdentityInterface(fields=['unbiased_template', 'unbiased_mask', 'unbiased_to_atlas_affine', 'unbiased_to_atlas_warp',
                                                                            'unbiased_to_atlas_inverse_warp', 'warped_unbiased']),
                                                name="inherit_unbiased_inputnode")

            inherit_unbiased_reg_node = pe.Node(Function(input_names=['reg_method', 'brain_extraction', 'keep_mask_after_extract', 'moving_image', 'moving_mask', 'fixed_image', 'fixed_mask', 'winsorize_lower_bound', 'winsorize_upper_bound','rabies_data_type'],
                                            output_names=['affine', 'warp',
                                                        'inverse_warp', 'warped_image'],
                                            function=run_antsRegistration),
                                name='inherit_unbiased_reg', mem_gb=2*opts.scale_min_memory)
            inherit_unbiased_reg_node.inputs.winsorize_lower_bound = winsorize_lower_bound
            inherit_unbiased_reg_node.inputs.winsorize_upper_bound = winsorize_upper_bound
            inherit_unbiased_reg_node.inputs.reg_method = 'Rigid' # this is the registration modelbuild conducts
            inherit_unbiased_reg_node.inputs.brain_extraction = False # brain extraction is not applied during template building
            inherit_unbiased_reg_node.inputs.keep_mask_after_extract = False # brain extraction is not applied during template building
            inherit_unbiased_reg_node.inputs.rabies_data_type = opts.data_type

            if commonspace_masking: # this parameter should be inherited from previous run
                workflow.connect([
                    (inputnode, inherit_unbiased_reg_node, [
                        ("moving_mask", "moving_mask"),
                        ]),
                    (generate_template, inherit_unbiased_reg_node, [
                        ("unbiased_mask", "fixed_mask"),
                        ]),
                    ])
            else:
                inherit_unbiased_reg_node.inputs.moving_mask = 'NULL'
                inherit_unbiased_reg_node.inputs.fixed_mask = 'NULL'

            workflow.connect([
                (inherit_unbiased_inputnode, generate_template, [
                    ("unbiased_template", "unbiased_template"),
                    ("unbiased_mask", "unbiased_mask"),
                    ]),
                (inherit_unbiased_inputnode, atlas_reg, [
                    ("unbiased_to_atlas_affine", "affine"),
                    ("unbiased_to_atlas_warp", "warp"),
                    ("unbiased_to_atlas_inverse_warp", "inverse_warp"),
                    ("warped_unbiased", "warped_image"),
                    ]),
                (inputnode, inherit_unbiased_reg_node, [
                    ("moving_image", "moving_image"),
                    ]),
                (generate_template, inherit_unbiased_reg_node, [
                    ("unbiased_template", "fixed_image"),
                    ]),
                (inherit_unbiased_reg_node, prep_commonspace_transform_node, [
                    ("affine", "anat_to_unbiased_affine"),
                    ("warp", "anat_to_unbiased_warp"),
                    ("inverse_warp", "anat_to_unbiased_inverse_warp"),
                    ]),
                (inherit_unbiased_reg_node, outputnode, [
                    ("affine", "anat_to_unbiased_affine"),
                    ("warp", "anat_to_unbiased_warp"),
                    ("inverse_warp", "anat_to_unbiased_inverse_warp"),
                    ]),
                (merged_join_common_reg, prep_commonspace_transform_node, [
                    ("file_list0", "anatspace_ref"),
                    ]),
                (inherit_unbiased_reg_node, PlotOverlap_Anat2Unbiased_node, [
                    ("warped_image", "moving"),
                    ]),
                ])
            
        else:
            workflow.connect([
                (merged_join_common_reg, commonspace_selectfiles, [
                    ("file_list0", "anat_list"),
                    ]),
                (inputnode, commonspace_selectfiles, [
                    ("moving_image", "filename"),
                    ]),
                (generate_template, commonspace_selectfiles, [
                    ("affine_list", "affine_list"),
                    ("warp_list", "warp_list"),
                    ("inverse_warp_list", "inverse_warp_list"),
                    ("warped_image_list", "warped_anat_list"),
                    ]),
                (commonspace_selectfiles, prep_commonspace_transform_node, [
                    ("anatspace_ref", "anatspace_ref"),
                    ("anat_to_unbiased_affine", "anat_to_unbiased_affine"),
                    ("anat_to_unbiased_warp", "anat_to_unbiased_warp"),
                    ("anat_to_unbiased_inverse_warp", "anat_to_unbiased_inverse_warp"),
                    ]),
                (commonspace_selectfiles, outputnode, [
                    ("anat_to_unbiased_affine", "anat_to_unbiased_affine"),
                    ("anat_to_unbiased_warp", "anat_to_unbiased_warp"),
                    ("anat_to_unbiased_inverse_warp", "anat_to_unbiased_inverse_warp"),
                    ]),
                (commonspace_selectfiles, PlotOverlap_Anat2Unbiased_node, [
                    ("warped_anat", "moving"),
                    ]),
                ])
            
            if output_datasinks:
                workflow.connect([
                    (commonspace_selectfiles, transforms_datasink, [
                        ("anat_to_unbiased_affine", "anat_to_unbiased_affine"),
                        ("anat_to_unbiased_warp", "anat_to_unbiased_warp"),
                        ("anat_to_unbiased_inverse_warp", "anat_to_unbiased_inverse_warp"),
                        ]),
                    ])


    return workflow


def join_iterables(workflow, joinsource_list, node_prefix, num_inputs=1):

    field_list=[]
    for j in range(num_inputs):
        field_list.append(f'file_list{j}')

    i=0
    for joinsource in joinsource_list:
        joinnode = pe.JoinNode(niu.IdentityInterface(fields=field_list),
                                            name=f"{node_prefix}_{joinsource}_joinnode",
                                            joinsource=joinsource,
                                            joinfield=field_list)
        if i==0:
            source_join = joinnode
        else:
            for field in field_list:
                workflow.connect([
                    (joinnode_prev, joinnode, [
                        (field, field),
                        ]),
                    ])

        joinnode_prev = joinnode
        i+=1

    merged_join = joinnode

    return workflow, source_join, merged_join


def select_commonspace_outputs(filename, anat_list, affine_list, warp_list, inverse_warp_list, warped_anat_list):
    from rabies.preprocess_pkg.commonspace_reg import select_from_list
    anat_to_unbiased_affine = select_from_list(filename, affine_list)
    anat_to_unbiased_warp = select_from_list(filename, warp_list)
    anat_to_unbiased_inverse_warp = select_from_list(
        filename, inverse_warp_list)
    warped_anat = select_from_list(filename, warped_anat_list)
    anatspace_ref = select_from_list(filename, anat_list)
    return anatspace_ref, warped_anat, anat_to_unbiased_affine, anat_to_unbiased_warp, anat_to_unbiased_inverse_warp


def select_from_list(filename, filelist):
    filelist = flatten_list(filelist)
    import pathlib
    filename_template = pathlib.Path(filename).name.rsplit(".nii")[0]
    selected_file = None
    for file in filelist:
        if filename_template in file:
            if selected_file is None:
                selected_file = file
            else:
                raise ValueError(
                    f"Found duplicates for filename {filename_template}.")

    if selected_file is None:
        raise ValueError(f"No file associated with {filename_template} were found from list {filelist}.")
    else:
        return selected_file


def prep_commonspace_transform(anatspace_ref, atlas_mask, anat_to_unbiased_affine,
                               anat_to_unbiased_warp,anat_to_unbiased_inverse_warp,
                               to_atlas_affine,to_atlas_warp,to_atlas_inverse_warp, fast_commonspace=False):

    if fast_commonspace:
        anat_to_commonspace_transform_list=[to_atlas_warp,to_atlas_affine]
        anat_to_commonspace_inverse_list=[0,0]
        commonspace_to_anat_transform_list=[to_atlas_affine,to_atlas_inverse_warp]
        commonspace_to_anat_inverse_list=[1,0]
    else:
        anat_to_commonspace_transform_list=[to_atlas_warp,to_atlas_affine,anat_to_unbiased_warp,anat_to_unbiased_affine]
        anat_to_commonspace_inverse_list=[0,0,0,0]
        commonspace_to_anat_transform_list=[anat_to_unbiased_affine,anat_to_unbiased_inverse_warp,to_atlas_affine,to_atlas_inverse_warp]
        commonspace_to_anat_inverse_list=[1,0,1,0]

    import os
    import pathlib  # Better path manipulation
    filename_split = pathlib.Path(
        anatspace_ref).name.rsplit(".nii")
    anatspace_mask = os.path.abspath(filename_split[0]+'_mask.nii.gz')
    from rabies.utils import antsApplyTransforms
    # resample the atlas brain mask to anat space
    antsApplyTransforms(transforms = commonspace_to_anat_transform_list, inverses = commonspace_to_anat_inverse_list, 
                       input_image = atlas_mask, ref_image = anatspace_ref, output_filename = anatspace_mask, interpolation='GenericLabel', rabies_data_type=None, clip_negative=False)

    return anatspace_mask, anat_to_commonspace_transform_list,anat_to_commonspace_inverse_list,commonspace_to_anat_transform_list,commonspace_to_anat_inverse_list


class GenerateTemplateInputSpec(BaseInterfaceInputSpec):
    #these input traits are later inherited in the GenerateTemplate class
    moving_image_list = traits.List(exists=True, mandatory=True,
                            desc="List of anatomical images used for commonspace registration.")
    moving_mask_list = traits.List(exists=True,
                            desc="List of masks accompanying each image. (optional)")
    stages = traits.Str(
        exists=True, mandatory=True, desc="Input to --stages parameter of modelbuild.sh")
    masking = traits.Bool(
        desc="Whether to use the masking option.")
    output_folder = traits.Str(
        exists=True, mandatory=True, desc="Path to output folder.")
    template_anat = File(exists=True, mandatory=True,
                         desc="Reference anatomical template to define the target space.")
    winsorize_lower_bound = traits.Float(
        desc="")
    winsorize_upper_bound = traits.Float(
        desc="")

class GenerateTemplateOutputSpec(TraitedSpec):
    unbiased_template = File(
        exists=True, mandatory=True, desc="Output template generated from commonspace registration.")
    unbiased_mask = File(desc="Output mask generated along with the template. (optional)")
    affine_list = traits.List(exists=True, mandatory=True,
                              desc="List of affine transforms from anat to template space.")
    warp_list = traits.List(exists=True, mandatory=True,
                            desc="List of non-linear transforms from anat to template space.")
    inverse_warp_list = traits.List(exists=True, mandatory=True,
                                    desc="List of the inverse non-linear transforms from anat to template space.")
    warped_image_list = traits.List(exists=True, mandatory=True,
                                   desc="List of anatomical images warped to template space..")


class GenerateTemplate(BaseInterface):
    """

    """

    input_spec = GenerateTemplateInputSpec
    output_spec = GenerateTemplateOutputSpec

    def _run_interface(self, runtime):
        import os
        import pandas as pd
        import pathlib
        import multiprocessing
        from nipype import logging
        log = logging.getLogger('nipype.workflow')

        cwd = os.getcwd()
        template_folder = self.inputs.output_folder
        winsorize_lower_bound = self.inputs.winsorize_lower_bound
        winsorize_upper_bound = self.inputs.winsorize_upper_bound
        stages = self.inputs.stages

        # change the - split for a comma ,
        stages = ','.join(stages.split('-'))

        for stage in stages.split(','):
            if stage not in ['rigid', 'similarity', 'affine', 'nlin']:
                raise ValueError(f"{stage} is not a proper modelbuild.sh stage - it must be one of 'rigid', 'similarity', 'affine', 'nlin'. Review your --commonspace_reg parameter.")

        command = f'mkdir -p {template_folder}'
        rc,c_out = run_command(command)

        merged = flatten_list(list(self.inputs.moving_image_list))
        # create a csv file of the input image list
        csv_path = cwd+'/commonspace_input_files.csv'
        df = pd.DataFrame(data=merged)
        df.to_csv(csv_path, header=False, sep=',', index=False)

        if self.inputs.masking:
            merged_masks = flatten_list(list(self.inputs.moving_mask_list))
            for mask in merged_masks:
                if 'NULL' in mask:
                    merged_masks = ['NULL'] # there should be no NULL mask
            if not (merged_masks[0]=='NULL'):
                mask_csv_path = cwd+'/commonspace_input_masks.csv'
                df = pd.DataFrame(data=merged_masks)
                df.to_csv(mask_csv_path, header=False, sep=',', index=False)
                masks = f'--masks {mask_csv_path}'
            else:
                log.info(f"Some 'NULL' input mask was found within moving_mask_list. No moving mask will be used by modelbuild.sh")
        else:
            merged_masks = ['NULL']
            masks=''
            unbiased_mask='NULL'

        if len(merged) == 1:
            log.info("Only a single scan was provided as input for commonspace registration. Commonspace registration "
                  "won't be run, and the output template will be the input scan.")

            # create an identity transform as a surrogate for the commonspace transforms
            import SimpleITK as sitk
            dimension = 3
            identity = sitk.Transform(dimension, sitk.sitkIdentity)

            file = merged[0]
            mask = merged_masks[0]
            filename_template = pathlib.Path(file).name.rsplit(".nii")[0]
            transform_file = template_folder+filename_template+'_identity.mat'
            sitk.WriteTransform(identity, transform_file)

            setattr(self, 'unbiased_template', file)
            setattr(self, 'unbiased_mask', mask)
            setattr(self, 'affine_list', [transform_file])
            setattr(self, 'warp_list', [transform_file])
            setattr(self, 'inverse_warp_list', [transform_file])
            setattr(self, 'warped_image_list', [file])

            return runtime


        # convert nipype plugin spec to match QBATCH
        plugin = os.environ['MODELBUILD_PLUGIN']
        if plugin=='MultiProc' or plugin=='Linear':
            cluster_type='local'
            num_threads = os.environ['MODELBUILD_NCORES']
        elif plugin=='SGE' or plugin=='SGEGraph':
            cluster_type='sge'
            num_threads = 1
        elif plugin=='PBS':
            cluster_type='PBS'
            num_threads = 1
        elif plugin=='SLURM' or plugin=='SLURMGraph':
            cluster_type='slurm'
            num_threads = 1
        else:
            raise ValueError("Plugin option must correspond to one of 'local', 'sge', 'pbs' or 'slurm'")

        # when the initial target template has the filename of a previous modelbuild, it will be mistaken
        # for an intermediate modelbuild step and include --close and --initial-transform to the registration.
        # To avoid this, the template file is renamed to another generic filename. 
        command = f'cp {self.inputs.template_anat} {template_folder}/modelbuild_starting_target.nii.gz'
        rc,c_out = run_command(command)
        log.debug(f"The --starting-target template original file is {self.inputs.template_anat}, and was renamed to {template_folder}/modelbuild_starting_target.nii.gz.")

        command = f'QBATCH_SYSTEM={cluster_type} QBATCH_CORES={num_threads} modelbuild.sh \
            --float --average-type median --gradient-step 0.25 --iterations 2 --starting-target {template_folder}/modelbuild_starting_target.nii.gz --stages {stages} \
            --output-dir {template_folder} --sharpen-type unsharp --block --debug {masks} {csv_path}'
        if winsorize_lower_bound>0.0: # only specify an input if they differ from the default
            command+=f" --winsorize_lower_bound {winsorize_lower_bound}"
        if winsorize_upper_bound<1.0: # only specify an input if they differ from the default
            command+=f" --winsorize_upper_bound {winsorize_upper_bound}"
        rc,c_out = run_command(command)

        last_stage = stages.split(',')[-1]

        unbiased_template = template_folder + \
            f'/{last_stage}/1/average/template_sharpen_shapeupdate.nii.gz'
        # verify that all outputs are present
        if not os.path.isfile(unbiased_template):
            raise ValueError(unbiased_template+" doesn't exists.")

        if self.inputs.masking and not (merged_masks[0]=='NULL'):
            unbiased_mask = template_folder + \
                f'/{last_stage}/1/average/mask_shapeupdate.nii.gz'
            # verify that all outputs are present
            if not os.path.isfile(unbiased_mask):
                raise ValueError(unbiased_mask+" doesn't exists.")

        affine_list = []
        warp_list = []
        inverse_warp_list = []
        warped_image_list = []

        import SimpleITK as sitk
        dimension = 3
        identity = sitk.Transform(dimension, sitk.sitkIdentity)

        nlin_out = last_stage=='nlin'
        i = 0
        for file in merged:
            file = str(file)
            filename_template = pathlib.Path(file).name.rsplit(".nii")[0]

            if not nlin_out:
                # create a surrogate transform that takes the place of non-linear since nlin stage is not run
                surrogate_transform_file = f'{template_folder}/{filename_template}_surrogate_identity_transform.mat'
                sitk.WriteTransform(identity, surrogate_transform_file)
                anat_to_unbiased_inverse_warp = surrogate_transform_file
                anat_to_unbiased_warp = surrogate_transform_file
            else:
                anat_to_unbiased_inverse_warp = f'{template_folder}/{last_stage}/1/transforms/{filename_template}_1InverseWarp.nii.gz'
                anat_to_unbiased_warp = f'{template_folder}/{last_stage}/1/transforms/{filename_template}_1Warp.nii.gz'

            if not os.path.isfile(anat_to_unbiased_inverse_warp):
                raise ValueError(
                    anat_to_unbiased_inverse_warp+" file doesn't exists.")
            if not os.path.isfile(anat_to_unbiased_warp):
                raise ValueError(anat_to_unbiased_warp+" file doesn't exists.")
            
            anat_to_unbiased_affine = f'{template_folder}/{last_stage}/1/transforms/{filename_template}_0GenericAffine.mat'
            if not os.path.isfile(anat_to_unbiased_affine):
                raise ValueError(anat_to_unbiased_affine
                                 + " file doesn't exists.")
            warped_image = f'{template_folder}/{last_stage}/1/resample/{filename_template}.nii.gz'
            if not os.path.isfile(warped_image):
                raise ValueError(warped_image
                                 + " file doesn't exists.")
            inverse_warp_list.append(anat_to_unbiased_inverse_warp)
            warp_list.append(anat_to_unbiased_warp)
            affine_list.append(anat_to_unbiased_affine)
            warped_image_list.append(warped_image)
            i += 1

        setattr(self, 'unbiased_template', unbiased_template)
        setattr(self, 'unbiased_mask', unbiased_mask)
        setattr(self, 'affine_list', affine_list)
        setattr(self, 'warp_list', warp_list)
        setattr(self, 'inverse_warp_list', inverse_warp_list)
        setattr(self, 'warped_image_list', warped_image_list)

        return runtime

    def _list_outputs(self):
        return {'unbiased_template': getattr(self, 'unbiased_template'),
                'unbiased_mask': getattr(self, 'unbiased_mask'),
                'affine_list': getattr(self, 'affine_list'),
                'warp_list': getattr(self, 'warp_list'),
                'inverse_warp_list': getattr(self, 'inverse_warp_list'),
                'warped_image_list': getattr(self, 'warped_image_list'), }


def inherit_unbiased_files(RABIES_output_path, opts):
    import os
    from nipype import logging
    log = logging.getLogger('nipype.workflow')
    from rabies.utils import get_workflow_dict
    import pickle
    cli_file = f'{RABIES_output_path}/rabies_preprocess.pkl'
    with open(cli_file, 'rb') as handle:
        inherit_preprocess_opts = pickle.load(handle)

    if inherit_preprocess_opts.commonspace_reg['fast_commonspace']:
        raise ValueError("fast_commonspace was conducted in the previous run; no unbiased template can be inherited.")

    log.warning("--inherit_unbiased_template will be applied. The following options are inherited from the previous run: \n \
            --anatomical_resampling \n \
            --commonspace_reg \n \
            --anat_template \n \
            --brain_mask \n \
            --WM_mask \n \
            --CSF_mask \n \
            --vascular_mask \n")
    
    opts.anatomical_resampling = inherit_preprocess_opts.anatomical_resampling
    opts.commonspace_reg = inherit_preprocess_opts.commonspace_reg
    opts.anat_template = inherit_preprocess_opts.anat_template
    opts.brain_mask = inherit_preprocess_opts.brain_mask
    opts.WM_mask = inherit_preprocess_opts.WM_mask
    opts.CSF_mask = inherit_preprocess_opts.CSF_mask
    opts.vascular_mask = inherit_preprocess_opts.vascular_mask

    preproc_workflow_file = f'{RABIES_output_path}/rabies_preprocess_workflow.pkl'
    node_dict = get_workflow_dict(preproc_workflow_file)

    inherit_dict = {}

    inherit_dict['registration_template'] = node_dict['main_wf.resample_template'].result.outputs.get()['registration_template']
    inherit_dict['registration_mask'] = node_dict['main_wf.resample_template'].result.outputs.get()['registration_mask']
    inherit_dict['commonspace_template'] = node_dict['main_wf.resample_template'].result.outputs.get()['commonspace_template']

    inherit_dict['unbiased_template'] = node_dict['main_wf.commonspace_reg_wf.generate_template'].result.outputs.get()['unbiased_template']
    inherit_dict['unbiased_mask'] = node_dict['main_wf.commonspace_reg_wf.generate_template'].result.outputs.get()['unbiased_mask']

    inherit_dict['unbiased_to_atlas_affine'] = node_dict['main_wf.commonspace_reg_wf.atlas_reg'].result.outputs.get()['affine']
    inherit_dict['unbiased_to_atlas_warp'] = node_dict['main_wf.commonspace_reg_wf.atlas_reg'].result.outputs.get()['warp']
    inherit_dict['unbiased_to_atlas_inverse_warp'] = node_dict['main_wf.commonspace_reg_wf.atlas_reg'].result.outputs.get()['inverse_warp']
    inherit_dict['warped_unbiased'] = node_dict['main_wf.commonspace_reg_wf.atlas_reg'].result.outputs.get()['warped_image']

    # check that the files exist
    for key in list(inherit_dict.keys()):
        try:
            if not os.path.isfile(inherit_dict[key]):
                raise ValueError(f"File not found for {key}. {inherit_dict[key]} was found instead.")
        except:
            if key=='unbiased_mask': 
                if opts.commonspace_reg['masking']: # the file is only present when masking was used
                    raise ValueError(f"File not found for {key}. {inherit_dict[key]} was found instead.")
            elif key=='unbiased_to_atlas_warp' or key=='unbiased_to_atlas_inverse_warp': # if registration was rigid or affine, no warp files
                if not inherit_dict[key]=='NULL':
                    raise ValueError(f"File not found for {key}. {inherit_dict[key]} was found instead.")
            else:
                raise ValueError(f"File not found for {key}. {inherit_dict[key]} was found instead.")
                
    return opts, inherit_dict