# -*- coding: utf-8 -*-
# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:

import os
from warnings import warn

from ....pipeline import engine as pe
from ....interfaces import utility as util
from ....interfaces import fsl as fsl


def tbss1_op_string(in_files):
    import nibabel as nb
    from nipype.utils import NUMPY_MMAP
    op_strings = []
    for infile in in_files:
        img = nb.load(infile, mmap=NUMPY_MMAP)
        dimtup = tuple(d - 2 for d in img.shape)
        dimtup = dimtup[0:3]
        op_str = '-min 1 -ero -roi 1 %d 1 %d 1 %d 0 1' % dimtup
        op_strings.append(op_str)
    return op_strings


def create_tbss_1_preproc(name='tbss_1_preproc'):
    """Preprocess FA data for TBSS: erodes a little and zero end slicers and
    creates masks(for use in FLIRT & FNIRT from FSL).
    A pipeline that does the same as tbss_1_preproc script in FSL

    Example
    -------

    >>> from nipype.workflows.dmri.fsl import tbss
    >>> tbss1 = tbss.create_tbss_1_preproc()
    >>> tbss1.inputs.inputnode.fa_list = ['s1_FA.nii', 's2_FA.nii', 's3_FA.nii']

    Inputs::

        inputnode.fa_list

    Outputs::

        outputnode.fa_list
        outputnode.mask_list
        outputnode.slices

    """

    # Define the inputnode
    inputnode = pe.Node(
        interface=util.IdentityInterface(fields=["fa_list"]), name="inputnode")

    # Prep the FA images
    prepfa = pe.MapNode(
        fsl.ImageMaths(suffix="_prep"),
        name="prepfa",
        iterfield=['in_file', 'op_string'])

    # Slicer
    slicer = pe.MapNode(
        fsl.Slicer(all_axial=True, image_width=1280),
        name='slicer',
        iterfield=['in_file'])

    # Create a mask
    getmask1 = pe.MapNode(
        fsl.ImageMaths(op_string="-bin", suffix="_mask"),
        name="getmask1",
        iterfield=['in_file'])
    getmask2 = pe.MapNode(
        fsl.MultiImageMaths(op_string="-dilD -dilD -sub 1 -abs -add %s"),
        name="getmask2",
        iterfield=['in_file', 'operand_files'])

    # $FSLDIR/bin/fslmaths FA/${f}_FA_mask -dilD -dilD -sub 1 -abs -add FA/${f}_FA_mask FA/${f}_FA_mask -odt char
    # Define the tbss1 workflow
    tbss1 = pe.Workflow(name=name)
    tbss1.connect([
        (inputnode, prepfa, [("fa_list", "in_file")]),
        (inputnode, prepfa, [(("fa_list", tbss1_op_string), "op_string")]),
        (prepfa, getmask1, [("out_file", "in_file")]),
        (getmask1, getmask2, [("out_file", "in_file"), ("out_file",
                                                        "operand_files")]),
        (prepfa, slicer, [('out_file', 'in_file')]),
    ])

    # Define the outputnode
    outputnode = pe.Node(
        interface=util.IdentityInterface(
            fields=["fa_list", "mask_list", "slices"]),
        name="outputnode")
    tbss1.connect([(prepfa, outputnode, [("out_file", "fa_list")]),
                   (getmask2, outputnode, [("out_file", "mask_list")]),
                   (slicer, outputnode, [('out_file', 'slices')])])
    return tbss1


def create_tbss_2_reg(name="tbss_2_reg"):
    """TBSS nonlinear registration:
    A pipeline that does the same as 'tbss_2_reg -t' script in FSL. '-n' option
    is not supported at the moment.

    Example
    -------

    >>> from nipype.workflows.dmri.fsl import tbss
    >>> tbss2 = create_tbss_2_reg(name="tbss2")
    >>> tbss2.inputs.inputnode.target = fsl.Info.standard_image("FMRIB58_FA_1mm.nii.gz")  # doctest: +SKIP
    >>> tbss2.inputs.inputnode.fa_list = ['s1_FA.nii', 's2_FA.nii', 's3_FA.nii']
    >>> tbss2.inputs.inputnode.mask_list = ['s1_mask.nii', 's2_mask.nii', 's3_mask.nii']

    Inputs::

        inputnode.fa_list
        inputnode.mask_list
        inputnode.target

    Outputs::

        outputnode.field_list

    """

    # Define the inputnode
    inputnode = pe.Node(
        interface=util.IdentityInterface(
            fields=["fa_list", "mask_list", "target"]),
        name="inputnode")

    # Flirt the FA image to the target
    flirt = pe.MapNode(
        interface=fsl.FLIRT(dof=12),
        iterfield=['in_file', 'in_weight'],
        name="flirt")

    fnirt = pe.MapNode(
        interface=fsl.FNIRT(fieldcoeff_file=True),
        iterfield=['in_file', 'inmask_file', 'affine_file'],
        name="fnirt")
    # Fnirt the FA image to the target
    if fsl.no_fsl():
        warn('NO FSL found')
    else:
        config_file = os.path.join(os.environ["FSLDIR"],
                                   "etc/flirtsch/FA_2_FMRIB58_1mm.cnf")
        fnirt.inputs.config_file = config_file

    # Define the registration workflow
    tbss2 = pe.Workflow(name=name)

    # Connect up the registration workflow
    tbss2.connect([
        (inputnode, flirt, [("fa_list", "in_file"), ("target", "reference"),
                            ("mask_list", "in_weight")]),
        (inputnode, fnirt, [("fa_list", "in_file"),
                            ("mask_list", "inmask_file"), ("target",
                                                           "ref_file")]),
        (flirt, fnirt, [("out_matrix_file", "affine_file")]),
    ])

    # Define the outputnode
    outputnode = pe.Node(
        interface=util.IdentityInterface(fields=['field_list']),
        name="outputnode")

    tbss2.connect([(fnirt, outputnode, [('fieldcoeff_file', 'field_list')])])
    return tbss2


def create_tbss_3_postreg(name='tbss_3_postreg', estimate_skeleton=True):
    """Post-registration processing: derive mean_FA and mean_FA_skeleton from
    mean of all subjects in study. Target is assumed to be FMRIB58_FA_1mm.
    A pipeline that does the same as 'tbss_3_postreg -S' script from FSL
    Setting 'estimate_skeleton to False will use precomputed FMRIB58_FA-skeleton_1mm
    skeleton (same as 'tbss_3_postreg -T').

    Example
    -------

    >>> from nipype.workflows.dmri.fsl import tbss
    >>> tbss3 = tbss.create_tbss_3_postreg()
    >>> tbss3.inputs.inputnode.fa_list = ['s1_wrapped_FA.nii', 's2_wrapped_FA.nii', 's3_wrapped_FA.nii']

    Inputs::

        inputnode.field_list
        inputnode.fa_list

    Outputs::

        outputnode.groupmask
        outputnode.skeleton_file
        outputnode.meanfa_file
        outputnode.mergefa_file

    """

    # Create the inputnode
    inputnode = pe.Node(
        interface=util.IdentityInterface(fields=['field_list', 'fa_list']),
        name='inputnode')

    # Apply the warpfield to the masked FA image
    applywarp = pe.MapNode(
        interface=fsl.ApplyWarp(),
        iterfield=['in_file', 'field_file'],
        name="applywarp")
    if fsl.no_fsl():
        warn('NO FSL found')
    else:
        applywarp.inputs.ref_file = fsl.Info.standard_image(
            "FMRIB58_FA_1mm.nii.gz")

    # Merge the FA files into a 4D file
    mergefa = pe.Node(fsl.Merge(dimension="t"), name="mergefa")

    # Get a group mask
    groupmask = pe.Node(
        fsl.ImageMaths(
            op_string="-max 0 -Tmin -bin",
            out_data_type="char",
            suffix="_mask"),
        name="groupmask")

    maskgroup = pe.Node(
        fsl.ImageMaths(op_string="-mas", suffix="_masked"), name="maskgroup")

    tbss3 = pe.Workflow(name=name)
    tbss3.connect([
        (inputnode, applywarp, [("fa_list", "in_file"), ("field_list",
                                                         "field_file")]),
        (applywarp, mergefa, [("out_file", "in_files")]),
        (mergefa, groupmask, [("merged_file", "in_file")]),
        (mergefa, maskgroup, [("merged_file", "in_file")]),
        (groupmask, maskgroup, [("out_file", "in_file2")]),
    ])

    # Create outputnode
    outputnode = pe.Node(
        interface=util.IdentityInterface(fields=[
            'groupmask', 'skeleton_file', 'meanfa_file', 'mergefa_file'
        ]),
        name='outputnode')

    if estimate_skeleton:
        # Take the mean over the fourth dimension
        meanfa = pe.Node(
            fsl.ImageMaths(op_string="-Tmean", suffix="_mean"), name="meanfa")

        # Use the mean FA volume to generate a tract skeleton
        makeskeleton = pe.Node(
            fsl.TractSkeleton(skeleton_file=True), name="makeskeleton")
        tbss3.connect(
            [(maskgroup, meanfa, [("out_file", "in_file")]),
             (meanfa, makeskeleton,
              [("out_file", "in_file")]), (groupmask, outputnode,
                                           [('out_file', 'groupmask')]),
             (makeskeleton, outputnode,
              [('skeleton_file', 'skeleton_file')]), (meanfa, outputnode, [
                  ('out_file', 'meanfa_file')
              ]), (maskgroup, outputnode, [('out_file', 'mergefa_file')])])
    else:
        # $FSLDIR/bin/fslmaths $FSLDIR/data/standard/FMRIB58_FA_1mm -mas mean_FA_mask mean_FA
        maskstd = pe.Node(
            fsl.ImageMaths(op_string="-mas", suffix="_masked"), name="maskstd")
        if fsl.no_fsl():
            warn('NO FSL found')
        else:
            maskstd.inputs.in_file = fsl.Info.standard_image(
                "FMRIB58_FA_1mm.nii.gz")

        # $FSLDIR/bin/fslmaths mean_FA -bin mean_FA_mask
        binmaskstd = pe.Node(
            fsl.ImageMaths(op_string="-bin"), name="binmaskstd")

        # $FSLDIR/bin/fslmaths all_FA -mas mean_FA_mask all_FA
        maskgroup2 = pe.Node(
            fsl.ImageMaths(op_string="-mas", suffix="_masked"),
            name="maskgroup2")

        tbss3.connect([(groupmask, maskstd, [("out_file", "in_file2")]),
                       (maskstd, binmaskstd, [("out_file", "in_file")]),
                       (maskgroup, maskgroup2, [("out_file", "in_file")]),
                       (binmaskstd, maskgroup2, [("out_file", "in_file2")])])

        if fsl.no_fsl():
            warn('NO FSL found')
        else:
            outputnode.inputs.skeleton_file = fsl.Info.standard_image(
                "FMRIB58_FA-skeleton_1mm.nii.gz")
        tbss3.connect([(binmaskstd, outputnode, [('out_file', 'groupmask')]),
                       (maskstd, outputnode, [('out_file', 'meanfa_file')]),
                       (maskgroup2, outputnode, [('out_file',
                                                  'mergefa_file')])])
    return tbss3


def tbss4_op_string(skeleton_thresh):
    op_string = "-thr %.1f -bin" % skeleton_thresh
    return op_string


def create_tbss_4_prestats(name='tbss_4_prestats'):
    """Post-registration processing:Creating skeleton mask using a threshold
     projecting all FA data onto skeleton.
    A pipeline that does the same as tbss_4_prestats script from FSL

    Example
    -------

    >>> from nipype.workflows.dmri.fsl import tbss
    >>> tbss4 = tbss.create_tbss_4_prestats(name='tbss4')
    >>> tbss4.inputs.inputnode.skeleton_thresh = 0.2

    Inputs::

        inputnode.skeleton_thresh
        inputnode.groupmask
        inputnode.skeleton_file
        inputnode.meanfa_file
        inputnode.mergefa_file

    Outputs::

        outputnode.all_FA_skeletonised
        outputnode.mean_FA_skeleton_mask
        outputnode.distance_map
        outputnode.skeleton_file

    """
    # Create inputnode
    inputnode = pe.Node(
        interface=util.IdentityInterface(fields=[
            'groupmask', 'skeleton_file', 'meanfa_file', 'mergefa_file',
            'skeleton_thresh'
        ]),
        name='inputnode')

    # Mask the skeleton at the threshold
    skeletonmask = pe.Node(fsl.ImageMaths(suffix="_mask"), name="skeletonmask")

    # Invert the brainmask then add in the tract skeleton
    invertmask = pe.Node(
        fsl.ImageMaths(suffix="_inv", op_string="-mul -1 -add 1 -add"),
        name="invertmask")

    # Generate a distance map with the tract skeleton
    distancemap = pe.Node(fsl.DistanceMap(), name="distancemap")

    # Project the FA values onto the skeleton
    projectfa = pe.Node(
        fsl.TractSkeleton(
            project_data=True, skeleton_file=True, use_cingulum_mask=True),
        name="projectfa")

    # Create tbss4 workflow
    tbss4 = pe.Workflow(name=name)
    tbss4.connect([
        (inputnode, invertmask, [("groupmask", "in_file")]),
        (inputnode, skeletonmask, [("skeleton_file", "in_file"),
                                   (('skeleton_thresh', tbss4_op_string),
                                    'op_string')]),
        (inputnode, projectfa, [('skeleton_thresh', 'threshold'),
                                ("meanfa_file", "in_file"), ("mergefa_file",
                                                             "data_file")]),
        (skeletonmask, invertmask, [("out_file", "in_file2")]),
        (invertmask, distancemap, [("out_file", "in_file")]),
        (distancemap, projectfa, [("distance_map", "distance_map")]),
    ])

    # Create the outputnode
    outputnode = pe.Node(
        interface=util.IdentityInterface(fields=[
            'projectedfa_file', 'skeleton_mask', 'distance_map',
            'skeleton_file'
        ]),
        name='outputnode')

    tbss4.connect(
        [(projectfa, outputnode, [('projected_data', 'projectedfa_file'),
                                  ('skeleton_file', 'skeleton_file')]),
         (distancemap, outputnode, [('distance_map', 'distance_map')]),
         (skeletonmask, outputnode, [('out_file', 'skeleton_mask')])])

    return tbss4


def create_tbss_all(name='tbss_all', estimate_skeleton=True):
    """Create a pipeline that combines create_tbss_* pipelines

    Example
    -------

    >>> from nipype.workflows.dmri.fsl import tbss
    >>> tbss_wf = tbss.create_tbss_all('tbss', estimate_skeleton=True)
    >>> tbss_wf.inputs.inputnode.skeleton_thresh = 0.2
    >>> tbss_wf.inputs.inputnode.fa_list = ['s1_wrapped_FA.nii', 's2_wrapped_FA.nii', 's3_wrapped_FA.nii']

    >>> tbss_wf = tbss.create_tbss_all('tbss', estimate_skeleton=False)
    >>> tbss_wf.inputs.inputnode.skeleton_thresh = 0.2
    >>> tbss_wf.inputs.inputnode.fa_list = ['s1_wrapped_FA.nii', 's2_wrapped_FA.nii', 's3_wrapped_FA.nii']


    Inputs::

        inputnode.fa_list
        inputnode.skeleton_thresh

    Outputs::

        outputnode.meanfa_file
        outputnode.projectedfa_file
        outputnode.skeleton_file
        outputnode.skeleton_mask

    """

    # Define the inputnode
    inputnode = pe.Node(
        interface=util.IdentityInterface(
            fields=['fa_list', 'skeleton_thresh']),
        name='inputnode')

    tbss1 = create_tbss_1_preproc(name='tbss1')
    tbss2 = create_tbss_2_reg(name='tbss2')
    if fsl.no_fsl():
        warn('NO FSL found')
    else:
        tbss2.inputs.inputnode.target = fsl.Info.standard_image(
            "FMRIB58_FA_1mm.nii.gz")
    tbss3 = create_tbss_3_postreg(
        name='tbss3', estimate_skeleton=estimate_skeleton)
    tbss4 = create_tbss_4_prestats(name='tbss4')

    tbss_all = pe.Workflow(name=name)
    tbss_all.connect(
        [(inputnode, tbss1, [('fa_list', 'inputnode.fa_list')]),
         (inputnode, tbss4,
          [('skeleton_thresh', 'inputnode.skeleton_thresh')]), (tbss1, tbss2, [
              ('outputnode.fa_list', 'inputnode.fa_list'),
              ('outputnode.mask_list', 'inputnode.mask_list')
          ]), (tbss1, tbss3, [('outputnode.fa_list', 'inputnode.fa_list')]),
         (tbss2, tbss3, [('outputnode.field_list', 'inputnode.field_list')]),
         (tbss3, tbss4, [('outputnode.groupmask',
                          'inputnode.groupmask'), ('outputnode.skeleton_file',
                                                   'inputnode.skeleton_file'),
                         ('outputnode.meanfa_file', 'inputnode.meanfa_file'),
                         ('outputnode.mergefa_file',
                          'inputnode.mergefa_file')])])

    # Define the outputnode
    outputnode = pe.Node(
        interface=util.IdentityInterface(fields=[
            'groupmask', 'skeleton_file3', 'meanfa_file', 'mergefa_file',
            'projectedfa_file', 'skeleton_file4', 'skeleton_mask',
            'distance_map'
        ]),
        name='outputnode')
    outputall_node = pe.Node(
        interface=util.IdentityInterface(fields=[
            'fa_list1', 'mask_list1', 'field_list2', 'groupmask3',
            'skeleton_file3', 'meanfa_file3', 'mergefa_file3',
            'projectedfa_file4', 'skeleton_mask4', 'distance_map4'
        ]),
        name='outputall_node')

    tbss_all.connect([
        (tbss3, outputnode, [
            ('outputnode.meanfa_file', 'meanfa_file'),
            ('outputnode.mergefa_file', 'mergefa_file'),
            ('outputnode.groupmask', 'groupmask'),
            ('outputnode.skeleton_file', 'skeleton_file3'),
        ]),
        (tbss4, outputnode, [
            ('outputnode.projectedfa_file', 'projectedfa_file'),
            ('outputnode.skeleton_file', 'skeleton_file4'),
            ('outputnode.skeleton_mask', 'skeleton_mask'),
            ('outputnode.distance_map', 'distance_map'),
        ]),
        (tbss1, outputall_node, [
            ('outputnode.fa_list', 'fa_list1'),
            ('outputnode.mask_list', 'mask_list1'),
        ]),
        (tbss2, outputall_node, [
            ('outputnode.field_list', 'field_list2'),
        ]),
        (tbss3, outputall_node, [
            ('outputnode.meanfa_file', 'meanfa_file3'),
            ('outputnode.mergefa_file', 'mergefa_file3'),
            ('outputnode.groupmask', 'groupmask3'),
            ('outputnode.skeleton_file', 'skeleton_file3'),
        ]),
        (tbss4, outputall_node, [
            ('outputnode.projectedfa_file', 'projectedfa_file4'),
            ('outputnode.skeleton_mask', 'skeleton_mask4'),
            ('outputnode.distance_map', 'distance_map4'),
        ]),
    ])
    return tbss_all


def create_tbss_non_FA(name='tbss_non_FA'):
    """
    A pipeline that implement tbss_non_FA in FSL

    Example
    -------

    >>> from nipype.workflows.dmri.fsl import tbss
    >>> tbss_MD = tbss.create_tbss_non_FA()
    >>> tbss_MD.inputs.inputnode.file_list = []
    >>> tbss_MD.inputs.inputnode.field_list = []
    >>> tbss_MD.inputs.inputnode.skeleton_thresh = 0.2
    >>> tbss_MD.inputs.inputnode.groupmask = './xxx'
    >>> tbss_MD.inputs.inputnode.meanfa_file = './xxx'
    >>> tbss_MD.inputs.inputnode.distance_map = []
    >>> tbss_MD.inputs.inputnode.all_FA_file = './xxx'

    Inputs::

        inputnode.file_list
        inputnode.field_list
        inputnode.skeleton_thresh
        inputnode.groupmask
        inputnode.meanfa_file
        inputnode.distance_map
        inputnode.all_FA_file

    Outputs::

        outputnode.projected_nonFA_file

    """

    # Define the inputnode
    inputnode = pe.Node(
        interface=util.IdentityInterface(fields=[
            'file_list', 'field_list', 'skeleton_thresh', 'groupmask',
            'meanfa_file', 'distance_map', 'all_FA_file'
        ]),
        name='inputnode')

    # Apply the warpfield to the non FA image
    applywarp = pe.MapNode(
        interface=fsl.ApplyWarp(),
        iterfield=['in_file', 'field_file'],
        name="applywarp")
    if fsl.no_fsl():
        warn('NO FSL found')
    else:
        applywarp.inputs.ref_file = fsl.Info.standard_image(
            "FMRIB58_FA_1mm.nii.gz")
    # Merge the non FA files into a 4D file
    merge = pe.Node(fsl.Merge(dimension="t"), name="merge")
    # merged_file="all_FA.nii.gz"
    maskgroup = pe.Node(
        fsl.ImageMaths(op_string="-mas", suffix="_masked"), name="maskgroup")
    projectfa = pe.Node(
        fsl.TractSkeleton(
            project_data=True,
            # projected_data = 'test.nii.gz',
            use_cingulum_mask=True),
        name="projectfa")

    tbss_non_FA = pe.Workflow(name=name)
    tbss_non_FA.connect([
        (inputnode, applywarp, [
            ('file_list', 'in_file'),
            ('field_list', 'field_file'),
        ]),
        (applywarp, merge, [("out_file", "in_files")]),
        (merge, maskgroup, [("merged_file", "in_file")]),
        (inputnode, maskgroup, [('groupmask', 'in_file2')]),
        (maskgroup, projectfa, [('out_file', 'alt_data_file')]),
        (inputnode, projectfa,
         [('skeleton_thresh', 'threshold'), ("meanfa_file", "in_file"),
          ("distance_map", "distance_map"), ("all_FA_file", 'data_file')]),
    ])

    # Define the outputnode
    outputnode = pe.Node(
        interface=util.IdentityInterface(fields=['projected_nonFA_file']),
        name='outputnode')
    tbss_non_FA.connect([
        (projectfa, outputnode, [
            ('projected_data', 'projected_nonFA_file'),
        ]),
    ])
    return tbss_non_FA
