# coding: utf8


def t1_b0_registration(
        participant_id, session_id,
        caps_directory, working_directory=None,
        name="t1_b0_registration"):
    """
    Perform rigid registration of the T1-weighted image onto the B0 image.

    Given a subject, this pipelines performs a registration between its
    T1-weighted image onto its 1mm upsampled B0 image. Once done, the estimated
    transformation matrix is used to align the binary mask of the segmentation
    of the white matter (obtained with FSL fast) from the anatomical space to
    the diffusion space. The same operation is performed on the Desikan and
    Destrieux parcellation (obtained with FreeSurfer recon-all) except that
    these latter are not resliced on the B0 image.

    These steps enable to prepare the data for the DWI processing pipelines.

    Args:
         participant_id (str): Subject (participant) ID in a BIDS format
            ('sub-<participant_label>').
         session_id (str): Session ID in a BIDS format ('ses-<session_label>').
         caps_directory (str): Directory where the results are stored
            in a CAPS hierarchy.
         working_directory (Optional[str]): Directory where the temporary
            results are stored. If not specified, it is automatically generated
            (generally in /tmp/).
         name (Optional[str]): Name of the pipelines.

    Inputnode:
        in_bias_corrected_bet_t1 (str): File containing the bias corrected
            brain extracted T1-weighted image. It corresponds to the
            out_brain_extracted file from FSL-T1 pipelines.
        in_preprocessed_dwi (str): File containing the preprocessed DWI
            dataset. It assumes that the reference b0 is the first volume in
            the dataset (which is the case if you are using Clinica).
        in_b0_mask (str): File containing the mask of the b0 image. It assumes
            that `in_b0_mask` has the same header as `in_preprocessed_dwi`
            (which is the case if you are using Clinica).
        in_white_matter_binary_mask (str): File containing the binary
            segmentation of the white matter (obtained with FSL fast). You can
            use the one generated by FreeSurfer recon-all but you must convert
            your image in FreeSurfer space to Native space first.
        in_desikan_parcellation (str): File containing the Desikan parcellation
            (obtained with FreeSurfer recon-all). The file is usually located
            in ${participant_id}/mri/aparc+aseg.mgz.
        in_destrieux_parcellation (str): File containing the Destrieux
            parcellation (obtained with FreeSurfer recon-all). The file is
            usually located in ${participant_id}/mri/aparc.a2009s+aseg.mgz.

    Outputnode:
        out_registered_t1 (str): File containing the registration of the
            T1-weighted image onto the diffusion space.
        out_flirt_matrix (str): File containing the transformation matrix
            estimated by FSL flirt.
        out_wm_mask_in_diffusion_space (str): File containing the segmentation
            of the white matter in diffusion space.
        out_mrtrix_matrix (str): File containing the transformation matrix in
            MRtrix format (can be used with the mrtransform command).
        out_desikan_in_diffusion_space (str): File containing the Desikan
            parcellation in diffusion space.
        out_destrieux_in_diffusion_space (str): File containing the Desikan
            parcellation in diffusion space.

    """
    from os.path import join
    import nipype.interfaces.fsl as fsl
    import nipype.interfaces.io as nio
    import nipype.interfaces.utility as niu
    import nipype.pipeline.engine as pe
    from nipype.interfaces.freesurfer.preprocess import MRIConvert
    from clinica.utils.freesurfer import freesurfer_volume_to_native_volume
    from clinica.utils.mri_registration import convert_flirt_transformation_to_mrtrix_transformation
    from clinica.utils.mri_registration import apply_mrtrix_transform_without_resampling
    import tempfile
    from clinica.utils.check_dependency import check_freesurfer, check_fsl, check_mrtrix

    if working_directory is None:
        working_directory = tempfile.mkdtemp()

    check_freesurfer()
    check_fsl()
    check_mrtrix()

    caps_identifier = participant_id + '_' + session_id

    inputnode = pe.Node(niu.IdentityInterface(
        fields=['in_bias_corrected_bet_t1', 'in_preprocessed_dwi',
                'in_b0_mask', 'in_white_matter_binary_mask',
                'in_desikan_parcellation', 'in_destrieux_parcellation']),
        name='inputnode')

    get_b0 = pe.Node(fsl.ExtractROI(args='0 1'), name='get_b0')

    upsample_b0 = pe.Node(
        MRIConvert(vox_size=(1, 1, 1), out_type='niigz'),
        name='upsample_b0')

    upsample_b0_mask = pe.Node(
        MRIConvert(vox_size=(1, 1, 1), out_type='niigz'),
        name='upsample_b0_mask')

    registration_t1_to_b0 = pe.Node(fsl.FLIRT(
        dof=6, interp='spline', cost='normmi', cost_func='normmi',
        out_matrix_file=caps_identifier + '_t1-to-b0_withResampling.mat'),
        name='registration_t1_to_b0')

    apply_flirt_registration = pe.Node(
        fsl.ApplyXfm(apply_xfm=True, interp='spline'),
        name='apply_flirt_registration')
    apply_flirt_registration.inputs.out_file = \
        caps_identifier + '_binarymask-whitematter_reslicedOnDiffusionSpace.nii.gz'

    convert_flirt_to_mrtrix = pe.Node(interface=niu.Function(
        input_names=['in_source_image', 'in_reference_image',
                     'in_flirt_matrix', 'name_output_matrix'],
        output_names=['out_mrtrix_matrix'],
        function=convert_flirt_transformation_to_mrtrix_transformation),
        name='convert_flirt_to_mrtrix')
    convert_flirt_to_mrtrix.inputs.name_output_matrix = \
        caps_identifier + '_t1-to-b0_withoutResampling.mat'

    desikan_in_native_space = pe.Node(interface=niu.Function(
        input_names=['freesurfer_volume', 'native_volume', 'name_output_volume'],
        output_names=['out_volume'],
        function=freesurfer_volume_to_native_volume),
        name='desikan_in_native_space')
    destrieux_in_native_space = pe.Node(interface=niu.Function(
        input_names=['freesurfer_volume', 'native_volume', 'name_output_volume'],
        output_names=['out_volume'],
        function=freesurfer_volume_to_native_volume),
        name='destrieux_in_native_space')

    desikan_in_diffusion_space = pe.Node(interface=niu.Function(
        input_names=['in_image', 'in_mrtrix_matrix', 'name_output_image'],
        output_names=['out_deformed_image'],
        function=apply_mrtrix_transform_without_resampling),
        name='desikan_in_diffusion_space')
    desikan_in_diffusion_space.inputs.name_output_image = \
        caps_identifier + '_parcellation-desikan_onDiffusionSpace.nii.gz'
    destrieux_in_diffusion_space = pe.Node(interface=niu.Function(
        input_names=['in_image', 'in_mrtrix_matrix', 'name_output_image'],
        output_names=['out_deformed_image'],
        function=apply_mrtrix_transform_without_resampling),
        name='destrieux_in_diffusion_space')
    destrieux_in_diffusion_space.inputs.name_output_image = \
        caps_identifier + '_parcellation-destrieux_onDiffusionSpace.nii.gz'

    outputnode = pe.Node(niu.IdentityInterface(
        fields=['out_registered_t1', 'out_flirt_matrix',
                'out_wm_mask_in_diffusion_space', 'out_mrtrix_matrix',
                'out_desikan_in_diffusion_space',
                'out_destrieux_in_diffusion_space']),
        name='outputnode')

    datasink = pe.Node(nio.DataSink(), name='datasink')
    datasink.inputs.base_directory = join(caps_directory, 'subjects',
                                          participant_id, session_id)
#    datasink.inputs.substitutions = [('fast_pve_0.nii.gz', caps_identifier + '_binary-csf.nii.gz')]

    wf = pe.Workflow(name=name)
    wf.base_dir = working_directory
    wf.connect([
        # Get b0 from DWI:
        (inputnode, get_b0, [('in_preprocessed_dwi', 'in_file')]),
        # Upsample at 1mm the b0 image:
        (get_b0, upsample_b0, [('roi_file', 'in_file')]),
        # Upsample at 1mm the b0 mask:
        (inputnode, upsample_b0_mask, [('in_b0_mask', 'in_file')]),
        # Register the T1 image onto the b0:
        (inputnode,        registration_t1_to_b0, [('in_bias_corrected_bet_t1', 'in_file')]),
        (upsample_b0,      registration_t1_to_b0, [('out_file', 'reference')]),
        (upsample_b0_mask, registration_t1_to_b0, [('out_file', 'ref_weight')]),
        # Apply flirt registration to WM mask:
        (inputnode,             apply_flirt_registration, [('in_white_matter_binary_mask', 'in_file')]),
        (upsample_b0,           apply_flirt_registration, [('out_file', 'reference')]),
        (registration_t1_to_b0, apply_flirt_registration, [('out_matrix_file', 'in_matrix_file')]),
        # Convert flirt matrix to MRtrix matrix:
        (inputnode,             convert_flirt_to_mrtrix, [('in_bias_corrected_bet_t1', 'in_source_image')]),
        (upsample_b0,           convert_flirt_to_mrtrix, [('out_file', 'in_reference_image')]),
        (registration_t1_to_b0, convert_flirt_to_mrtrix, [('out_matrix_file', 'in_flirt_matrix')]),
        # Convert FreeSurfer parcellations into native space:
        (inputnode, desikan_in_native_space, [('in_desikan_parcellation', 'freesurfer_volume'),
                                              ('in_bias_corrected_bet_t1', 'native_volume')]),
        (inputnode, destrieux_in_native_space, [('in_destrieux_parcellation', 'freesurfer_volume'),
                                                ('in_bias_corrected_bet_t1', 'native_volume')]),
        # Apply registration without resampling on Desikan & Destrieux parcellations:
        (desikan_in_native_space, desikan_in_diffusion_space, [('out_volume', 'in_image')]),  # noqa
        (convert_flirt_to_mrtrix, desikan_in_diffusion_space, [('out_mrtrix_matrix', 'in_mrtrix_matrix')]),  # noqa
        (destrieux_in_native_space, destrieux_in_diffusion_space, [('out_volume', 'in_image')]),  # noqa
        (convert_flirt_to_mrtrix,   destrieux_in_diffusion_space, [('out_mrtrix_matrix', 'in_mrtrix_matrix')]),  # noqa
        # Outputnode:
        (registration_t1_to_b0,        outputnode, [('out_file',                          'out_registered_t1')]),  # noqa
        (registration_t1_to_b0,        outputnode, [('out_matrix_file',                    'out_flirt_matrix')]),  # noqa
        (apply_flirt_registration,     outputnode, [('out_file',             'out_wm_mask_in_diffusion_space')]),  # noqa
        (convert_flirt_to_mrtrix,      outputnode, [('out_mrtrix_matrix',                 'out_mrtrix_matrix')]),  # noqa
        (desikan_in_diffusion_space,   outputnode, [('out_deformed_image',   'out_desikan_in_diffusion_space')]),  # noqa
        (destrieux_in_diffusion_space, outputnode, [('out_deformed_image', 'out_destrieux_in_diffusion_space')]),  # noqa
        # Datasink:
        (registration_t1_to_b0,        datasink, [('out_file',                          'dwi.@out_registered_t1')]),  # noqa
        (registration_t1_to_b0,        datasink, [('out_matrix_file',                    'dwi.@out_flirt_matrix')]),  # noqa
        (apply_flirt_registration,     datasink, [('out_file',              'dwi.@out_wm_mask_in_diffusion_mask')]),  # noqa
        (convert_flirt_to_mrtrix,      datasink, [('out_mrtrix_matrix',                 'dwi.@out_mrtrix_matrix')]),  # noqa
        (desikan_in_diffusion_space,   datasink, [('out_deformed_image',   'dwi.@out_desikan_in_diffusion_space')]),  # noqa
        (destrieux_in_diffusion_space, datasink, [('out_deformed_image', 'dwi.@out_destrieux_in_diffusion_space')])   # noqa
    ])
    return wf
