# -*- 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 utility 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, File, traits, isdefined
from ...utils.filemanip import fname_presuffix
from .base import CommandLineDtitk, DTITKRenameMixin
import os

__docformat__ = "restructuredtext"


class TVAdjustVoxSpInputSpec(CommandLineInputSpec):
    in_file = File(
        desc="tensor volume to modify", exists=True, mandatory=True, argstr="-in %s"
    )
    out_file = File(
        desc="output path",
        argstr="-out %s",
        name_source="in_file",
        name_template="%s_avs",
        keep_extension=True,
    )
    target_file = File(
        desc="target volume to match", argstr="-target %s", xor=["voxel_size", "origin"]
    )
    voxel_size = traits.Tuple(
        (traits.Float(), traits.Float(), traits.Float()),
        desc="xyz voxel size (superseded by target)",
        argstr="-vsize %g %g %g",
        xor=["target_file"],
    )
    origin = traits.Tuple(
        (traits.Float(), traits.Float(), traits.Float()),
        desc="xyz origin (superseded by target)",
        argstr="-origin %g %g %g",
        xor=["target_file"],
    )


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


class TVAdjustVoxSp(CommandLineDtitk):
    """
     Adjusts the voxel space of a tensor volume.

    Example
    -------
    >>> from nipype.interfaces import dtitk
    >>> node = dtitk.TVAdjustVoxSp()
    >>> node.inputs.in_file = 'im1.nii'
    >>> node.inputs.target_file = 'im2.nii'
    >>> node.cmdline
    'TVAdjustVoxelspace -in im1.nii -out im1_avs.nii -target im2.nii'
    >>> node.run() # doctest: +SKIP

    """

    input_spec = TVAdjustVoxSpInputSpec
    output_spec = TVAdjustVoxSpOutputSpec
    _cmd = "TVAdjustVoxelspace"


class SVAdjustVoxSpInputSpec(CommandLineInputSpec):
    in_file = File(
        desc="scalar volume to modify", exists=True, mandatory=True, argstr="-in %s"
    )
    out_file = File(
        desc="output path",
        argstr="-out %s",
        name_source="in_file",
        name_template="%s_avs",
        keep_extension=True,
    )
    target_file = File(
        desc="target volume to match", argstr="-target %s", xor=["voxel_size", "origin"]
    )
    voxel_size = traits.Tuple(
        (traits.Float(), traits.Float(), traits.Float()),
        desc="xyz voxel size (superseded by target)",
        argstr="-vsize %g %g %g",
        xor=["target_file"],
    )
    origin = traits.Tuple(
        (traits.Float(), traits.Float(), traits.Float()),
        desc="xyz origin (superseded by target)",
        argstr="-origin %g %g %g",
        xor=["target_file"],
    )


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


class SVAdjustVoxSp(CommandLineDtitk):
    """
    Adjusts the voxel space of a scalar volume.

    Example
    -------
    >>> from nipype.interfaces import dtitk
    >>> node = dtitk.SVAdjustVoxSp()
    >>> node.inputs.in_file = 'im1.nii'
    >>> node.inputs.target_file = 'im2.nii'
    >>> node.cmdline
    'SVAdjustVoxelspace -in im1.nii -out im1_avs.nii -target im2.nii'
    >>> node.run() # doctest: +SKIP

    """

    input_spec = SVAdjustVoxSpInputSpec
    output_spec = SVAdjustVoxSpOutputSpec
    _cmd = "SVAdjustVoxelspace"


class TVResampleInputSpec(CommandLineInputSpec):
    in_file = File(
        desc="tensor volume to resample", exists=True, mandatory=True, argstr="-in %s"
    )
    out_file = File(
        desc="output path",
        name_source="in_file",
        name_template="%s_resampled",
        keep_extension=True,
        argstr="-out %s",
    )
    target_file = File(
        desc="specs read from the target volume",
        argstr="-target %s",
        xor=["array_size", "voxel_size", "origin"],
    )
    align = traits.Enum(
        "center",
        "origin",
        argstr="-align %s",
        desc="how to align output volume to input volume",
    )
    interpolation = traits.Enum(
        "LEI", "EI", argstr="-interp %s", desc="Log Euclidean Euclidean Interpolation"
    )
    array_size = traits.Tuple(
        (traits.Int(), traits.Int(), traits.Int()),
        desc="resampled array size",
        xor=["target_file"],
        argstr="-size %d %d %d",
    )
    voxel_size = traits.Tuple(
        (traits.Float(), traits.Float(), traits.Float()),
        desc="resampled voxel size",
        xor=["target_file"],
        argstr="-vsize %g %g %g",
    )
    origin = traits.Tuple(
        (traits.Float(), traits.Float(), traits.Float()),
        desc="xyz origin",
        xor=["target_file"],
        argstr="-origin %g %g %g",
    )


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


class TVResample(CommandLineDtitk):
    """
    Resamples a tensor volume.

    Example
    -------
    >>> from nipype.interfaces import dtitk
    >>> node = dtitk.TVResample()
    >>> node.inputs.in_file = 'im1.nii'
    >>> node.inputs.target_file = 'im2.nii'
    >>> node.cmdline
    'TVResample -in im1.nii -out im1_resampled.nii -target im2.nii'
    >>> node.run() # doctest: +SKIP

    """

    input_spec = TVResampleInputSpec
    output_spec = TVResampleOutputSpec
    _cmd = "TVResample"


class SVResampleInputSpec(CommandLineInputSpec):
    in_file = File(
        desc="image to resample", exists=True, mandatory=True, argstr="-in %s"
    )
    out_file = File(
        desc="output path",
        name_source="in_file",
        name_template="%s_resampled",
        keep_extension=True,
        argstr="-out %s",
    )
    target_file = File(
        desc="specs read from the target volume",
        argstr="-target %s",
        xor=["array_size", "voxel_size", "origin"],
    )
    align = traits.Enum(
        "center",
        "origin",
        argstr="-align %s",
        desc="how to align output volume to input volume",
    )
    array_size = traits.Tuple(
        (traits.Int(), traits.Int(), traits.Int()),
        desc="resampled array size",
        xor=["target_file"],
        argstr="-size %d %d %d",
    )
    voxel_size = traits.Tuple(
        (traits.Float(), traits.Float(), traits.Float()),
        desc="resampled voxel size",
        xor=["target_file"],
        argstr="-vsize %g %g %g",
    )
    origin = traits.Tuple(
        (traits.Float(), traits.Float(), traits.Float()),
        desc="xyz origin",
        xor=["target_file"],
        argstr="-origin %g %g %g",
    )


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


class SVResample(CommandLineDtitk):
    """
    Resamples a scalar volume.

    Example
    -------
    >>> from nipype.interfaces import dtitk
    >>> node = dtitk.SVResample()
    >>> node.inputs.in_file = 'im1.nii'
    >>> node.inputs.target_file = 'im2.nii'
    >>> node.cmdline
    'SVResample -in im1.nii -out im1_resampled.nii -target im2.nii'
    >>> node.run() # doctest: +SKIP

    """

    input_spec = SVResampleInputSpec
    output_spec = SVResampleOutputSpec
    _cmd = "SVResample"


class TVtoolInputSpec(CommandLineInputSpec):
    in_file = File(
        desc="scalar volume to resample", exists=True, argstr="-in %s", mandatory=True
    )
    """NOTE: there are a lot more options here; not implementing all of them"""
    in_flag = traits.Enum("fa", "tr", "ad", "rd", "pd", "rgb", argstr="-%s", desc="")
    out_file = File(argstr="-out %s", genfile=True)


class TVtoolOutputSpec(TraitedSpec):
    out_file = File()


class TVtool(CommandLineDtitk):
    """
    Calculates a tensor metric volume from a tensor volume.

    Example
    -------
    >>> from nipype.interfaces import dtitk
    >>> node = dtitk.TVtool()
    >>> node.inputs.in_file = 'im1.nii'
    >>> node.inputs.in_flag = 'fa'
    >>> node.cmdline
    'TVtool -in im1.nii -fa -out im1_fa.nii'
    >>> node.run() # doctest: +SKIP

    """

    input_spec = TVtoolInputSpec
    output_spec = TVtoolOutputSpec
    _cmd = "TVtool"

    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
        return fname_presuffix(
            os.path.basename(self.inputs.in_file), suffix="_" + self.inputs.in_flag
        )


"""Note: SVTool not implemented at this time"""


class BinThreshInputSpec(CommandLineInputSpec):
    in_file = File(
        desc="Image to threshold/binarize",
        exists=True,
        position=0,
        argstr="%s",
        mandatory=True,
    )
    out_file = File(
        desc="output path",
        position=1,
        argstr="%s",
        keep_extension=True,
        name_source="in_file",
        name_template="%s_thrbin",
    )
    lower_bound = traits.Float(
        0.01,
        usedefault=True,
        position=2,
        argstr="%g",
        mandatory=True,
        desc="lower bound of binarization range",
    )
    upper_bound = traits.Float(
        100,
        usedefault=True,
        position=3,
        argstr="%g",
        mandatory=True,
        desc="upper bound of binarization range",
    )
    inside_value = traits.Float(
        1,
        position=4,
        argstr="%g",
        usedefault=True,
        mandatory=True,
        desc="value for voxels in " "binarization range",
    )
    outside_value = traits.Float(
        0,
        position=5,
        argstr="%g",
        usedefault=True,
        mandatory=True,
        desc="value for voxels" "outside of binarization range",
    )


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


class BinThresh(CommandLineDtitk):
    """
    Binarizes an image.

    Example
    -------
    >>> from nipype.interfaces import dtitk
    >>> node = dtitk.BinThresh()
    >>> node.inputs.in_file = 'im1.nii'
    >>> node.inputs.lower_bound = 0
    >>> node.inputs.upper_bound = 100
    >>> node.inputs.inside_value = 1
    >>> node.inputs.outside_value = 0
    >>> node.cmdline
    'BinaryThresholdImageFilter im1.nii im1_thrbin.nii 0 100 1 0'
    >>> node.run() # doctest: +SKIP

    """

    input_spec = BinThreshInputSpec
    output_spec = BinThreshOutputSpec
    _cmd = "BinaryThresholdImageFilter"


class BinThreshTask(DTITKRenameMixin, BinThresh):
    pass


class SVAdjustVoxSpTask(DTITKRenameMixin, SVAdjustVoxSp):
    pass


class SVResampleTask(DTITKRenameMixin, SVResample):
    pass


class TVAdjustOriginTask(DTITKRenameMixin, TVAdjustVoxSp):
    pass


class TVAdjustVoxSpTask(DTITKRenameMixin, TVAdjustVoxSp):
    pass


class TVResampleTask(DTITKRenameMixin, TVResample):
    pass


class TVtoolTask(DTITKRenameMixin, TVtool):
    pass