# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
"""Nipype interfaces for ANTs' commands."""
import os
from glob import glob
from nipype.interfaces import base
from nipype.interfaces.ants.base import ANTSCommandInputSpec, ANTSCommand
from nipype.interfaces.base import traits, isdefined


class _ImageMathInputSpec(ANTSCommandInputSpec):
    dimension = traits.Int(
        3, usedefault=True, position=1, argstr="%d", desc="dimension of output image"
    )
    output_image = base.File(
        position=2,
        argstr="%s",
        name_source=["op1"],
        name_template="%s_maths",
        desc="output image file",
        keep_extension=True,
    )
    operation = base.Str(
        mandatory=True, position=3, argstr="%s", desc="operations and intputs"
    )
    op1 = base.File(
        exists=True, mandatory=True, position=-2, argstr="%s", desc="first operator"
    )
    op2 = traits.Either(
        base.File(exists=True),
        base.Str,
        position=-1,
        argstr="%s",
        desc="second operator",
    )
    copy_header = traits.Bool(
        True,
        usedefault=True,
        desc="copy headers of the original image into the output (corrected) file",
    )


class _ImageMathOuputSpec(base.TraitedSpec):
    output_image = base.File(exists=True, desc="output image file")


class ImageMath(ANTSCommand):
    """
    Operations over images.

    Example
    -------
    >>> maths = ImageMath(dimension=3, op1=nifti_fname, operation='+', op2='2')
    >>> result = maths.run()
    >>> np.all(nb.load(result.outputs.output_image).get_sform() ==
    ...        nb.load(nifti_fname).get_sform())
    True

    """

    _cmd = "ImageMath"
    input_spec = _ImageMathInputSpec
    output_spec = _ImageMathOuputSpec

    def _list_outputs(self):
        outputs = super(ImageMath, self)._list_outputs()
        if self.inputs.copy_header:  # Fix headers
            _copy_header(self.inputs.op1, outputs["output_image"], set_dtype=False)
        return outputs


class _ResampleImageBySpacingInputSpec(ANTSCommandInputSpec):
    dimension = traits.Int(
        3, usedefault=True, position=1, argstr="%d", desc="dimension of output image"
    )
    input_image = base.File(
        exists=True, mandatory=True, position=2, argstr="%s", desc="input image file"
    )
    output_image = base.File(
        position=3,
        argstr="%s",
        name_source=["input_image"],
        name_template="%s_resampled",
        desc="output image file",
        keep_extension=True,
    )
    out_spacing = traits.Either(
        traits.List(traits.Float, minlen=2, maxlen=3),
        traits.Tuple(traits.Float, traits.Float, traits.Float),
        traits.Tuple(traits.Float, traits.Float),
        position=4,
        argstr="%s",
        mandatory=True,
        desc="output spacing",
    )
    apply_smoothing = traits.Bool(
        False, argstr="%d", position=5, desc="smooth before resampling"
    )
    addvox = traits.Int(
        argstr="%d",
        position=6,
        requires=["apply_smoothing"],
        desc="addvox pads each dimension by addvox",
    )
    nn_interp = traits.Bool(
        argstr="%d", desc="nn interpolation", position=-1, requires=["addvox"]
    )


class _ResampleImageBySpacingOutputSpec(base.TraitedSpec):
    output_image = traits.File(exists=True, desc="resampled file")


class ResampleImageBySpacing(ANTSCommand):
    """
    Resample an image with a given spacing.

    Examples
    --------
    >>> res = ResampleImageBySpacing(dimension=3)
    >>> res.inputs.input_image = nifti_fname
    >>> res.inputs.output_image = 'output.nii.gz'
    >>> res.inputs.out_spacing = (4, 4, 4)
    >>> res.cmdline  #doctest: +ELLIPSIS
    'ResampleImageBySpacing 3 .../test.nii.gz output.nii.gz 4 4 4'

    >>> res = ResampleImageBySpacing(dimension=3)
    >>> res.inputs.input_image = nifti_fname
    >>> res.inputs.output_image = 'output.nii.gz'
    >>> res.inputs.out_spacing = (4, 4, 4)
    >>> res.inputs.apply_smoothing = True
    >>> res.cmdline  #doctest: +ELLIPSIS
    'ResampleImageBySpacing 3 .../test.nii.gz output.nii.gz 4 4 4 1'

    >>> res = ResampleImageBySpacing(dimension=3)
    >>> res.inputs.input_image = nifti_fname
    >>> res.inputs.output_image = 'output.nii.gz'
    >>> res.inputs.out_spacing = (0.4, 0.4, 0.4)
    >>> res.inputs.apply_smoothing = True
    >>> res.inputs.addvox = 2
    >>> res.inputs.nn_interp = False
    >>> res.cmdline  #doctest: +ELLIPSIS
    'ResampleImageBySpacing 3 .../test.nii.gz output.nii.gz 0.4 0.4 0.4 1 2 0'

    """

    _cmd = "ResampleImageBySpacing"
    input_spec = _ResampleImageBySpacingInputSpec
    output_spec = _ResampleImageBySpacingOutputSpec

    def _format_arg(self, name, trait_spec, value):
        if name == "out_spacing":
            if len(value) != self.inputs.dimension:
                raise ValueError("out_spacing dimensions should match dimension")

            value = " ".join(["%g" % d for d in value])

        return super(ResampleImageBySpacing, self)._format_arg(name, trait_spec, value)


class _ThresholdImageInputSpec(ANTSCommandInputSpec):
    dimension = traits.Int(
        3, usedefault=True, position=1, argstr="%d", desc="dimension of output image"
    )
    input_image = base.File(
        exists=True, mandatory=True, position=2, argstr="%s", desc="input image file"
    )
    output_image = base.File(
        position=3,
        argstr="%s",
        name_source=["input_image"],
        name_template="%s_resampled",
        desc="output image file",
        keep_extension=True,
    )

    mode = traits.Enum(
        "Otsu",
        "Kmeans",
        argstr="%s",
        position=4,
        requires=["num_thresholds"],
        xor=["th_low", "th_high"],
        desc="whether to run Otsu / Kmeans thresholding",
    )
    num_thresholds = traits.Int(position=5, argstr="%d", desc="number of thresholds")
    input_mask = base.File(
        exists=True,
        requires=["num_thresholds"],
        argstr="%s",
        desc="input mask for Otsu, Kmeans",
    )

    th_low = traits.Float(position=4, argstr="%f", xor=["mode"], desc="lower threshold")
    th_high = traits.Float(
        position=5, argstr="%f", xor=["mode"], desc="upper threshold"
    )
    inside_value = traits.Float(
        1, position=6, argstr="%f", requires=["th_low"], desc="inside value"
    )
    outside_value = traits.Float(
        0, position=7, argstr="%f", requires=["th_low"], desc="outside value"
    )
    copy_header = traits.Bool(
        True,
        mandatory=True,
        usedefault=True,
        desc="copy headers of the original image into the output (corrected) file",
    )


class _ThresholdImageOutputSpec(base.TraitedSpec):
    output_image = traits.File(exists=True, desc="resampled file")


class ThresholdImage(ANTSCommand):
    """
    Apply thresholds on images.

    Examples
    --------
    >>> thres = ThresholdImage(dimension=3)
    >>> thres.inputs.input_image = nifti_fname
    >>> thres.inputs.output_image = 'output.nii.gz'
    >>> thres.inputs.th_low = 0.5
    >>> thres.inputs.th_high = 1.0
    >>> thres.inputs.inside_value = 1.0
    >>> thres.inputs.outside_value = 0.0
    >>> thres.cmdline  #doctest: +ELLIPSIS
    'ThresholdImage 3 .../test.nii.gz output.nii.gz 0.500000 1.000000 1.000000 0.000000'

    >>> result = thres.run()
    >>> os.path.exists(result.outputs.output_image)
    True

    >>> thres = ThresholdImage(dimension=3)
    >>> thres.inputs.input_image = nifti_fname
    >>> thres.inputs.output_image = 'output.nii.gz'
    >>> thres.inputs.mode = 'Kmeans'
    >>> thres.inputs.num_thresholds = 4
    >>> thres.cmdline  #doctest: +ELLIPSIS
    'ThresholdImage 3 .../test.nii.gz output.nii.gz Kmeans 4'

    """

    _cmd = "ThresholdImage"
    input_spec = _ThresholdImageInputSpec
    output_spec = _ThresholdImageOutputSpec

    def _list_outputs(self):
        outputs = super(ThresholdImage, self)._list_outputs()
        if self.inputs.copy_header:  # Fix headers
            _copy_header(
                self.inputs.input_image, outputs["output_image"], set_dtype=False
            )
        return outputs


class _AIInputSpec(ANTSCommandInputSpec):
    dimension = traits.Int(
        3, usedefault=True, argstr="-d %d", desc="dimension of output image"
    )
    verbose = traits.Bool(
        False, usedefault=True, argstr="-v %d", desc="enable verbosity"
    )

    fixed_image = traits.File(
        exists=True,
        mandatory=True,
        desc="Image to which the moving_image should be transformed",
    )
    moving_image = traits.File(
        exists=True,
        mandatory=True,
        desc="Image that will be transformed to fixed_image",
    )

    fixed_image_mask = traits.File(exists=True, argstr="-x %s", desc="fixed mage mask")
    moving_image_mask = traits.File(
        exists=True, requires=["fixed_image_mask"], desc="moving mage mask"
    )

    metric_trait = (
        traits.Enum("Mattes", "GC", "MI"),
        traits.Int(32),
        traits.Enum("Regular", "Random", "None"),
        traits.Range(value=0.2, low=0.0, high=1.0),
    )
    metric = traits.Tuple(
        *metric_trait, argstr="-m %s", mandatory=True, desc="the metric(s) to use."
    )

    transform = traits.Tuple(
        traits.Enum("Affine", "Rigid", "Similarity"),
        traits.Range(value=0.1, low=0.0, exclude_low=True),
        argstr="-t %s[%f]",
        usedefault=True,
        desc="Several transform options are available",
    )

    principal_axes = traits.Bool(
        False,
        usedefault=True,
        argstr="-p %d",
        xor=["blobs"],
        desc="align using principal axes",
    )
    search_factor = traits.Tuple(
        traits.Float(20),
        traits.Range(value=0.12, low=0.0, high=1.0),
        usedefault=True,
        argstr="-s [%f,%f]",
        desc="search factor",
    )

    search_grid = traits.Either(
        traits.Tuple(
            traits.Float, traits.Tuple(traits.Float, traits.Float, traits.Float)
        ),
        traits.Tuple(traits.Float, traits.Tuple(traits.Float, traits.Float)),
        argstr="-g %s",
        desc="Translation search grid in mm",
    )

    convergence = traits.Tuple(
        traits.Range(low=1, high=10000, value=10),
        traits.Float(1e-6),
        traits.Range(low=1, high=100, value=10),
        usedefault=True,
        argstr="-c [%d,%f,%d]",
        desc="convergence",
    )

    output_transform = traits.File(
        "initialization.mat", usedefault=True, argstr="-o %s", desc="output file name"
    )


class _AIOuputSpec(base.TraitedSpec):
    output_transform = traits.File(exists=True, desc="output file name")


class AI(ANTSCommand):
    """Replaces ``AffineInitializer``."""

    _cmd = "antsAI"
    input_spec = _AIInputSpec
    output_spec = _AIOuputSpec

    def _run_interface(self, runtime, correct_return_codes=(0,)):
        runtime = super(AI, self)._run_interface(runtime, correct_return_codes)

        setattr(
            self,
            "_output",
            {
                "output_transform": os.path.join(
                    runtime.cwd, os.path.basename(self.inputs.output_transform)
                )
            },
        )
        return runtime

    def _format_arg(self, opt, spec, val):
        if opt == "metric":
            val = "%s[{fixed_image},{moving_image},%d,%s,%f]" % val
            val = val.format(
                fixed_image=self.inputs.fixed_image,
                moving_image=self.inputs.moving_image,
            )
            return spec.argstr % val

        if opt == "search_grid":
            val1 = "x".join(["%f" % v for v in val[1]])
            fmtval = "[%s]" % ",".join([str(val[0]), val1])
            return spec.argstr % fmtval

        if opt == "fixed_image_mask":
            if isdefined(self.inputs.moving_image_mask):
                return spec.argstr % ("[%s,%s]" % (val, self.inputs.moving_image_mask))

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

    def _list_outputs(self):
        return getattr(self, "_output")


class _AntsJointFusionInputSpec(ANTSCommandInputSpec):
    dimension = traits.Enum(
        3,
        2,
        4,
        argstr="-d %d",
        desc="This option forces the image to be treated "
        "as a specified-dimensional image. If not "
        "specified, the program tries to infer the "
        "dimensionality from the input image.",
    )
    target_image = traits.List(
        base.InputMultiPath(base.File(exists=True)),
        argstr="-t %s",
        mandatory=True,
        desc="The target image (or "
        "multimodal target images) assumed to be "
        "aligned to a common image domain.",
    )
    atlas_image = traits.List(
        base.InputMultiPath(base.File(exists=True)),
        argstr="-g %s...",
        mandatory=True,
        desc="The atlas image (or "
        "multimodal atlas images) assumed to be "
        "aligned to a common image domain.",
    )
    atlas_segmentation_image = base.InputMultiPath(
        base.File(exists=True),
        argstr="-l %s...",
        mandatory=True,
        desc="The atlas segmentation "
        "images. For performing label fusion the number "
        "of specified segmentations should be identical "
        "to the number of atlas image sets.",
    )
    alpha = traits.Float(
        default_value=0.1,
        usedefault=True,
        argstr="-a %s",
        desc=(
            "Regularization "
            "term added to matrix Mx for calculating the inverse. Default = 0.1"
        ),
    )
    beta = traits.Float(
        default_value=2.0,
        usedefault=True,
        argstr="-b %s",
        desc=(
            "Exponent for mapping "
            "intensity difference to the joint error. Default = 2.0"
        ),
    )
    retain_label_posterior_images = traits.Bool(
        False,
        argstr="-r",
        usedefault=True,
        requires=["atlas_segmentation_image"],
        desc=(
            "Retain label posterior probability images. Requires "
            "atlas segmentations to be specified. Default = false"
        ),
    )
    retain_atlas_voting_images = traits.Bool(
        False,
        argstr="-f",
        usedefault=True,
        desc=("Retain atlas voting images. Default = false"),
    )
    constrain_nonnegative = traits.Bool(
        False,
        argstr="-c",
        usedefault=True,
        desc=("Constrain solution to non-negative weights."),
    )
    patch_radius = traits.ListInt(
        minlen=3,
        maxlen=3,
        argstr="-p %s",
        desc=("Patch radius for similarity measures." "Default: 2x2x2"),
    )
    patch_metric = traits.Enum(
        "PC",
        "MSQ",
        argstr="-m %s",
        desc=(
            "Metric to be used in determining the most similar "
            "neighborhood patch. Options include Pearson's "
            "correlation (PC) and mean squares (MSQ). Default = "
            "PC (Pearson correlation)."
        ),
    )
    search_radius = traits.List(
        [3, 3, 3],
        minlen=1,
        maxlen=3,
        argstr="-s %s",
        usedefault=True,
        desc=(
            "Search radius for similarity measures. Default = 3x3x3. "
            "One can also specify an image where the value at the "
            "voxel specifies the isotropic search radius at that voxel."
        ),
    )
    exclusion_image_label = traits.List(
        traits.Str(),
        argstr="-e %s",
        requires=["exclusion_image"],
        desc=("Specify a label for the exclusion region."),
    )
    exclusion_image = traits.List(
        base.File(exists=True),
        desc=("Specify an exclusion region for the given label."),
    )
    mask_image = base.File(
        argstr="-x %s",
        exists=True,
        desc="If a mask image "
        "is specified, fusion is only performed in the mask region.",
    )
    out_label_fusion = base.File(
        argstr="%s", hash_files=False, desc="The output label fusion image."
    )
    out_intensity_fusion_name_format = traits.Str(
        argstr="",
        desc="Optional intensity fusion "
        "image file name format. "
        '(e.g. "antsJointFusionIntensity_%d.nii.gz")',
    )
    out_label_post_prob_name_format = traits.Str(
        "antsJointFusionPosterior_%d.nii.gz",
        requires=["out_label_fusion", "out_intensity_fusion_name_format"],
        desc="Optional label posterior probability " "image file name format.",
    )
    out_atlas_voting_weight_name_format = traits.Str(
        "antsJointFusionVotingWeight_%d.nii.gz",
        requires=[
            "out_label_fusion",
            "out_intensity_fusion_name_format",
            "out_label_post_prob_name_format",
        ],
        desc="Optional atlas voting weight image " "file name format.",
    )
    verbose = traits.Bool(False, argstr="-v", desc=("Verbose output."))


class _AntsJointFusionOutputSpec(base.TraitedSpec):
    out_label_fusion = base.File(exists=True)
    out_intensity_fusion = base.OutputMultiPath(base.File(exists=True))
    out_label_post_prob = base.OutputMultiPath(base.File(exists=True))
    out_atlas_voting_weight = base.OutputMultiPath(base.File(exists=True))


class AntsJointFusion(ANTSCommand):
    """Run ``antsJoinFusion`` (finds the consensus segmentation)."""

    input_spec = _AntsJointFusionInputSpec
    output_spec = _AntsJointFusionOutputSpec
    _cmd = "antsJointFusion"

    def _format_arg(self, opt, spec, val):
        if opt == "exclusion_image_label":
            retval = []
            for ii in range(len(self.inputs.exclusion_image_label)):
                retval.append(
                    "-e {0}[{1}]".format(
                        self.inputs.exclusion_image_label[ii],
                        self.inputs.exclusion_image[ii],
                    )
                )
            retval = " ".join(retval)
        elif opt == "patch_radius":
            retval = "-p {0}".format(self._format_xarray(val))
        elif opt == "search_radius":
            retval = "-s {0}".format(self._format_xarray(val))
        elif opt == "out_label_fusion":
            if isdefined(self.inputs.out_intensity_fusion_name_format):
                if isdefined(self.inputs.out_label_post_prob_name_format):
                    if isdefined(self.inputs.out_atlas_voting_weight_name_format):
                        retval = "-o [{0}, {1}, {2}, {3}]".format(
                            self.inputs.out_label_fusion,
                            self.inputs.out_intensity_fusion_name_format,
                            self.inputs.out_label_post_prob_name_format,
                            self.inputs.out_atlas_voting_weight_name_format,
                        )
                    else:
                        retval = "-o [{0}, {1}, {2}]".format(
                            self.inputs.out_label_fusion,
                            self.inputs.out_intensity_fusion_name_format,
                            self.inputs.out_label_post_prob_name_format,
                        )
                else:
                    retval = "-o [{0}, {1}]".format(
                        self.inputs.out_label_fusion,
                        self.inputs.out_intensity_fusion_name_format,
                    )
            else:
                retval = "-o {0}".format(self.inputs.out_label_fusion)
        elif opt == "out_intensity_fusion_name_format":
            retval = ""
            if not isdefined(self.inputs.out_label_fusion):
                retval = "-o {0}".format(self.inputs.out_intensity_fusion_name_format)
        elif opt == "atlas_image":
            atlas_image_cmd = " ".join(
                [
                    "-g [{0}]".format(", ".join("'%s'" % fn for fn in ai))
                    for ai in self.inputs.atlas_image
                ]
            )
            retval = atlas_image_cmd
        elif opt == "target_image":
            target_image_cmd = " ".join(
                [
                    "-t [{0}]".format(", ".join("'%s'" % fn for fn in ai))
                    for ai in self.inputs.target_image
                ]
            )
            retval = target_image_cmd
        elif opt == "atlas_segmentation_image":
            if len(val) != len(self.inputs.atlas_image):
                raise ValueError(
                    "Number of specified segmentations should be identical to the number "
                    "of atlas image sets {0}!={1}".format(
                        len(val), len(self.inputs.atlas_image)
                    )
                )

            atlas_segmentation_image_cmd = " ".join(
                ["-l {0}".format(fn) for fn in self.inputs.atlas_segmentation_image]
            )
            retval = atlas_segmentation_image_cmd
        else:

            return super(AntsJointFusion, self)._format_arg(opt, spec, val)
        return retval

    def _list_outputs(self):
        outputs = self._outputs().get()
        if isdefined(self.inputs.out_label_fusion):
            outputs["out_label_fusion"] = os.path.abspath(self.inputs.out_label_fusion)
        if isdefined(self.inputs.out_intensity_fusion_name_format):
            outputs["out_intensity_fusion"] = glob(
                os.path.abspath(
                    self.inputs.out_intensity_fusion_name_format.replace("%d", "*")
                )
            )
        if isdefined(self.inputs.out_label_post_prob_name_format):
            outputs["out_label_post_prob"] = glob(
                os.path.abspath(
                    self.inputs.out_label_post_prob_name_format.replace("%d", "*")
                )
            )
        if isdefined(self.inputs.out_atlas_voting_weight_name_format):
            outputs["out_atlas_voting_weight"] = glob(
                os.path.abspath(
                    self.inputs.out_atlas_voting_weight_name_format.replace("%d", "*")
                )
            )
        return outputs


def _copy_header(header_file, in_file, set_dtype=True):
    """Copy header from input image to an output image."""
    import nibabel as nb

    in_img = nb.load(header_file)
    out_img = nb.load(in_file, mmap=False)
    new_img = out_img.__class__(out_img.dataobj, in_img.affine, in_img.header)
    if set_dtype:
        new_img.set_data_dtype(out_img.get_data_dtype())
    new_img.to_filename(in_file)
    return in_file
