# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
"""
Nipype interface for seg_EM.

The em module provides higher-level interfaces to some of the operations
that can be performed with the seg_em command-line program.

Examples
--------
See the docstrings of the individual classes for examples.
"""

from ..base import (TraitedSpec, File, traits, CommandLineInputSpec,
                    InputMultiPath)
from .base import NiftySegCommand
from ..niftyreg.base import get_custom_path


class EMInputSpec(CommandLineInputSpec):
    """Input Spec for EM."""
    in_file = File(
        argstr='-in %s',
        exists=True,
        mandatory=True,
        desc='Input image to segment',
        position=4)

    mask_file = File(
        argstr='-mask %s',
        exists=True,
        desc='Filename of the ROI for label fusion')

    # Priors
    no_prior = traits.Int(
        argstr='-nopriors %s',
        mandatory=True,
        desc='Number of classes to use without prior',
        xor=['prior_4D', 'priors'])

    prior_4D = File(
        argstr='-prior4D %s',
        exists=True,
        mandatory=True,
        desc='4D file containing the priors',
        xor=['no_prior', 'priors'])

    priors = InputMultiPath(
        argstr='%s',
        mandatory=True,
        desc='List of priors filepaths.',
        xor=['no_prior', 'prior_4D'])

    # iterations
    max_iter = traits.Int(
        argstr='-max_iter %s',
        default_value=100,
        usedefault=True,
        desc='Maximum number of iterations')

    min_iter = traits.Int(
        argstr='-min_iter %s',
        default_value=0,
        usedefault=True,
        desc='Minimum number of iterations')

    # other options
    bc_order_val = traits.Int(
        argstr='-bc_order %s',
        default_value=3,
        usedefault=True,
        desc='Polynomial order for the bias field')

    mrf_beta_val = traits.Float(
        argstr='-mrf_beta %s', desc='Weight of the Markov Random Field')

    desc = 'Bias field correction will run only if the ratio of improvement \
is below bc_thresh. (default=0 [OFF])'

    bc_thresh_val = traits.Float(
        argstr='-bc_thresh %s',
        default_value=0,
        usedefault=True,
        desc=desc)

    desc = 'Amount of regularization over the diagonal of the covariance \
matrix [above 1]'

    reg_val = traits.Float(argstr='-reg %s', desc=desc)

    desc = 'Outlier detection as in (Van Leemput TMI 2003). <fl1> is the \
Mahalanobis threshold [recommended between 3 and 7] <fl2> is a convergence \
ratio below which the outlier detection is going to be done [recommended 0.01]'

    outlier_val = traits.Tuple(
        traits.Float(), traits.Float(), argstr='-outlier %s %s', desc=desc)

    desc = 'Relax Priors [relaxation factor: 0<rf<1 (recommended=0.5), \
gaussian regularization: gstd>0 (recommended=2.0)] /only 3D/'

    relax_priors = traits.Tuple(
        traits.Float(), traits.Float(), argstr='-rf %s %s', desc=desc)

    # outputs
    out_file = File(
        name_source=['in_file'],
        name_template='%s_em.nii.gz',
        argstr='-out %s',
        desc='Output segmentation')
    out_bc_file = File(
        name_source=['in_file'],
        name_template='%s_bc_em.nii.gz',
        argstr='-bc_out %s',
        desc='Output bias corrected image')
    out_outlier_file = File(
        name_source=['in_file'],
        name_template='%s_outlier_em.nii.gz',
        argstr='-out_outlier %s',
        desc='Output outlierness image')


class EMOutputSpec(TraitedSpec):
    """Output Spec for EM."""
    out_file = File(desc="Output segmentation")
    out_bc_file = File(desc="Output bias corrected image")
    out_outlier_file = File(desc='Output outlierness image')


class EM(NiftySegCommand):
    """Interface for executable seg_EM from NiftySeg platform.

    seg_EM is a general purpose intensity based image segmentation tool. In
    it's simplest form, it takes in one 2D or 3D image and segments it in n
    classes.

    `Source code <http://cmictig.cs.ucl.ac.uk/wiki/index.php/NiftySeg>`_ |
    `Documentation <http://cmictig.cs.ucl.ac.uk/wiki/index.php/NiftySeg_documentation>`_

    Examples
    --------
    >>> from nipype.interfaces import niftyseg
    >>> node = niftyseg.EM()
    >>> node.inputs.in_file = 'im1.nii'
    >>> node.inputs.no_prior = 4
    >>> node.cmdline
    'seg_EM -in im1.nii -bc_order 3 -bc_thresh 0 -max_iter 100 -min_iter 0 -nopriors 4 \
-bc_out im1_bc_em.nii.gz -out im1_em.nii.gz -out_outlier im1_outlier_em.nii.gz'

    """
    _cmd = get_custom_path('seg_EM', env_dir='NIFTYSEGDIR')
    _suffix = '_em'
    input_spec = EMInputSpec
    output_spec = EMOutputSpec

    def _format_arg(self, opt, spec, val):
        """Convert input to appropriate format for seg_EM."""
        if opt == 'priors':
            _nb_priors = len(self.inputs.priors)
            return '-priors %d %s' % (_nb_priors, ' '.join(self.inputs.priors))
        else:
            return super(EM, self)._format_arg(opt, spec, val)
