# emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
"""
The dwi module of niftyfit, which wraps the fitting methods in NiftyFit.
"""

from ..base import File, TraitedSpec, traits, isdefined, CommandLineInputSpec
from .base import NiftyFitCommand
from ..niftyreg.base import get_custom_path


class FitDwiInputSpec(CommandLineInputSpec):
    """Input Spec for FitDwi."""

    # Inputs options
    source_file = File(
        position=1,
        exists=True,
        argstr="-source %s",
        mandatory=True,
        desc="The source image containing the dwi data.",
    )
    desc = "The file containing the bvalues of the source DWI."
    bval_file = File(
        position=2, exists=True, argstr="-bval %s", mandatory=True, desc=desc
    )
    desc = "The file containing the bvectors of the source DWI."
    bvec_file = File(
        position=3, exists=True, argstr="-bvec %s", mandatory=True, desc=desc
    )
    te_file = File(
        exists=True, argstr="-TE %s", desc="Filename of TEs (ms).", xor=["te_file"]
    )
    te_value = File(
        exists=True, argstr="-TE %s", desc="Value of TEs (ms).", xor=["te_file"]
    )
    mask_file = File(exists=True, desc="The image mask", argstr="-mask %s")
    desc = "Filename of parameter priors for -ball and -nod."
    prior_file = File(exists=True, argstr="-prior %s", desc=desc)
    desc = "Rotate the output tensors according to the q/s form of the image \
(resulting tensors will be in mm coordinates, default: 0)."

    rot_sform_flag = traits.Int(desc=desc, argstr="-rotsform %d")

    # generic output options:
    error_file = File(
        name_source=["source_file"],
        name_template="%s_error.nii.gz",
        desc="Filename of parameter error maps.",
        argstr="-error %s",
    )
    res_file = File(
        name_source=["source_file"],
        name_template="%s_resmap.nii.gz",
        desc="Filename of model residual map.",
        argstr="-res %s",
    )
    syn_file = File(
        name_source=["source_file"],
        name_template="%s_syn.nii.gz",
        desc="Filename of synthetic image.",
        argstr="-syn %s",
    )
    nodiff_file = File(
        name_source=["source_file"],
        name_template="%s_no_diff.nii.gz",
        desc="Filename of average no diffusion image.",
        argstr="-nodiff %s",
    )

    # Output options, with templated output names based on the source image
    mcmap_file = File(
        name_source=["source_file"],
        name_template="%s_mcmap.nii.gz",
        desc="Filename of multi-compartment model parameter map " "(-ivim,-ball,-nod)",
        argstr="-mcmap %s",
        requires=["nodv_flag"],
    )

    # Model Specific Output options:
    mdmap_file = File(
        name_source=["source_file"],
        name_template="%s_mdmap.nii.gz",
        desc="Filename of MD map/ADC",
        argstr="-mdmap %s",
    )
    famap_file = File(
        name_source=["source_file"],
        name_template="%s_famap.nii.gz",
        desc="Filename of FA map",
        argstr="-famap %s",
    )
    v1map_file = File(
        name_source=["source_file"],
        name_template="%s_v1map.nii.gz",
        desc="Filename of PDD map [x,y,z]",
        argstr="-v1map %s",
    )
    rgbmap_file = File(
        name_source=["source_file"],
        name_template="%s_rgbmap.nii.gz",
        desc="Filename of colour-coded FA map",
        argstr="-rgbmap %s",
        requires=["dti_flag"],
    )

    desc = "Use lower triangular (tenmap2) or diagonal, off-diagonal tensor \
format"

    ten_type = traits.Enum("lower-tri", "diag-off-diag", desc=desc, usedefault=True)

    tenmap_file = File(
        name_source=["source_file"],
        name_template="%s_tenmap.nii.gz",
        desc="Filename of tensor map [diag,offdiag].",
        argstr="-tenmap %s",
        requires=["dti_flag"],
    )
    tenmap2_file = File(
        name_source=["source_file"],
        name_template="%s_tenmap2.nii.gz",
        desc="Filename of tensor map [lower tri]",
        argstr="-tenmap2 %s",
        requires=["dti_flag"],
    )

    # Methods options
    desc = "Fit single exponential to non-directional data [default with \
no b-vectors]"

    mono_flag = traits.Bool(
        desc=desc,
        argstr="-mono",
        position=4,
        xor=[
            "ivim_flag",
            "dti_flag",
            "ball_flag",
            "ballv_flag",
            "nod_flag",
            "nodv_flag",
        ],
    )
    ivim_flag = traits.Bool(
        desc="Fit IVIM model to non-directional data.",
        argstr="-ivim",
        position=4,
        xor=[
            "mono_flag",
            "dti_flag",
            "ball_flag",
            "ballv_flag",
            "nod_flag",
            "nodv_flag",
        ],
    )
    desc = "Fit the tensor model [default with b-vectors]."
    dti_flag = traits.Bool(
        desc=desc,
        argstr="-dti",
        position=4,
        xor=[
            "mono_flag",
            "ivim_flag",
            "ball_flag",
            "ballv_flag",
            "nod_flag",
            "nodv_flag",
        ],
    )
    ball_flag = traits.Bool(
        desc="Fit the ball and stick model.",
        argstr="-ball",
        position=4,
        xor=[
            "mono_flag",
            "ivim_flag",
            "dti_flag",
            "ballv_flag",
            "nod_flag",
            "nodv_flag",
        ],
    )
    desc = "Fit the ball and stick model with optimised PDD."
    ballv_flag = traits.Bool(
        desc=desc,
        argstr="-ballv",
        position=4,
        xor=[
            "mono_flag",
            "ivim_flag",
            "dti_flag",
            "ball_flag",
            "nod_flag",
            "nodv_flag",
        ],
    )
    nod_flag = traits.Bool(
        desc="Fit the NODDI model",
        argstr="-nod",
        position=4,
        xor=[
            "mono_flag",
            "ivim_flag",
            "dti_flag",
            "ball_flag",
            "ballv_flag",
            "nodv_flag",
        ],
    )
    nodv_flag = traits.Bool(
        desc="Fit the NODDI model with optimised PDD",
        argstr="-nodv",
        position=4,
        xor=[
            "mono_flag",
            "ivim_flag",
            "dti_flag",
            "ball_flag",
            "ballv_flag",
            "nod_flag",
        ],
    )

    # Experimental options
    desc = "Maximum number of non-linear LSQR iterations [100x2 passes])"
    maxit_val = traits.Int(desc=desc, argstr="-maxit %d", requires=["gn_flag"])
    desc = "LM parameters (initial value, decrease rate) [100,1.2]."
    lm_vals = traits.Tuple(
        traits.Float, traits.Float, argstr="-lm %f %f", requires=["gn_flag"], desc=desc
    )
    desc = "Use Gauss-Newton algorithm [Levenberg-Marquardt]."
    gn_flag = traits.Bool(desc=desc, argstr="-gn", xor=["wls_flag"])
    desc = "Use Variational Bayes fitting with known prior (currently \
identity covariance...)."

    vb_flag = traits.Bool(desc=desc, argstr="-vb")
    cov_file = File(
        exists=True,
        desc="Filename of ithe nc*nc covariance matrix [I]",
        argstr="-cov %s",
    )
    wls_flag = traits.Bool(desc=desc, argstr="-wls", xor=["gn_flag"])
    desc = "Use location-weighted least squares for DTI fitting [3x3 Gaussian]"
    swls_val = traits.Float(desc=desc, argstr="-swls %f")
    slice_no = traits.Int(desc="Fit to single slice number.", argstr="-slice %d")
    voxel = traits.Tuple(
        traits.Int,
        traits.Int,
        traits.Int,
        desc="Fit to single voxel only.",
        argstr="-voxel %d %d %d",
    )
    diso_val = traits.Float(
        desc="Isotropic diffusivity for -nod [3e-3]", argstr="-diso %f"
    )
    dpr_val = traits.Float(
        desc="Parallel diffusivity for -nod [1.7e-3].", argstr="-dpr %f"
    )
    wm_t2_val = traits.Float(desc="White matter T2 value [80ms].", argstr="-wmT2 %f")
    csf_t2_val = traits.Float(desc="CSF T2 value [400ms].", argstr="-csfT2 %f")
    desc = "Threshold for perfusion/diffsuion effects [100]."
    perf_thr = traits.Float(desc=desc, argstr="-perfthreshold %f")

    # MCMC options:
    mcout = File(
        name_source=["source_file"],
        name_template="%s_mcout.txt",
        desc="Filename of mc samples (ascii text file)",
        argstr="-mcout %s",
    )
    mcsamples = traits.Int(
        desc="Number of samples to keep [100].", argstr="-mcsamples %d"
    )
    mcmaxit = traits.Int(
        desc="Number of iterations to run [10,000].", argstr="-mcmaxit %d"
    )
    acceptance = traits.Float(
        desc="Fraction of iterations to accept [0.23].", argstr="-accpetance %f"
    )


class FitDwiOutputSpec(TraitedSpec):
    """Output Spec for FitDwi."""

    error_file = File(desc="Filename of parameter error maps")
    res_file = File(desc="Filename of model residual map")
    syn_file = File(desc="Filename of synthetic image")
    nodiff_file = File(desc="Filename of average no diffusion image.")
    mdmap_file = File(desc="Filename of MD map/ADC")
    famap_file = File(desc="Filename of FA map")
    v1map_file = File(desc="Filename of PDD map [x,y,z]")
    rgbmap_file = File(desc="Filename of colour FA map")
    tenmap_file = File(desc="Filename of tensor map")
    tenmap2_file = File(desc="Filename of tensor map [lower tri]")

    mcmap_file = File(
        desc="Filename of multi-compartment model " "parameter map (-ivim,-ball,-nod)."
    )
    mcout = File(desc="Filename of mc samples (ascii text file)")


class FitDwi(NiftyFitCommand):
    """Interface for executable fit_dwi from Niftyfit platform.

    Use NiftyFit to perform diffusion model fitting.

    Diffusion-weighted MR Fitting.
    Fits DWI parameter maps to multi-shell, multi-directional data.

    `Source code <https://cmiclab.cs.ucl.ac.uk/CMIC/NiftyFit-Release>`_

    Examples
    --------

    >>> from nipype.interfaces import niftyfit
    >>> fit_dwi = niftyfit.FitDwi(dti_flag=True)
    >>> fit_dwi.inputs.source_file = 'dwi.nii.gz'
    >>> fit_dwi.inputs.bvec_file = 'bvecs'
    >>> fit_dwi.inputs.bval_file = 'bvals'
    >>> fit_dwi.inputs.rgbmap_file = 'rgb.nii.gz'
    >>> fit_dwi.cmdline
    'fit_dwi -source dwi.nii.gz -bval bvals -bvec bvecs -dti \
-error dwi_error.nii.gz -famap dwi_famap.nii.gz -mcout dwi_mcout.txt \
-mdmap dwi_mdmap.nii.gz -nodiff dwi_no_diff.nii.gz -res dwi_resmap.nii.gz \
-rgbmap rgb.nii.gz -syn dwi_syn.nii.gz -tenmap2 dwi_tenmap2.nii.gz \
-v1map dwi_v1map.nii.gz'

    """

    _cmd = get_custom_path("fit_dwi", env_dir="NIFTYFITDIR")
    input_spec = FitDwiInputSpec
    output_spec = FitDwiOutputSpec
    _suffix = "_fit_dwi"

    def _format_arg(self, name, trait_spec, value):
        if name == "tenmap_file" and self.inputs.ten_type != "diag-off-diag":
            return ""
        if name == "tenmap2_file" and self.inputs.ten_type != "lower-tri":
            return ""
        return super(FitDwi, self)._format_arg(name, trait_spec, value)


class DwiToolInputSpec(CommandLineInputSpec):
    """Input Spec for DwiTool."""

    desc = "The source image containing the fitted model."
    source_file = File(
        position=1, exists=True, desc=desc, argstr="-source %s", mandatory=True
    )
    desc = "The file containing the bvalues of the source DWI."
    bval_file = File(
        position=2, exists=True, desc=desc, argstr="-bval %s", mandatory=True
    )
    desc = "The file containing the bvectors of the source DWI."
    bvec_file = File(position=3, exists=True, desc=desc, argstr="-bvec %s")
    b0_file = File(
        position=4,
        exists=True,
        desc="The B0 image corresponding to the source DWI",
        argstr="-b0 %s",
    )
    mask_file = File(position=5, exists=True, desc="The image mask", argstr="-mask %s")

    # Output options, with templated output names based on the source image
    desc = "Filename of multi-compartment model parameter map \
(-ivim,-ball,-nod)"

    mcmap_file = File(
        name_source=["source_file"],
        name_template="%s_mcmap.nii.gz",
        desc=desc,
        argstr="-mcmap %s",
    )
    desc = "Filename of synthetic image. Requires: bvec_file/b0_file."
    syn_file = File(
        name_source=["source_file"],
        name_template="%s_syn.nii.gz",
        desc=desc,
        argstr="-syn %s",
        requires=["bvec_file", "b0_file"],
    )
    mdmap_file = File(
        name_source=["source_file"],
        name_template="%s_mdmap.nii.gz",
        desc="Filename of MD map/ADC",
        argstr="-mdmap %s",
    )
    famap_file = File(
        name_source=["source_file"],
        name_template="%s_famap.nii.gz",
        desc="Filename of FA map",
        argstr="-famap %s",
    )
    v1map_file = File(
        name_source=["source_file"],
        name_template="%s_v1map.nii.gz",
        desc="Filename of PDD map [x,y,z]",
        argstr="-v1map %s",
    )
    rgbmap_file = File(
        name_source=["source_file"],
        name_template="%s_rgbmap.nii.gz",
        desc="Filename of colour FA map.",
        argstr="-rgbmap %s",
    )
    logdti_file = File(
        name_source=["source_file"],
        name_template="%s_logdti2.nii.gz",
        desc="Filename of output logdti map.",
        argstr="-logdti2 %s",
    )

    # Methods options
    desc = "Input is a single exponential to non-directional data \
[default with no b-vectors]"

    mono_flag = traits.Bool(
        desc=desc,
        position=6,
        argstr="-mono",
        xor=[
            "ivim_flag",
            "dti_flag",
            "dti_flag2",
            "ball_flag",
            "ballv_flag",
            "nod_flag",
            "nodv_flag",
        ],
    )
    desc = "Inputs is an IVIM model to non-directional data."
    ivim_flag = traits.Bool(
        desc=desc,
        position=6,
        argstr="-ivim",
        xor=[
            "mono_flag",
            "dti_flag",
            "dti_flag2",
            "ball_flag",
            "ballv_flag",
            "nod_flag",
            "nodv_flag",
        ],
    )
    dti_flag = traits.Bool(
        desc="Input is a tensor model diag/off-diag.",
        position=6,
        argstr="-dti",
        xor=[
            "mono_flag",
            "ivim_flag",
            "dti_flag2",
            "ball_flag",
            "ballv_flag",
            "nod_flag",
            "nodv_flag",
        ],
    )
    dti_flag2 = traits.Bool(
        desc="Input is a tensor model lower triangular",
        position=6,
        argstr="-dti2",
        xor=[
            "mono_flag",
            "ivim_flag",
            "dti_flag",
            "ball_flag",
            "ballv_flag",
            "nod_flag",
            "nodv_flag",
        ],
    )
    ball_flag = traits.Bool(
        desc="Input is a ball and stick model.",
        position=6,
        argstr="-ball",
        xor=[
            "mono_flag",
            "ivim_flag",
            "dti_flag",
            "dti_flag2",
            "ballv_flag",
            "nod_flag",
            "nodv_flag",
        ],
    )
    desc = "Input is a ball and stick model with optimised PDD."
    ballv_flag = traits.Bool(
        desc=desc,
        position=6,
        argstr="-ballv",
        xor=[
            "mono_flag",
            "ivim_flag",
            "dti_flag",
            "dti_flag2",
            "ball_flag",
            "nod_flag",
            "nodv_flag",
        ],
    )
    nod_flag = traits.Bool(
        desc="Input is a NODDI model",
        position=6,
        argstr="-nod",
        xor=[
            "mono_flag",
            "ivim_flag",
            "dti_flag",
            "dti_flag2",
            "ball_flag",
            "ballv_flag",
            "nodv_flag",
        ],
    )
    nodv_flag = traits.Bool(
        desc="Input is a NODDI model with optimised PDD",
        position=6,
        argstr="-nodv",
        xor=[
            "mono_flag",
            "ivim_flag",
            "dti_flag",
            "dti_flag2",
            "ball_flag",
            "ballv_flag",
            "nod_flag",
        ],
    )

    # Experimental options
    diso_val = traits.Float(
        desc="Isotropic diffusivity for -nod [3e-3]", argstr="-diso %f"
    )
    dpr_val = traits.Float(
        desc="Parallel diffusivity for -nod [1.7e-3].", argstr="-dpr %f"
    )


class DwiToolOutputSpec(TraitedSpec):
    """Output Spec for DwiTool."""

    desc = "Filename of multi-compartment model parameter map \
(-ivim,-ball,-nod)"

    mcmap_file = File(desc=desc)
    syn_file = File(desc="Filename of synthetic image")
    mdmap_file = File(desc="Filename of MD map/ADC")
    famap_file = File(desc="Filename of FA map")
    v1map_file = File(desc="Filename of PDD map [x,y,z]")
    rgbmap_file = File(desc="Filename of colour FA map")
    logdti_file = File(desc="Filename of output logdti map")


class DwiTool(NiftyFitCommand):
    """Interface for executable dwi_tool from Niftyfit platform.

    Use DwiTool.

    Diffusion-Weighted MR Prediction.
    Predicts DWI from previously fitted models and calculates model derived
    maps.

    `Source code <https://cmiclab.cs.ucl.ac.uk/CMIC/NiftyFit-Release>`_

    Examples
    --------

    >>> from nipype.interfaces import niftyfit
    >>> dwi_tool = niftyfit.DwiTool(dti_flag=True)
    >>> dwi_tool.inputs.source_file = 'dwi.nii.gz'
    >>> dwi_tool.inputs.bvec_file = 'bvecs'
    >>> dwi_tool.inputs.bval_file = 'bvals'
    >>> dwi_tool.inputs.mask_file = 'mask.nii.gz'
    >>> dwi_tool.inputs.b0_file = 'b0.nii.gz'
    >>> dwi_tool.inputs.rgbmap_file = 'rgb_map.nii.gz'
    >>> dwi_tool.cmdline
    'dwi_tool -source dwi.nii.gz -bval bvals -bvec bvecs -b0 b0.nii.gz \
-mask mask.nii.gz -dti -famap dwi_famap.nii.gz -logdti2 dwi_logdti2.nii.gz \
-mcmap dwi_mcmap.nii.gz -mdmap dwi_mdmap.nii.gz -rgbmap rgb_map.nii.gz \
-syn dwi_syn.nii.gz -v1map dwi_v1map.nii.gz'

    """

    _cmd = get_custom_path("dwi_tool", env_dir="NIFTYFITDIR")
    input_spec = DwiToolInputSpec
    output_spec = DwiToolOutputSpec
    _suffix = "_dwi_tool"

    def _format_arg(self, name, trait_spec, value):
        if name == "syn_file":
            if not isdefined(self.inputs.bvec_file) or not isdefined(
                self.inputs.b0_file
            ):
                return ""
        if name in ["logdti_file", "rgbmap_file"]:
            if not isdefined(self.inputs.dti_flag) and not isdefined(
                self.inputs.dti_flag2
            ):
                return ""
        return super(DwiTool, self)._format_arg(name, trait_spec, value)