# -*- coding: utf-8 -*-
"""ANTS Apply Transforms interface
"""
from __future__ import (print_function, division, unicode_literals,
                        absolute_import)
from builtins import range, str
import os

from .base import ANTSCommand, ANTSCommandInputSpec
from ..base import TraitedSpec, File, traits, isdefined, InputMultiPath
from ...utils.filemanip import split_filename


class WarpTimeSeriesImageMultiTransformInputSpec(ANTSCommandInputSpec):
    dimension = traits.Enum(
        4,
        3,
        argstr='%d',
        usedefault=True,
        desc='image dimension (3 or 4)',
        position=1)
    input_image = File(
        argstr='%s',
        mandatory=True,
        copyfile=True,
        desc=('image to apply transformation to (generally a '
              'coregistered functional)'))
    out_postfix = traits.Str(
        '_wtsimt',
        argstr='%s',
        usedefault=True,
        desc=('Postfix that is prepended to all output '
              'files (default = _wtsimt)'))
    reference_image = File(
        argstr='-R %s',
        xor=['tightest_box'],
        desc='reference image space that you wish to warp INTO')
    tightest_box = traits.Bool(
        argstr='--tightest-bounding-box',
        desc=('computes tightest bounding box (overrided by '
              'reference_image if given)'),
        xor=['reference_image'])
    reslice_by_header = traits.Bool(
        argstr='--reslice-by-header',
        desc=('Uses orientation matrix and origin encoded in '
              'reference image file header. Not typically used '
              'with additional transforms'))
    use_nearest = traits.Bool(
        argstr='--use-NN', desc='Use nearest neighbor interpolation')
    use_bspline = traits.Bool(
        argstr='--use-Bspline', desc='Use 3rd order B-Spline interpolation')
    transformation_series = InputMultiPath(
        File(exists=True),
        argstr='%s',
        desc='transformation file(s) to be applied',
        mandatory=True,
        copyfile=False)
    invert_affine = traits.List(
        traits.Int,
        desc=(
            'List of Affine transformations to invert.'
            'E.g.: [1,4,5] inverts the 1st, 4th, and 5th Affines '
            'found in transformation_series. Note that indexing '
            'starts with 1 and does not include warp fields. Affine '
            'transformations are distinguished '
            'from warp fields by the word "affine" included in their filenames.'
        ))


class WarpTimeSeriesImageMultiTransformOutputSpec(TraitedSpec):
    output_image = File(exists=True, desc='Warped image')


class WarpTimeSeriesImageMultiTransform(ANTSCommand):
    """Warps a time-series from one space to another

    Examples
    --------

    >>> from nipype.interfaces.ants import WarpTimeSeriesImageMultiTransform
    >>> wtsimt = WarpTimeSeriesImageMultiTransform()
    >>> wtsimt.inputs.input_image = 'resting.nii'
    >>> wtsimt.inputs.reference_image = 'ants_deformed.nii.gz'
    >>> wtsimt.inputs.transformation_series = ['ants_Warp.nii.gz','ants_Affine.txt']
    >>> wtsimt.cmdline
    'WarpTimeSeriesImageMultiTransform 4 resting.nii resting_wtsimt.nii -R ants_deformed.nii.gz ants_Warp.nii.gz \
ants_Affine.txt'

    >>> wtsimt = WarpTimeSeriesImageMultiTransform()
    >>> wtsimt.inputs.input_image = 'resting.nii'
    >>> wtsimt.inputs.reference_image = 'ants_deformed.nii.gz'
    >>> wtsimt.inputs.transformation_series = ['ants_Warp.nii.gz','ants_Affine.txt']
    >>> wtsimt.inputs.invert_affine = [1] # # this will invert the 1st Affine file: ants_Affine.txt
    >>> wtsimt.cmdline
    'WarpTimeSeriesImageMultiTransform 4 resting.nii resting_wtsimt.nii -R ants_deformed.nii.gz ants_Warp.nii.gz \
-i ants_Affine.txt'
    """

    _cmd = 'WarpTimeSeriesImageMultiTransform'
    input_spec = WarpTimeSeriesImageMultiTransformInputSpec
    output_spec = WarpTimeSeriesImageMultiTransformOutputSpec

    def _format_arg(self, opt, spec, val):
        if opt == 'out_postfix':
            _, name, ext = split_filename(
                os.path.abspath(self.inputs.input_image))
            return name + val + ext
        if opt == 'transformation_series':
            series = []
            affine_counter = 0
            affine_invert = []
            for transformation in val:
                if 'Affine' in transformation and \
                        isdefined(self.inputs.invert_affine):
                    affine_counter += 1
                    if affine_counter in self.inputs.invert_affine:
                        series += ['-i']
                        affine_invert.append(affine_counter)
                series += [transformation]

            if isdefined(self.inputs.invert_affine):
                diff_inv = set(self.inputs.invert_affine) - set(affine_invert)
                if diff_inv:
                    raise Exceptions(
                        "Review invert_affine, not all indexes from invert_affine were used, "
                        "check the description for the full definition")

            return ' '.join(series)
        return super(WarpTimeSeriesImageMultiTransform, self)._format_arg(
            opt, spec, val)

    def _list_outputs(self):
        outputs = self._outputs().get()
        _, name, ext = split_filename(os.path.abspath(self.inputs.input_image))
        outputs['output_image'] = os.path.join(os.getcwd(), ''.join(
            (name, self.inputs.out_postfix, ext)))
        return outputs

    def _run_interface(self, runtime, correct_return_codes=[0]):
        runtime = super(WarpTimeSeriesImageMultiTransform,
                        self)._run_interface(
                            runtime, correct_return_codes=[0, 1])
        if "100 % complete" not in runtime.stdout:
            self.raise_exception(runtime)
        return runtime


class WarpImageMultiTransformInputSpec(ANTSCommandInputSpec):
    dimension = traits.Enum(
        3,
        2,
        argstr='%d',
        usedefault=True,
        desc='image dimension (2 or 3)',
        position=1)
    input_image = File(
        argstr='%s',
        mandatory=True,
        desc=('image to apply transformation to (generally a '
              'coregistered functional)'),
        position=2)
    output_image = File(
        genfile=True,
        hash_files=False,
        argstr='%s',
        desc='name of the output warped image',
        position=3,
        xor=['out_postfix'])
    out_postfix = File(
        "_wimt",
        usedefault=True,
        hash_files=False,
        desc=('Postfix that is prepended to all output '
              'files (default = _wimt)'),
        xor=['output_image'])
    reference_image = File(
        argstr='-R %s',
        xor=['tightest_box'],
        desc='reference image space that you wish to warp INTO')
    tightest_box = traits.Bool(
        argstr='--tightest-bounding-box',
        desc=('computes tightest bounding box (overrided by '
              'reference_image if given)'),
        xor=['reference_image'])
    reslice_by_header = traits.Bool(
        argstr='--reslice-by-header',
        desc=('Uses orientation matrix and origin encoded in '
              'reference image file header. Not typically used '
              'with additional transforms'))
    use_nearest = traits.Bool(
        argstr='--use-NN', desc='Use nearest neighbor interpolation')
    use_bspline = traits.Bool(
        argstr='--use-BSpline', desc='Use 3rd order B-Spline interpolation')
    transformation_series = InputMultiPath(
        File(exists=True),
        argstr='%s',
        desc='transformation file(s) to be applied',
        mandatory=True,
        position=-1)
    invert_affine = traits.List(
        traits.Int,
        desc=(
            'List of Affine transformations to invert.'
            'E.g.: [1,4,5] inverts the 1st, 4th, and 5th Affines '
            'found in transformation_series. Note that indexing '
            'starts with 1 and does not include warp fields. Affine '
            'transformations are distinguished '
            'from warp fields by the word "affine" included in their filenames.'
        ))


class WarpImageMultiTransformOutputSpec(TraitedSpec):
    output_image = File(exists=True, desc='Warped image')


class WarpImageMultiTransform(ANTSCommand):
    """Warps an image from one space to another

    Examples
    --------

    >>> from nipype.interfaces.ants import WarpImageMultiTransform
    >>> wimt = WarpImageMultiTransform()
    >>> wimt.inputs.input_image = 'structural.nii'
    >>> wimt.inputs.reference_image = 'ants_deformed.nii.gz'
    >>> wimt.inputs.transformation_series = ['ants_Warp.nii.gz','ants_Affine.txt']
    >>> wimt.cmdline
    'WarpImageMultiTransform 3 structural.nii structural_wimt.nii -R ants_deformed.nii.gz ants_Warp.nii.gz \
ants_Affine.txt'

    >>> wimt = WarpImageMultiTransform()
    >>> wimt.inputs.input_image = 'diffusion_weighted.nii'
    >>> wimt.inputs.reference_image = 'functional.nii'
    >>> wimt.inputs.transformation_series = ['func2anat_coreg_Affine.txt','func2anat_InverseWarp.nii.gz', \
    'dwi2anat_Warp.nii.gz','dwi2anat_coreg_Affine.txt']
    >>> wimt.inputs.invert_affine = [1]  # this will invert the 1st Affine file: 'func2anat_coreg_Affine.txt'
    >>> wimt.cmdline
    'WarpImageMultiTransform 3 diffusion_weighted.nii diffusion_weighted_wimt.nii -R functional.nii \
-i func2anat_coreg_Affine.txt func2anat_InverseWarp.nii.gz dwi2anat_Warp.nii.gz dwi2anat_coreg_Affine.txt'

    """

    _cmd = 'WarpImageMultiTransform'
    input_spec = WarpImageMultiTransformInputSpec
    output_spec = WarpImageMultiTransformOutputSpec

    def _gen_filename(self, name):
        if name == 'output_image':
            _, name, ext = split_filename(
                os.path.abspath(self.inputs.input_image))
            return ''.join((name, self.inputs.out_postfix, ext))
        return None

    def _format_arg(self, opt, spec, val):
        if opt == 'transformation_series':
            series = []
            affine_counter = 0
            affine_invert = []
            for transformation in val:
                if "affine" in transformation.lower() and \
                        isdefined(self.inputs.invert_affine):
                    affine_counter += 1
                    if affine_counter in self.inputs.invert_affine:
                        series += ['-i']
                        affine_invert.append(affine_counter)
                series += [transformation]

            if isdefined(self.inputs.invert_affine):
                diff_inv = set(self.inputs.invert_affine) - set(affine_invert)
                if diff_inv:
                    raise Exceptions(
                        "Review invert_affine, not all indexes from invert_affine were used, "
                        "check the description for the full definition")

            return ' '.join(series)

        return super(WarpImageMultiTransform, self)._format_arg(opt, spec, val)

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


class ApplyTransformsInputSpec(ANTSCommandInputSpec):
    dimension = traits.Enum(
        2,
        3,
        4,
        argstr='--dimensionality %d',
        desc=('This option forces the image to be treated '
              'as a specified-dimensional image. If not '
              'specified, antsWarp tries to infer the '
              'dimensionality from the input image.'))
    input_image_type = traits.Enum(
        0,
        1,
        2,
        3,
        argstr='--input-image-type %d',
        desc=('Option specifying the input image '
              'type of scalar (default), vector, '
              'tensor, or time series.'))
    input_image = File(
        argstr='--input %s',
        mandatory=True,
        desc=('image to apply transformation to (generally a '
              'coregistered functional)'),
        exists=True)
    output_image = traits.Str(
        argstr='--output %s',
        desc='output file name',
        genfile=True,
        hash_files=False)
    out_postfix = traits.Str(
        "_trans",
        usedefault=True,
        desc=('Postfix that is appended to all output '
              'files (default = _trans)'))
    reference_image = File(
        argstr='--reference-image %s',
        mandatory=True,
        desc='reference image space that you wish to warp INTO',
        exists=True)
    interpolation = traits.Enum(
        'Linear',
        'NearestNeighbor',
        'CosineWindowedSinc',
        'WelchWindowedSinc',
        'HammingWindowedSinc',
        'LanczosWindowedSinc',
        'MultiLabel',
        'Gaussian',
        'BSpline',
        argstr='%s',
        usedefault=True)
    interpolation_parameters = traits.Either(
        traits.Tuple(traits.Int()),  # BSpline (order)
        traits.Tuple(
            traits.Float(),  # Gaussian/MultiLabel (sigma, alpha)
            traits.Float()))
    transforms = traits.Either(
        InputMultiPath(File(exists=True)),
        'identity',
        argstr='%s',
        mandatory=True,
        desc='transform files: will be applied in reverse order. For '
        'example, the last specified transform will be applied first.')
    invert_transform_flags = InputMultiPath(traits.Bool())
    default_value = traits.Float(
        0.0, argstr='--default-value %g', usedefault=True)
    print_out_composite_warp_file = traits.Bool(
        False,
        requires=["output_image"],
        desc='output a composite warp file instead of a transformed image')
    float = traits.Bool(
        argstr='--float %d',
        default_value=False,
        usedefault=True,
        desc='Use float instead of double for computations.')


class ApplyTransformsOutputSpec(TraitedSpec):
    output_image = File(exists=True, desc='Warped image')


class ApplyTransforms(ANTSCommand):
    """ApplyTransforms, applied to an input image, transforms it according to a
    reference image and a transform (or a set of transforms).

    Examples
    --------

    >>> from nipype.interfaces.ants import ApplyTransforms
    >>> at = ApplyTransforms()
    >>> at.inputs.input_image = 'moving1.nii'
    >>> at.inputs.reference_image = 'fixed1.nii'
    >>> at.inputs.transforms = 'identity'
    >>> at.cmdline
    'antsApplyTransforms --default-value 0 --float 0 --input moving1.nii \
--interpolation Linear --output moving1_trans.nii \
--reference-image fixed1.nii -t identity'

    >>> at = ApplyTransforms()
    >>> at.inputs.dimension = 3
    >>> at.inputs.input_image = 'moving1.nii'
    >>> at.inputs.reference_image = 'fixed1.nii'
    >>> at.inputs.output_image = 'deformed_moving1.nii'
    >>> at.inputs.interpolation = 'Linear'
    >>> at.inputs.default_value = 0
    >>> at.inputs.transforms = ['ants_Warp.nii.gz', 'trans.mat']
    >>> at.inputs.invert_transform_flags = [False, False]
    >>> at.cmdline
    'antsApplyTransforms --default-value 0 --dimensionality 3 --float 0 --input moving1.nii \
--interpolation Linear --output deformed_moving1.nii --reference-image fixed1.nii \
--transform [ ants_Warp.nii.gz, 0 ] --transform [ trans.mat, 0 ]'

    >>> at1 = ApplyTransforms()
    >>> at1.inputs.dimension = 3
    >>> at1.inputs.input_image = 'moving1.nii'
    >>> at1.inputs.reference_image = 'fixed1.nii'
    >>> at1.inputs.output_image = 'deformed_moving1.nii'
    >>> at1.inputs.interpolation = 'BSpline'
    >>> at1.inputs.interpolation_parameters = (5,)
    >>> at1.inputs.default_value = 0
    >>> at1.inputs.transforms = ['ants_Warp.nii.gz', 'trans.mat']
    >>> at1.inputs.invert_transform_flags = [False, False]
    >>> at1.cmdline
    'antsApplyTransforms --default-value 0 --dimensionality 3 --float 0 --input moving1.nii \
--interpolation BSpline[ 5 ] --output deformed_moving1.nii --reference-image fixed1.nii \
--transform [ ants_Warp.nii.gz, 0 ] --transform [ trans.mat, 0 ]'
    """
    _cmd = 'antsApplyTransforms'
    input_spec = ApplyTransformsInputSpec
    output_spec = ApplyTransformsOutputSpec

    def _gen_filename(self, name):
        if name == 'output_image':
            output = self.inputs.output_image
            if not isdefined(output):
                _, name, ext = split_filename(self.inputs.input_image)
                output = name + self.inputs.out_postfix + ext
            return output
        return None

    def _get_transform_filenames(self):
        retval = []
        for ii in range(len(self.inputs.transforms)):
            if isdefined(self.inputs.invert_transform_flags):
                if len(self.inputs.transforms) == len(
                        self.inputs.invert_transform_flags):
                    invert_code = 1 if self.inputs.invert_transform_flags[
                        ii] else 0
                    retval.append("--transform [ %s, %d ]" %
                                  (self.inputs.transforms[ii], invert_code))
                else:
                    raise Exception((
                        "ERROR: The useInverse list must have the same number "
                        "of entries as the transformsFileName list."))
            else:
                retval.append("--transform %s" % self.inputs.transforms[ii])
        return " ".join(retval)

    def _get_output_warped_filename(self):
        if isdefined(self.inputs.print_out_composite_warp_file):
            return "--output [ %s, %d ]" % (
                self._gen_filename("output_image"),
                int(self.inputs.print_out_composite_warp_file))
        else:
            return "--output %s" % (self._gen_filename("output_image"))

    def _format_arg(self, opt, spec, val):
        if opt == "output_image":
            return self._get_output_warped_filename()
        elif opt == "transforms":
            if val == 'identity':
                return '-t identity'
            return self._get_transform_filenames()
        elif opt == 'interpolation':
            if self.inputs.interpolation in ['BSpline', 'MultiLabel', 'Gaussian'] and \
                    isdefined(self.inputs.interpolation_parameters):
                return '--interpolation %s[ %s ]' % (
                    self.inputs.interpolation, ', '.join([
                        str(param)
                        for param in self.inputs.interpolation_parameters
                    ]))
            else:
                return '--interpolation %s' % self.inputs.interpolation
        return super(ApplyTransforms, self)._format_arg(opt, spec, val)

    def _list_outputs(self):
        outputs = self._outputs().get()
        outputs['output_image'] = os.path.abspath(
            self._gen_filename('output_image'))
        return outputs


class ApplyTransformsToPointsInputSpec(ANTSCommandInputSpec):
    dimension = traits.Enum(
        2,
        3,
        4,
        argstr='--dimensionality %d',
        desc=('This option forces the image to be treated '
              'as a specified-dimensional image. If not '
              'specified, antsWarp tries to infer the '
              'dimensionality from the input image.'))
    input_file = File(
        argstr='--input %s',
        mandatory=True,
        desc=
        ("Currently, the only input supported is a csv file with"
         " columns including x,y (2D), x,y,z (3D) or x,y,z,t,label (4D) column headers."
         " The points should be defined in physical space."
         " If in doubt how to convert coordinates from your files to the space"
         " required by antsApplyTransformsToPoints try creating/drawing a simple"
         " label volume with only one voxel set to 1 and all others set to 0."
         " Write down the voxel coordinates. Then use ImageMaths LabelStats to find"
         " out what coordinates for this voxel antsApplyTransformsToPoints is"
         " expecting."),
        exists=True)
    output_file = traits.Str(
        argstr='--output %s',
        desc='Name of the output CSV file',
        name_source=['input_file'],
        hash_files=False,
        name_template='%s_transformed.csv')
    transforms = traits.List(
        File(exists=True),
        argstr='%s',
        mandatory=True,
        desc='transforms that will be applied to the points')
    invert_transform_flags = traits.List(
        traits.Bool(),
        desc='list indicating if a transform should be reversed')


class ApplyTransformsToPointsOutputSpec(TraitedSpec):
    output_file = File(
        exists=True, desc='csv file with transformed coordinates')


class ApplyTransformsToPoints(ANTSCommand):
    """ApplyTransformsToPoints, applied to an CSV file, transforms coordinates
    using provided transform (or a set of transforms).

    Examples
    --------

    >>> from nipype.interfaces.ants import ApplyTransforms
    >>> at = ApplyTransformsToPoints()
    >>> at.inputs.dimension = 3
    >>> at.inputs.input_file = 'moving.csv'
    >>> at.inputs.transforms = ['trans.mat', 'ants_Warp.nii.gz']
    >>> at.inputs.invert_transform_flags = [False, False]
    >>> at.cmdline
    'antsApplyTransformsToPoints --dimensionality 3 --input moving.csv --output moving_transformed.csv \
--transform [ trans.mat, 0 ] --transform [ ants_Warp.nii.gz, 0 ]'


    """
    _cmd = 'antsApplyTransformsToPoints'
    input_spec = ApplyTransformsToPointsInputSpec
    output_spec = ApplyTransformsToPointsOutputSpec

    def _get_transform_filenames(self):
        retval = []
        for ii in range(len(self.inputs.transforms)):
            if isdefined(self.inputs.invert_transform_flags):
                if len(self.inputs.transforms) == len(
                        self.inputs.invert_transform_flags):
                    invert_code = 1 if self.inputs.invert_transform_flags[
                        ii] else 0
                    retval.append("--transform [ %s, %d ]" %
                                  (self.inputs.transforms[ii], invert_code))
                else:
                    raise Exception((
                        "ERROR: The useInverse list must have the same number "
                        "of entries as the transformsFileName list."))
            else:
                retval.append("--transform %s" % self.inputs.transforms[ii])
        return " ".join(retval)

    def _format_arg(self, opt, spec, val):
        if opt == "transforms":
            return self._get_transform_filenames()
        return super(ApplyTransformsToPoints, self)._format_arg(opt, spec, val)
