# -*- 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:
"""DTITK registration interfaces

DTI-TK developed by Gary Hui Zhang, gary.zhang@ucl.ac.uk
For additional help, visit http://dti-tk.sf.net

The high-dimensional tensor-based DTI registration algorithm

Zhang, H., Avants, B.B, Yushkevich, P.A., Woo, J.H., Wang, S., McCluskey, L.H.,
 Elman, L.B., Melhem, E.R., Gee, J.C., High-dimensional spatial normalization
 of diffusion tensor images improves the detection of white matter differences
 in amyotrophic lateral sclerosis, IEEE Transactions on Medical Imaging,
 26(11):1585-1597, November 2007. PMID: 18041273.

The original piecewise-affine tensor-based DTI registration algorithm at the
core of DTI-TK

Zhang, H., Yushkevich, P.A., Alexander, D.C., Gee, J.C., Deformable
 registration of diffusion tensor MR images with explicit orientation
 optimization, Medical Image Analysis, 10(5):764-785, October 2006. PMID:
 16899392.

"""

from ..base import TraitedSpec, CommandLineInputSpec, traits, File, isdefined
from ...utils.filemanip import fname_presuffix, split_filename
from .base import CommandLineDtitk, DTITKRenameMixin
import os

__docformat__ = 'restructuredtext'


class RigidInputSpec(CommandLineInputSpec):
    fixed_file = File(desc="fixed tensor volume", exists=True,
                      mandatory=True, position=0, argstr="%s", copyfile=False)
    moving_file = File(desc="moving tensor volume", exists=True,
                       mandatory=True, position=1, argstr="%s", copyfile=False)
    similarity_metric = traits.Enum('EDS', 'GDS', 'DDS', 'NMI',
                                    mandatory=True, position=2, argstr="%s",
                                    desc="similarity metric", usedefault=True)
    sampling_xyz = traits.Tuple((4, 4, 4), mandatory=True, position=3,
                                argstr="%g %g %g", usedefault=True,
                                desc="dist between samp points (mm) (x,y,z)")
    ftol = traits.Float(mandatory=True, position=4, argstr="%g",
                        desc="cost function tolerance", default_value=0.01,
                        usedefault=True)
    initialize_xfm = File(copyfile=True, desc="Initialize w/DTITK-FORMAT"
                          "affine", position=5, argstr="%s", exists=True)


class RigidOutputSpec(TraitedSpec):
    out_file = File(exists=True)
    out_file_xfm = File(exists=True)


class Rigid(CommandLineDtitk):
    """Performs rigid registration between two tensor volumes

    Example
    -------

    >>> from nipype.interfaces import dtitk
    >>> node = dtitk.Rigid()
    >>> node.inputs.fixed_file = 'im1.nii'
    >>> node.inputs.moving_file = 'im2.nii'
    >>> node.inputs.similarity_metric = 'EDS'
    >>> node.inputs.sampling_xyz = (4,4,4)
    >>> node.inputs.ftol = 0.01
    >>> node.cmdline
    'dti_rigid_reg im1.nii im2.nii EDS 4 4 4 0.01'
    >>> node.run() # doctest: +SKIP
    """
    input_spec = RigidInputSpec
    output_spec = RigidOutputSpec
    _cmd = 'dti_rigid_reg'

    '''def _format_arg(self, name, spec, value):
        if name == 'initialize_xfm':
            value = 1
        return super(Rigid, self)._format_arg(name, spec, value)'''

    def _run_interface(self, runtime):
        runtime = super(Rigid, self)._run_interface(runtime)
        if '''.aff doesn't exist or can't be opened''' in runtime.stderr:
            self.raise_exception(runtime)
        return runtime

    def _list_outputs(self):
        outputs = self.output_spec().get()
        moving = self.inputs.moving_file
        outputs['out_file_xfm'] = fname_presuffix(moving, suffix='.aff',
                                                  use_ext=False)
        outputs['out_file'] = fname_presuffix(moving, suffix='_aff')
        return outputs


class Affine(Rigid):
    """Performs affine registration between two tensor volumes

    Example
    -------

    >>> from nipype.interfaces import dtitk
    >>> node = dtitk.Affine()
    >>> node.inputs.fixed_file = 'im1.nii'
    >>> node.inputs.moving_file = 'im2.nii'
    >>> node.inputs.similarity_metric = 'EDS'
    >>> node.inputs.sampling_xyz = (4,4,4)
    >>> node.inputs.ftol = 0.01
    >>> node.inputs.initialize_xfm = 'im_affine.aff'
    >>> node.cmdline
    'dti_affine_reg im1.nii im2.nii EDS 4 4 4 0.01 im_affine.aff'
    >>> node.run() # doctest: +SKIP
    """
    _cmd = 'dti_affine_reg'


class DiffeoInputSpec(CommandLineInputSpec):
    fixed_file = File(desc="fixed tensor volume",
                      exists=True, position=0, argstr="%s")
    moving_file = File(desc="moving tensor volume",
                       exists=True, position=1, argstr="%s", copyfile=False)
    mask_file = File(desc="mask", exists=True,  position=2, argstr="%s")
    legacy = traits.Enum(1, desc="legacy parameter; always set to 1",
                         usedefault=True, mandatory=True,
                         position=3, argstr="%d")
    n_iters = traits.Int(6, desc="number of iterations",
                         mandatory=True,
                         position=4, argstr="%d", usedefault=True)
    ftol = traits.Float(0.002, desc="iteration for the optimization to stop",
                        mandatory=True, position=5, argstr="%g",
                        usedefault=True)


class DiffeoOutputSpec(TraitedSpec):
    out_file = File(exists=True)
    out_file_xfm = File(exists=True)


class Diffeo(CommandLineDtitk):
    """Performs diffeomorphic registration between two tensor volumes

    Example
    -------

    >>> from nipype.interfaces import dtitk
    >>> node = dtitk.Diffeo()
    >>> node.inputs.fixed_file = 'im1.nii'
    >>> node.inputs.moving_file = 'im2.nii'
    >>> node.inputs.mask_file = 'mask.nii'
    >>> node.inputs.legacy = 1
    >>> node.inputs.n_iters = 6
    >>> node.inputs.ftol = 0.002
    >>> node.cmdline
    'dti_diffeomorphic_reg im1.nii im2.nii mask.nii 1 6 0.002'
    >>> node.run() # doctest: +SKIP
    """
    input_spec = DiffeoInputSpec
    output_spec = DiffeoOutputSpec
    _cmd = 'dti_diffeomorphic_reg'

    def _list_outputs(self):
        outputs = self.output_spec().get()
        moving = self.inputs.moving_file
        outputs['out_file_xfm'] = fname_presuffix(moving, suffix='_diffeo.df')
        outputs['out_file'] = fname_presuffix(moving, suffix='_diffeo')
        return outputs


class ComposeXfmInputSpec(CommandLineInputSpec):
    in_df = File(desc='diffeomorphic warp file', exists=True,
                 argstr="-df %s", mandatory=True)
    in_aff = File(desc='affine transform file', exists=True,
                  argstr="-aff %s", mandatory=True)
    out_file = File(desc='output path',
                    argstr="-out %s",  genfile=True)


class ComposeXfmOutputSpec(TraitedSpec):
    out_file = File(exists=True)


class ComposeXfm(CommandLineDtitk):
    """
     Combines diffeomorphic and affine transforms

    Example
    -------

    >>> from nipype.interfaces import dtitk
    >>> node = dtitk.ComposeXfm()
    >>> node.inputs.in_df = 'im_warp.df.nii'
    >>> node.inputs.in_aff= 'im_affine.aff'
    >>> node.cmdline
    'dfRightComposeAffine -aff im_affine.aff -df im_warp.df.nii -out
     im_warp_affdf.df.nii'
    >>> node.run() # doctest: +SKIP
    """
    input_spec = ComposeXfmInputSpec
    output_spec = ComposeXfmOutputSpec
    _cmd = 'dfRightComposeAffine'

    def _list_outputs(self):
        outputs = self._outputs().get()
        out_file = self.inputs.out_file
        if not isdefined(out_file):
            out_file = self._gen_filename('out_file')
        outputs['out_file'] = os.path.abspath(out_file)
        return outputs

    def _gen_filename(self, name):
        if name != 'out_file':
            return
        path, base, ext = split_filename(self.inputs.in_df)
        suffix = '_affdf'
        if base.endswith('.df'):
            suffix += '.df'
            base = base[:-3]
        return fname_presuffix(base, suffix=suffix + ext, use_ext=False)


class AffSymTensor3DVolInputSpec(CommandLineInputSpec):
    in_file = File(desc='moving tensor volume', exists=True,
                   argstr="-in %s", mandatory=True)
    out_file = File(desc='output filename',
                    argstr="-out %s", name_source="in_file",
                    name_template="%s_affxfmd", keep_extension=True)
    transform = File(exists=True, argstr="-trans %s",
                     xor=['target', 'translation', 'euler', 'deformation'],
                     desc='transform to apply: specify an input transformation'
                     ' file; parameters input will be ignored',)
    interpolation = traits.Enum('LEI', 'EI', usedefault=True,
                                argstr="-interp %s",
                                desc='Log Euclidean/Euclidean Interpolation')
    reorient = traits.Enum('PPD', 'NO', 'FS', argstr='-reorient %s',
                           usedefault=True, desc='Reorientation strategy: '
                           'preservation of principal direction, no '
                           'reorientation, or finite strain')
    target = File(exists=True, argstr="-target %s", xor=['transform'],
                  desc='output volume specification read from the target '
                  'volume if specified')
    translation = traits.Tuple((traits.Float(), traits.Float(),
                                traits.Float()),
                               desc='translation (x,y,z) in mm',
                               argstr='-translation %g %g %g',
                               xor=['transform'])
    euler = traits.Tuple((traits.Float(), traits.Float(), traits.Float()),
                         desc='(theta, phi, psi) in degrees',
                         xor=['transform'], argstr='-euler %g %g %g')
    deformation = traits.Tuple((traits.Float(),) * 6,
                               desc='(xx,yy,zz,xy,yz,xz)', xor=['transform'],
                               argstr='-deformation %g %g %g %g %g %g')


class AffSymTensor3DVolOutputSpec(TraitedSpec):
    out_file = File(exists=True)


class AffSymTensor3DVol(CommandLineDtitk):
    """
    Applies affine transform to a tensor volume

    Example
    -------

    >>> from nipype.interfaces import dtitk
    >>> node = dtitk.AffSymTensor3DVol()
    >>> node.inputs.in_file = 'im1.nii'
    >>> node.inputs.transform = 'im_affine.aff'
    >>> node.cmdline
    'affineSymTensor3DVolume -in im1.nii -interp LEI -out im1_affxfmd.nii
     -reorient PPD -trans im_affine.aff'
    >>> node.run() # doctest: +SKIP
    """
    input_spec = AffSymTensor3DVolInputSpec
    output_spec = AffSymTensor3DVolOutputSpec
    _cmd = 'affineSymTensor3DVolume'


class AffScalarVolInputSpec(CommandLineInputSpec):
    in_file = File(desc='moving scalar volume', exists=True,
                   argstr="-in %s", mandatory=True)
    out_file = File(desc='output filename',
                    argstr="-out %s", name_source="in_file",
                    name_template="%s_affxfmd", keep_extension=True)
    transform = File(exists=True, argstr="-trans %s",
                     xor=['target', 'translation', 'euler', 'deformation'],
                     desc='transform to apply: specify an input transformation'
                     ' file; parameters input will be ignored',)
    interpolation = traits.Enum('trilinear', 'NN',
                                usedefault=True, argstr="-interp %s",
                                desc='trilinear or nearest neighbor'
                                ' interpolation')
    target = File(exists=True, argstr="-target %s", xor=['transform'],
                  desc='output volume specification read from the target '
                  'volume if specified')
    translation = traits.Tuple((traits.Float(), traits.Float(),
                                traits.Float()),
                               desc='translation (x,y,z) in mm',
                               argstr='-translation %g %g %g',
                               xor=['transform'])
    euler = traits.Tuple((traits.Float(), traits.Float(), traits.Float()),
                         desc='(theta, phi, psi) in degrees',
                         xor=['transform'], argstr='-euler %g %g %g')
    deformation = traits.Tuple((traits.Float(),) * 6,
                               desc='(xx,yy,zz,xy,yz,xz)', xor=['transform'],
                               argstr='-deformation %g %g %g %g %g %g')


class AffScalarVolOutputSpec(TraitedSpec):
    out_file = File(desc='moved volume', exists=True)


class AffScalarVol(CommandLineDtitk):
    """
    Applies affine transform to a scalar volume

    Example
    -------

    >>> from nipype.interfaces import dtitk
    >>> node = dtitk.AffScalarVol()
    >>> node.inputs.in_file = 'im1.nii'
    >>> node.inputs.transform = 'im_affine.aff'
    >>> node.cmdline
    'affineScalarVolume -in im1.nii -interp 0 -out im1_affxfmd.nii -trans
     im_affine.aff'
    >>> node.run() # doctest: +SKIP
    """
    input_spec = AffScalarVolInputSpec
    output_spec = AffScalarVolOutputSpec
    _cmd = 'affineScalarVolume'

    def _format_arg(self, name, spec, value):
        if name == 'interpolation':
            value = {'trilinear': 0, 'NN': 1}[value]
        return super(AffScalarVol, self)._format_arg(name, spec, value)


class DiffeoSymTensor3DVolInputSpec(CommandLineInputSpec):
    in_file = File(desc='moving tensor volume', exists=True,
                   argstr="-in %s", mandatory=True)
    out_file = File(desc='output filename',
                    argstr="-out %s", name_source="in_file",
                    name_template="%s_diffeoxfmd", keep_extension=True)
    transform = File(exists=True, argstr="-trans %s",
                     mandatory=True, desc='transform to apply')
    df = traits.Str('FD', argstr="-df %s", usedefault=True)
    interpolation = traits.Enum('LEI', 'EI', usedefault=True,
                                argstr="-interp %s",
                                desc='Log Euclidean/Euclidean Interpolation')
    reorient = traits.Enum('PPD', 'FS', argstr='-reorient %s',
                           usedefault=True, desc='Reorientation strategy: '
                           'preservation of principal direction or finite '
                           'strain')
    target = File(exists=True, argstr="-target %s", xor=['voxel_size'],
                  desc='output volume specification read from the target '
                  'volume if specified')
    voxel_size = traits.Tuple((traits.Float(), traits.Float(), traits.Float()),
                              desc='xyz voxel size (superseded by target)',
                              argstr="-vsize %g %g %g", xor=['target'])
    flip = traits.Tuple((traits.Int(), traits.Int(), traits.Int()),
                        argstr="-flip %d %d %d")
    resampling_type = traits.Enum('backward', 'forward',
                                  desc='use backward or forward resampling',
                                  argstr="-type %s")


class DiffeoSymTensor3DVolOutputSpec(TraitedSpec):
    out_file = File(exists=True)


class DiffeoSymTensor3DVol(CommandLineDtitk):
    """
    Applies diffeomorphic transform to a tensor volume

    Example
    -------

    >>> from nipype.interfaces import dtitk
    >>> node = dtitk.DiffeoSymTensor3DVol()
    >>> node.inputs.in_file = 'im1.nii'
    >>> node.inputs.transform = 'im_warp.df.nii'
    >>> node.cmdline
    'deformationSymTensor3DVolume -df FD -in im1.nii -interp LEI -out
     im1_diffeoxfmd.nii -reorient PPD -trans im_warp.df.nii'
    >>> node.run() # doctest: +SKIP
    """

    input_spec = DiffeoSymTensor3DVolInputSpec
    output_spec = DiffeoSymTensor3DVolOutputSpec
    _cmd = 'deformationSymTensor3DVolume'

    def _format_arg(self, name, spec, value):
        if name == 'resampling_type':
            value = {'forward': 0, 'backward': 1}[value]
        return super(DiffeoSymTensor3DVol, self)._format_arg(name, spec, value)


class DiffeoScalarVolInputSpec(CommandLineInputSpec):
    in_file = File(desc='moving scalar volume', exists=True,
                   argstr="-in %s", mandatory=True)
    out_file = File(desc='output filename',
                    argstr="-out %s", name_source="in_file",
                    name_template="%s_diffeoxfmd", keep_extension=True)
    transform = File(exists=True, argstr="-trans %s",
                     mandatory=True, desc='transform to apply')
    target = File(exists=True, argstr="-target %s", xor=['voxel_size'],
                  desc='output volume specification read from the target '
                  'volume if specified')
    voxel_size = traits.Tuple((traits.Float(), traits.Float(), traits.Float()),
                              desc='xyz voxel size (superseded by target)',
                              argstr="-vsize %g %g %g", xor=['target'])
    flip = traits.Tuple((traits.Int(), traits.Int(), traits.Int()),
                        argstr="-flip %d %d %d")
    resampling_type = traits.Enum('backward', 'forward',
                                  desc='use backward or forward resampling',
                                  argstr="-type %s")
    interpolation = traits.Enum('trilinear', 'NN',
                                desc='trilinear, or nearest neighbor',
                                argstr="-interp %s",
                                usedefault=True)


class DiffeoScalarVolOutputSpec(TraitedSpec):
    out_file = File(desc='moved volume', exists=True)


class DiffeoScalarVol(CommandLineDtitk):
    """
    Applies diffeomorphic transform to a scalar volume

    Example
    -------

    >>> from nipype.interfaces import dtitk
    >>> node = dtitk.DiffeoScalarVol()
    >>> node.inputs.in_file = 'im1.nii'
    >>> node.inputs.transform = 'im_warp.df.nii'
    >>> node.cmdline
    'deformationScalarVolume -in im1.nii -interp 0 -out im1_diffeoxfmd.nii
     -trans im_warp.df.nii'
    >>> node.run() # doctest: +SKIP
    """

    input_spec = DiffeoScalarVolInputSpec
    output_spec = DiffeoScalarVolOutputSpec
    _cmd = 'deformationScalarVolume'

    def _format_arg(self, name, spec, value):
        if name == 'resampling_type':
            value = {'forward': 0, 'backward': 1}[value]
        elif name == 'interpolation':
            value = {'trilinear': 0, 'NN': 1}[value]
        return super(DiffeoScalarVol, self)._format_arg(name, spec, value)


class RigidTask(DTITKRenameMixin, Rigid):
    pass


class AffineTask(DTITKRenameMixin, Affine):
    pass


class DiffeoTask(DTITKRenameMixin, Diffeo):
    pass


class ComposeXfmTask(DTITKRenameMixin, ComposeXfm):
    pass


class affScalarVolTask(DTITKRenameMixin, AffScalarVol):
    pass


class affSymTensor3DVolTask(DTITKRenameMixin, AffSymTensor3DVol):
    pass


class diffeoScalarVolTask(DTITKRenameMixin, DiffeoScalarVol):
    pass


class diffeoSymTensor3DVolTask(DTITKRenameMixin, DiffeoSymTensor3DVol):
    pass
